Update dataset/datasets.py

dataset/datasets.py

@@ -1,81 +1,107 @@
-import os
 import torch
+import os
 import torchvision.transforms as transforms
 import matplotlib.pyplot as plt
-from torch.utils.data import Dataset
-from PIL import Image
 import numpy as np
 
+from utils.utils import generate_mask
+
+
 class TrainDataset(torch.utils.data.Dataset):
-    def __init__(self, data_path, transform=None, mults_amount=1):
+    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()
-
-        self.file_list = os.listdir(os.path.join(data_path, 'color'))
-
     def __len__(self):
-        return len(self.file_list)
-
+        return len(self.data)
+    
     def __getitem__(self, idx):
-        image_name = self.file_list[idx]
-
-        # Extract mult_number from the image name (last part before the extension)
-        mult_number = int(image_name.split('_')[-1].split('.')[0])
-
-        # Construct corresponding bw and dfm image names
-        bw_name = f"{image_name.split('.')[0]}_{mult_number}_dfm.png"
-        color_img_path = os.path.join(self.data_path, 'color', image_name)
-        bw_img_path = os.path.join(self.data_path, 'bw', bw_name)
-
-        color_img = Image.open(color_img_path).convert('RGB')
-        bw_img = Image.open(bw_img_path).convert('L')
+        image_name = self.data[idx]
+        
+        color_img = plt.imread(os.path.join(self.data_path, 'color', image_name))
+        
 
+        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 = 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)
+        
         if self.transform:
-            color_img = self.transform(color_img)
-            bw_img = self.transform(bw_img)
-
-        # Assuming generate_mask is a function generating a mask for hint
-        mask = generate_mask(bw_img.size[1], bw_img.size[0])
-        hint = torch.cat((bw_img * mask, mask), 0)
-
-        return bw_img, color_img, hint
+            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.endswith('_dfm.png')]
-        print("Files in FineTuning Dataset:", self.data)  # Debug print
-        # rest of your code remains unchanged
-
+    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.mult_amount = mult_amount
+        self.mults_amount = mult_amount
+        
+        np.random.shuffle(self.color_data)
+        
         self.ToTensor = transforms.ToTensor()
-
-        self.file_list = [x for x in os.listdir(os.path.join(data_path, 'real_manga')) if x.endswith('_dfm.png')]
-
     def __len__(self):
-        return len(self.file_list)
-
+        return len(self.data)
+    
     def __getitem__(self, idx):
-        image_name = self.file_list[idx]
-
-        # Extract mult_number from the image name (last part before '_dfm.png')
-        mult_number = int(image_name.split('_')[-2])
-
-        # Construct corresponding bw and dfm image names
-        bw_name = f"{image_name.split('_')[0]}_{mult_number}.png"
-        color_img_path = os.path.join(self.data_path, 'color', bw_name)
-        bw_img_path = os.path.join(self.data_path, 'real_manga', image_name)
-
-        color_img = Image.open(color_img_path).convert('RGB')
-        bw_img = Image.open(bw_img_path).convert('L')
-
+        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'
+        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)
+        
         if self.transform:
-            color_img = self.transform(color_img)
-            bw_img = self.transform(bw_img)
-
-        return bw_img, color_img
-
+            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