Update dataset/datasets.py

dataset/datasets.py

@@ -27,20 +27,14 @@ class TrainDataset(torch.utils.data.Dataset):
 
         if self.mults_amount > 1:
             mult_number = np.random.choice(range(self.mults_amount))
-            bw_name = image_name[:image_name.rfind('.')] + '_' + str(mult_number) + '.png'
-            dfm_name = image_name[:image_name.rfind('.')] + '_' + str(mult_number) + '_dfm.png'
         else:
-            bw_name = image_name
-            dfm_name = os.path.splitext(image_name)[0] + '_dfm.png'
+            mult_number = 0
 
-        bw_img = np.expand_dims(plt.imread(os.path.join(self.data_path, 'bw', bw_name)), 2)
-        dfm_img = np.expand_dims(plt.imread(os.path.join(self.data_path, 'bw', dfm_name)), 2)
+        bw_name = f"{os.path.splitext(image_name)[0]}_{mult_number}.png"
+        dfm_name = f"{os.path.splitext(image_name)[0]}_{mult_number}_dfm.png"
 
-        # Guardar imágenes en el directorio bw_directory
-        bw_img_path = os.path.join(self.bw_directory, bw_name)
-        dfm_img_path = os.path.join(self.bw_directory, dfm_name)
-        plt.imsave(bw_img_path, bw_img.squeeze(), cmap='gray')
-        plt.imsave(dfm_img_path, dfm_img.squeeze(), cmap='gray')
+        bw_img = np.expand_dims(plt.imread(os.path.join(self.bw_directory, bw_name)), 2)
+        dfm_img = np.expand_dims(plt.imread(os.path.join(self.bw_directory, dfm_name)), 2)
 
         # Normalización y generación de máscara
         bw_img = self.ToTensor(bw_img)
@@ -54,6 +48,7 @@ class TrainDataset(torch.utils.data.Dataset):
 
         return bw_img, color_img, hint, dfm_img
 
+# Resto del código...
 
 class FineTuningDataset(torch.utils.data.Dataset):
     def __init__(self, data_path, transform=None, mult_amount=1):