optimised

sets/CasiaFace.py

@@ -0,0 +1,21 @@
+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