22 lines
781 B
Python
22 lines
781 B
Python
import os
|
|
from torchvision import datasets
|
|
from torch.utils.data import Dataset
|
|
import torch
|
|
from .Data import Data
|
|
|
|
class CasiaSet(Data):
|
|
def __init__(self, resolution: int = 224, sample_size = 190):
|
|
super().__init__(resolution = resolution, sample_size = sample_size)
|
|
|
|
def get_set(self) -> Data:
|
|
path_str = "./datasets/casia-set"
|
|
path = os.path.abspath(path_str)
|
|
|
|
if not os.path.exists(path):
|
|
raise FileNotFoundError(f"Unpacked dataset missing at {self.final_path}. Run Extractor.py first!")
|
|
print(f"Loading unpacked CASIA dataset from: {self.final_path}")
|
|
set = datasets.ImageFolder(root=path, transform=None)
|
|
# we set the target here
|
|
self.target = torch.tensor(set.targets)
|
|
return set
|