21 lines
556 B
Python
21 lines
556 B
Python
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
|