Carvana dataset loader

milesial · Jul 30, 2020 · 4ad8323 · 4ad8323
1 parent 84f8392
commit 4ad8323
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/utils/dataset.py b/utils/dataset.py
@@ -9,10 +9,11 @@
 
 
 class BasicDataset(Dataset):
-    def __init__(self, imgs_dir, masks_dir, scale=1):
+    def __init__(self, imgs_dir, masks_dir, scale=1, mask_suffix=''):
         self.imgs_dir = imgs_dir
         self.masks_dir = masks_dir
         self.scale = scale
+        self.mask_suffix = mask_suffix
         assert 0 < scale <= 1, 'Scale must be between 0 and 1'
 
         self.ids = [splitext(file)[0] for file in listdir(imgs_dir)
@@ -43,7 +44,7 @@ def preprocess(cls, pil_img, scale):
 
     def __getitem__(self, i):
         idx = self.ids[i]
-        mask_file = glob(self.masks_dir + idx + '.*')
+        mask_file = glob(self.masks_dir + idx + self.mask_suffix + '.*')
         img_file = glob(self.imgs_dir + idx + '.*')
 
         assert len(mask_file) == 1, \
@@ -63,3 +64,8 @@ def __getitem__(self, i):
             'image': torch.from_numpy(img).type(torch.FloatTensor),
             'mask': torch.from_numpy(mask).type(torch.FloatTensor)
         }
+
+
+class CarvanaDataset(BasicDataset):
+    def __init__(self, imgs_dir, masks_dir, scale=1):
+        super().__init__(imgs_dir, masks_dir, scale, mask_suffix='_mask')