Update dataset/datasets.py

dataset/datasets.py

@@ -45,6 +45,28 @@ class TrainDataset(torch.utils.data.Dataset):
 
         return bw_img, color_img, dfm_img
 
+        class FineTuningDataset(Dataset):
+    def __init__(self, data_path, transform=None):
+        self.data_path = data_path
+        self.transform = transform
+
+        self.data = []
+        for file in os.listdir(data_path):
+            if file.endswith('.png'):
+                self.data.append(file)
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        image_path = os.path.join(self.data_path, self.data[idx])
+        image = Image.open(image_path)
+
+        if self.transform:
+            image = self.transform(image)
+
+        return image
+
 
 def main():
     train_dataset = TrainDataset(data_path='./train', transform=transforms.ToTensor())