Update dataset/datasets.py

dataset/datasets.py

@@ -8,11 +8,10 @@ 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):
         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):
@@ -23,14 +22,8 @@ class TrainDataset(torch.utils.data.Dataset):
 
         color_img = Image.open(os.path.join(self.data_path, 'color', image_name)).convert('RGB')
 
-        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 = 'dfm_' + self.data[idx]
+        bw_name = 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)