Update dataset/datasets.py

dataset/datasets.py

@@ -1,11 +1,4 @@
-import torch
-import os
 from PIL import Image
-import torchvision.transforms as transforms
-import matplotlib.pyplot as plt
-import numpy as np
-
-from utils.utils import generate_mask
 
 class TrainDataset(torch.utils.data.Dataset):
     def __init__(self, data_path, transform=None):
@@ -14,6 +7,11 @@ class TrainDataset(torch.utils.data.Dataset):
         self.transform = transform
         self.ToTensor = transforms.ToTensor()
 
+        # Directorio para guardar las imágenes en blanco y negro
+        self.bw_directory = os.path.join(data_path, 'bw')
+        if not os.path.exists(self.bw_directory):
+            os.makedirs(self.bw_directory)
+
     def __len__(self):
         return len(self.data)
 
@@ -25,8 +23,8 @@ class TrainDataset(torch.utils.data.Dataset):
         bw_name = "bw_" + image_name
         dfm_name = 'dfm_' + image_name
 
-        bw_img = np.expand_dims(np.array(Image.open(os.path.join(self.data_path, 'bw', bw_name)).convert('L')), 2)
-        dfm_img = np.expand_dims(np.array(Image.open(os.path.join(self.data_path, 'bw', dfm_name)).convert('L')), 2)
+        bw_img = np.expand_dims(np.array(Image.open(os.path.join(self.bw_directory, bw_name)).convert('L')), 2)
+        dfm_img = np.expand_dims(np.array(Image.open(os.path.join(self.bw_directory, dfm_name)).convert('L')), 2)
 
         bw_img = np.concatenate([bw_img, dfm_img], axis=2)
 
@@ -56,6 +54,11 @@ class FineTuningDataset(torch.utils.data.Dataset):
         self.data_path = data_path
         self.transform = transform
 
+        # Directorio para guardar las imágenes en blanco y negro
+        self.bw_directory = os.path.join(data_path, 'bw')
+        if not os.path.exists(self.bw_directory):
+            os.makedirs(self.bw_directory)
+
         np.random.shuffle(self.color_data)
         self.ToTensor = transforms.ToTensor()
 
@@ -64,8 +67,11 @@ class FineTuningDataset(torch.utils.data.Dataset):
 
     def __getitem__(self, idx):
         color_img = Image.open(os.path.join(self.data_path, 'color', self.color_data[idx])).convert('RGB')
-        bw_img = np.expand_dims(np.array(Image.open(os.path.join(self.data_path, 'real_manga', self.data[idx])).convert('L')), 2)
-        dfm_img = np.expand_dims(np.array(Image.open(os.path.join(self.data_path, 'real_manga', 'dfm_' + self.data[idx])).convert('L')), 2)
+        bw_name = "bw_" + self.data[idx]
+        dfm_name = "dfm_" + self.data[idx]
+
+        bw_img = np.expand_dims(np.array(Image.open(os.path.join(self.bw_directory, bw_name)).convert('L')), 2)
+        dfm_img = np.expand_dims(np.array(Image.open(os.path.join(self.bw_directory, dfm_name)).convert('L')), 2)
 
         if self.transform:
             result = self.transform(image=np.array(color_img))