Update dataset/datasets.py

dataset/datasets.py

@@ -64,13 +64,13 @@ class FineTuningDataset(torch.utils.data.Dataset):
         self.color
 
         bw_img = plt.imread(get_file_path(bw_name, self.data_path, self.mults_amount))
-dfm_img = plt.imread(get_file_path(dfm_name, self.data_path, self.mults_amount))
+        dfm_img = plt.imread(get_file_path(dfm_name, self.data_path, self.mults_amount))
 
-bw_img = bw_img.astype('float32')
-dfm_img = dfm_img.astype('float32')
+        bw_img = bw_img.astype('float32')
+        dfm_img = dfm_img.astype('float32')
 
-color_img = bw_img[:, :, 0]
-hint = dfm_img
+        color_img = bw_img[:, :, 0]
+        hint = dfm_img
 
-return bw_img, color_img, hint, dfm_img
+        return bw_img, color_img, hint, dfm_img