from core.datasets.image.ImageDataset import ImageDataset
from sklearn.datasets import load_digits


class DigitsDataset(ImageDataset):

    def __init__(self, seed):
        super().__init__(
            name='digits',
            num_classes=10,
            num_features=64,
            sample_shape=(8, 8, 1),
            seed=seed
        )

    def _get_x_y(self):
        dataset = load_digits()
        x = dataset.data
        y = dataset.target
        return x, y
