Update dataset/datasets.py

dataset/datasets.py

@@ -6,10 +6,12 @@ from torch.utils.data import Dataset
 from PIL import Image
 import numpy as np
 
-class TrainDataset(Dataset):
-    def __init__(self, data_path, transform=None):
+class TrainDataset(torch.utils.data.Dataset):
+    def __init__(self, data_path, transform=None, mults_amount=1):
+        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()
 
         self.file_list = os.listdir(os.path.join(data_path, 'color'))