Update dataset/datasets.py

dataset/datasets.py

@@ -22,7 +22,7 @@ class TrainDataset(torch.utils.data.Dataset):
 
         color_img = Image.open(os.path.join(self.data_path, 'color', image_name)).convert('RGB')
 
-        bw_name = "_bw" + image_name
+        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)