Update dataset/datasets.py

dataset/datasets.py

@@ -28,10 +28,10 @@ class TrainDataset(torch.utils.data.Dataset):
             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'
+            dfm_name = image_name[:image_name.rfind('.')] + '_' + str(mult_number) + '.png'
         else:
             bw_name = self.data[idx]
-            dfm_name =  os.path.splitext(self.data[idx])[0] + '0_dfm.png'
+            dfm_name =  os.path.splitext(self.data[idx])[0] + '0.png'
             
             
         bw_img =  np.expand_dims(plt.imread(os.path.join(self.data_path, 'bw', bw_name)), 2)
@@ -81,10 +81,10 @@ class FineTuningDataset(torch.utils.data.Dataset):
             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'
+            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'
+            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)