Merge pull request #2 from Dutra-Apex/fix_num_samples_error

Update dataset.py

src/data/dataset.py

@@ -18,10 +18,10 @@ def __init__(self, args):
         self.mask_type = args.mask_type
 
         # image and mask 
-        self.image_path = []
-        for ext in ['*.jpg', '*.png']: 
-            self.image_path.extend(glob(os.path.join(args.dir_image, args.data_train, ext)))
-        self.mask_path = glob(os.path.join(args.dir_mask, args.mask_type, '*.png'))
+        self.image_path, self.mask_path = [], []
+        for ext in ['*.jpg', '*.png', '*.jpeg']: 
+            self.image_path.extend(glob(f'{args.dir_image}*{ext}))
+            self.mask_path.extend(glob(f'{args.dir_image}*{ext}))
 
         # augmentation 
         self.img_trans = transforms.Compose([
@@ -77,4 +77,4 @@ def __getitem__(self, index):
     data = InpaintingData(args)
     print(len(data), len(data.mask_path))
     img, mask, filename = data[0]
-    print(img.size(), mask.size(), filename)
+    print(img.size(), mask.size(), filename)