from torchvision import datasets
from torch.utils.data import Dataset
import torch
from .Data import Data

class CelebA(Data):
    def __init__(self, resolution: int = 224, sample_size = 30):
        super().__init__(resolution, sample_size = sample_size)

    def get_set(self):
        set = datasets.CelebA(
            root = "../data", 
            split='all', 
            target_type='identity', 
            download=False, 
            transform=None
        )
        # set the target first
        self.target = set.identity
        return set