Update dataset/datasets.py

dataset/datasets.py

@@ -62,3 +62,15 @@ 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.find('_dfm') == -1] 
         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))
+
+bw_img = bw_img.astype('float32')
+dfm_img = dfm_img.astype('float32')
+
+color_img = bw_img[:, :, 0]
+hint = dfm_img
+
+return bw_img, color_img, hint, dfm_img
+