Update dataset/datasets.py

dataset/datasets.py

@@ -8,103 +8,135 @@ from utils.utils import generate_mask
 
 
 class TrainDataset(torch.utils.data.Dataset):
-    def __init__(self, data_path, transform = None, mults_amount = 1):
-        self.data = os.listdir(os.path.join(data_path, 'color'))
+    def __init__(self, data_path, transform=None, mults_amount=1):
+        self.data = os.listdir(os.path.join(data_path, "color"))
         self.data_path = data_path
         self.transform = transform
         self.mults_amount = mults_amount
-        
+
         self.ToTensor = transforms.ToTensor()
+
     def __len__(self):
         return len(self.data)
-    
+
     def __getitem__(self, idx):
-    image_name = self.data[idx]
-    
+        image_name = self.data[idx]
+
     try:
-        color_img = plt.imread(os.path.join(self.data_path, 'color', image_name))
+        color_img = plt.imread(os.path.join(self.data_path, "color", image_name))
     except SyntaxError:
         print(f"Archivo {image_name} no es un PNG válido. Saltando...")
         return None  # O alguna otra acción que prefieras
 
         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'
+
+            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 = self.data[idx]
-            dfm_name =  os.path.splitext(self.data[idx])[0] + '0_dfm.png'
-            
-            
-        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_img = np.concatenate([bw_img, dfm_img], axis = 2)
-        
+            dfm_name = os.path.splitext(self.data[idx])[0] + "0_dfm.png"
+
+        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_img = np.concatenate([bw_img, dfm_img], axis=2)
+
         if self.transform:
-            result = self.transform(image = color_img, mask = bw_img)
-            color_img = result['image']
-            bw_img = result['mask']
-          
+            result = self.transform(image=color_img, mask=bw_img)
+            color_img = result["image"]
+            bw_img = result["mask"]
+
         dfm_img = bw_img[:, :, 1]
         bw_img = bw_img[:, :, 0]
-        
+
         color_img = self.ToTensor(color_img)
         bw_img = self.ToTensor(bw_img)
-        
+
         dfm_img = self.ToTensor(dfm_img)
-        
+
         color_img = (color_img - 0.5) / 0.5
-        
+
         mask = generate_mask(bw_img.shape[1], bw_img.shape[2])
         hint = torch.cat((color_img * mask, mask), 0)
-        
+
         return bw_img, color_img, hint, dfm_img
-    
+
+
 class FineTuningDataset(torch.utils.data.Dataset):
-    def __init__(self, data_path, transform = None, mult_amount = 1):
-        self.data = [x for x in os.listdir(os.path.join(data_path, 'real_manga')) if x.find('_dfm') == -1] 
-        self.color_data = [x for x in os.listdir(os.path.join(data_path, 'color'))]
+    def __init__(self, data_path, transform=None, mult_amount=1):
+        self.data = [
+            x
+            for x in os.listdir(os.path.join(data_path, "real_manga"))
+            if x.find("_dfm") == -1
+        ]
+        self.color_data = [x for x in os.listdir(os.path.join(data_path, "color"))]
         self.data_path = data_path
         self.transform = transform
         self.mults_amount = mult_amount
-        
+
         np.random.shuffle(self.color_data)
-        
+
         self.ToTensor = transforms.ToTensor()
+
     def __len__(self):
         return len(self.data)
-    
+
     def __getitem__(self, idx):
-        color_img = plt.imread(os.path.join(self.data_path, 'color', self.color_data[idx]))
-        
+        color_img = plt.imread(
+            os.path.join(self.data_path, "color", self.color_data[idx])
+        )
+
         image_name = self.data[idx]
         if self.mults_amount > 1:
             mult_number = np.random.choice(range(self.mults_amount))
-            
-            bw_name = image_name[:image_name.rfind('.')] + '_' + str(self.mults_amount) + '.png'
-            dfm_name = image_name[:image_name.rfind('.')] + '_' + str(self.mults_amount) + '_dfm.png'
+
+            bw_name = (
+                image_name[: image_name.rfind(".")]
+                + "_"
+                + str(self.mults_amount)
+                + ".png"
+            )
+            dfm_name = (
+                image_name[: image_name.rfind(".")]
+                + "_"
+                + str(self.mults_amount)
+                + "_dfm.png"
+            )
         else:
             bw_name = self.data[idx]
-            dfm_name =  os.path.splitext(self.data[idx])[0] + '_dfm.png'
-        
-        
-        bw_img =  np.expand_dims(plt.imread(os.path.join(self.data_path, 'real_manga', image_name)), 2)
-        dfm_img =  np.expand_dims(plt.imread(os.path.join(self.data_path, 'real_manga', dfm_name)), 2)
-        
+            dfm_name = os.path.splitext(self.data[idx])[0] + "_dfm.png"
+
+        bw_img = np.expand_dims(
+            plt.imread(os.path.join(self.data_path, "real_manga", image_name)), 2
+        )
+        dfm_img = np.expand_dims(
+            plt.imread(os.path.join(self.data_path, "real_manga", dfm_name)), 2
+        )
+
         if self.transform:
-            result = self.transform(image = color_img)
-            color_img = result['image']
-            
-            result = self.transform(image = bw_img,  mask = dfm_img)
-            bw_img = result['image']
-            dfm_img = result['mask']
-        
+            result = self.transform(image=color_img)
+            color_img = result["image"]
+
+            result = self.transform(image=bw_img, mask=dfm_img)
+            bw_img = result["image"]
+            dfm_img = result["mask"]
+
         color_img = self.ToTensor(color_img)
         bw_img = self.ToTensor(bw_img)
         dfm_img = self.ToTensor(dfm_img)
-        
+
         color_img = (color_img - 0.5) / 0.5
-        
+
         return bw_img, dfm_img, color_img