Update app.py

app.py

@@ -86,7 +86,6 @@ class Pix2PixDataset(torch.utils.data.Dataset):
         return len(self.originals)
 
     def __getitem__(self, idx):
-        # Get original and target images
         original_img = self.originals[idx]['image']
         target_img = self.targets[idx]['image']
         
@@ -94,12 +93,12 @@ class Pix2PixDataset(torch.utils.data.Dataset):
         original = original_img.convert('RGB')
         target = target_img.convert('RGB')
         
-        # Extract the filename from the original image's path
-        original_img_path = original_img.filename  # Ensure this is the correct way to get the path
+        # Extract the filename from the image_path in the CSV
+        original_img_path = self.data.iloc[idx]['image_path']
         original_img_filename = os.path.basename(original_img_path)
-    
+        
         # Match the image filename with the `image_path` column in the CSV
-        matched_row = self.data[self.data['image_path'].str.contains(original_img_filename)]
+        matched_row = self.data[self.data['image_path'].str.endswith(original_img_filename)]
         
         if matched_row.empty:
             raise ValueError(f"No matching entry found in the CSV for image {original_img_filename}")