Update app.py

app.py

@@ -37,17 +37,29 @@ def train_model(epochs):
     ])
 
     # Create dataset and dataloader
-    class Pix2PixDataset(torch.utils.data.Dataset):
-        def __init__(self, ds):
-            self.ds = ds
+    # Create dataset and dataloader
+class Pix2PixDataset(torch.utils.data.Dataset):
+    def __init__(self, ds):
+        self.originals = [x for x in ds["train"] if x['label'] == 'original']
+        self.targets = [x for x in ds["train"] if x['label'] == 'target']
+
+        # Ensure original and target images match by their index
+        assert len(self.originals) == len(self.targets), "Mismatch in number of original and target images."
+
+    def __len__(self):
+        return len(self.originals)
+
+    def __getitem__(self, idx):
+        # Load original and target images for the given index
+        original_img = self.originals[idx]['image']
+        target_img = self.targets[idx]['image']
 
-        def __len__(self):
-            return len(self.ds["train"])
+        # Apply the necessary transforms
+        original = Image.open(original_img).convert('RGB')
+        target = Image.open(target_img).convert('RGB')
 
-        def __getitem__(self, idx):
-            original = Image.open(self.ds["train"][idx]['original_image']).convert('RGB')
-            target = Image.open(self.ds["train"][idx]['target_image']).convert('RGB')
-            return transform(original), transform(target)
+        # Return transformed original and target images
+        return transform(original), transform(target)
 
     dataset = Pix2PixDataset(ds)
     dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)