re-structured codes

Author: Fangchang Ma

nyu_dataloader.py renamed to dataloaders/dataloader.py

@@ -3,11 +3,9 @@
 import numpy as np
 import torch.utils.data as data
 import h5py
-import transforms
+import dataloaders.transforms as transforms
 
-IMG_EXTENSIONS = [
-    '.h5',
-]
+IMG_EXTENSIONS = ['.h5',]
 
 def is_image_file(filename):
     return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)
@@ -22,106 +20,60 @@ def make_dataset(dir, class_to_idx):
     images = []
     dir = os.path.expanduser(dir)
     for target in sorted(os.listdir(dir)):
-        # print(target)
         d = os.path.join(dir, target)
         if not os.path.isdir(d):
             continue
-
         for root, _, fnames in sorted(os.walk(d)):
             for fname in sorted(fnames):
                 if is_image_file(fname):
                     path = os.path.join(root, fname)
                     item = (path, class_to_idx[target])
                     images.append(item)
-
     return images
 
 def h5_loader(path):
     h5f = h5py.File(path, "r")
     rgb = np.array(h5f['rgb'])
     rgb = np.transpose(rgb, (1, 2, 0))
     depth = np.array(h5f['depth'])
-
     return rgb, depth
 
-iheight, iwidth = 480, 640 # raw image size
-oheight, owidth = 228, 304 # image size after pre-processing
-color_jitter = transforms.ColorJitter(0.4, 0.4, 0.4)
-
-def train_transform(rgb, depth):
-    s = np.random.uniform(1.0, 1.5) # random scaling
-    # print("scale factor s={}".format(s))
-    depth_np = depth / s
-    angle = np.random.uniform(-5.0, 5.0) # random rotation degrees
-    do_flip = np.random.uniform(0.0, 1.0) < 0.5 # random horizontal flip
-
-    # perform 1st part of data augmentation
-    transform = transforms.Compose([
-        transforms.Resize(250.0 / iheight), # this is for computational efficiency, since rotation is very slow
-        transforms.Rotate(angle),
-        transforms.Resize(s),
-        transforms.CenterCrop((oheight, owidth)),
-        transforms.HorizontalFlip(do_flip)
-    ])
-    rgb_np = transform(rgb)
-
-    # random color jittering 
-    rgb_np = color_jitter(rgb_np)
-
-    rgb_np = np.asfarray(rgb_np, dtype='float') / 255
-    depth_np = transform(depth_np)
-
-    return rgb_np, depth_np
-
-def val_transform(rgb, depth):
-    depth_np = depth
-
-    # perform 1st part of data augmentation
-    transform = transforms.Compose([
-        transforms.Resize(240.0 / iheight),
-        transforms.CenterCrop((oheight, owidth)),
-    ])
-    rgb_np = transform(rgb)
-    rgb_np = np.asfarray(rgb_np, dtype='float') / 255
-    depth_np = transform(depth_np)
-
-    return rgb_np, depth_np
-
-def rgb2grayscale(rgb):
-    return rgb[:,:,0] * 0.2989 + rgb[:,:,1] * 0.587 + rgb[:,:,2] * 0.114
-
+# def rgb2grayscale(rgb):
+#     return rgb[:,:,0] * 0.2989 + rgb[:,:,1] * 0.587 + rgb[:,:,2] * 0.114
 
 to_tensor = transforms.ToTensor()
 
-class NYUDataset(data.Dataset):
+class MyDataloader(data.Dataset):
     modality_names = ['rgb', 'rgbd', 'd'] # , 'g', 'gd'
 
     def __init__(self, root, type, sparsifier=None, modality='rgb', loader=h5_loader):
         classes, class_to_idx = find_classes(root)
         imgs = make_dataset(root, class_to_idx)
-        if len(imgs) == 0:
-            raise(RuntimeError("Found 0 images in subfolders of: " + root + "\n"
-                               "Supported image extensions are: " + ",".join(IMG_EXTENSIONS)))
-
+        assert len(imgs)>0, "Found 0 images in subfolders of: " + root + "\n"
+        print("Found {} images in {} folder.".format(len(imgs), type))
         self.root = root
         self.imgs = imgs
         self.classes = classes
         self.class_to_idx = class_to_idx
         if type == 'train':
-            self.transform = train_transform
+            self.transform = self.train_transform
         elif type == 'val':
-            self.transform = val_transform
+            self.transform = self.val_transform
         else:
             raise (RuntimeError("Invalid dataset type: " + type + "\n"
                                 "Supported dataset types are: train, val"))
         self.loader = loader
         self.sparsifier = sparsifier
 
-        if modality in self.modality_names:
-            self.modality = modality
-        else:
-            raise (RuntimeError("Invalid modality type: " + modality + "\n"
-                                "Supported dataset types are: " + ''.join(self.modality_names)))
+        assert (modality in self.modality_names), "Invalid modality type: " + modality + "\n" + \
+                                "Supported dataset types are: " + ''.join(self.modality_names)
+        self.modality = modality
+
+    def train_transform(self, rgb, depth):
+        raise (RuntimeError("train_transform() is not implemented. "))
+
+    def val_transform(rgb, depth):
+        raise (RuntimeError("val_transform() is not implemented."))
 
     def create_sparse_depth(self, rgb, depth):
         if self.sparsifier is None:
@@ -134,7 +86,6 @@ def create_sparse_depth(self, rgb, depth):
 
     def create_rgbd(self, rgb, depth):
         sparse_depth = self.create_sparse_depth(rgb, depth)
-        # rgbd = np.dstack((rgb[:,:,0], rgb[:,:,1], rgb[:,:,2], sparse_depth))
         rgbd = np.append(rgb, np.expand_dims(sparse_depth, axis=2), axis=2)
         return rgbd
 
@@ -150,14 +101,7 @@ def __getraw__(self, index):
         rgb, depth = self.loader(path)
         return rgb, depth
 
-    def __get_all_item__(self, index):
-        """
-        Args:
-            index (int): Index
-
-        Returns:
-            tuple: (input_tensor, depth_tensor, input_np, depth_np) 
-        """
+    def __getitem__(self, index):
         rgb, depth = self.__getraw__(index)
         if self.transform is not None:
             rgb_np, depth_np = self.transform(rgb, depth)
@@ -181,19 +125,40 @@ def __get_all_item__(self, index):
         depth_tensor = to_tensor(depth_np)
         depth_tensor = depth_tensor.unsqueeze(0)
 
-        return input_tensor, depth_tensor, input_np, depth_np
-
-    def __getitem__(self, index):
-        """
-        Args:
-            index (int): Index
-
-        Returns:
-            tuple: (input_tensor, depth_tensor) 
-        """
-        input_tensor, depth_tensor, input_np, depth_np = self.__get_all_item__(index)
-
         return input_tensor, depth_tensor
 
     def __len__(self):
         return len(self.imgs)
+
+    # def __get_all_item__(self, index):
+    #     """
+    #     Args:
+    #         index (int): Index
+
+    #     Returns:
+    #         tuple: (input_tensor, depth_tensor, input_np, depth_np)
+    #     """
+    #     rgb, depth = self.__getraw__(index)
+    #     if self.transform is not None:
+    #         rgb_np, depth_np = self.transform(rgb, depth)
+    #     else:
+    #         raise(RuntimeError("transform not defined"))
+
+    #     # color normalization
+    #     # rgb_tensor = normalize_rgb(rgb_tensor)
+    #     # rgb_np = normalize_np(rgb_np)
+
+    #     if self.modality == 'rgb':
+    #         input_np = rgb_np
+    #     elif self.modality == 'rgbd':
+    #         input_np = self.create_rgbd(rgb_np, depth_np)
+    #     elif self.modality == 'd':
+    #         input_np = self.create_sparse_depth(rgb_np, depth_np)
+
+    #     input_tensor = to_tensor(input_np)
+    #     while input_tensor.dim() < 3:
+    #         input_tensor = input_tensor.unsqueeze(0)
+    #     depth_tensor = to_tensor(depth_np)
+    #     depth_tensor = depth_tensor.unsqueeze(0)
+
+    #     return input_tensor, depth_tensor, input_np, depth_np

dense_to_sparse.py renamed to dataloaders/dense_to_sparse.py

@@ -0,0 +1,48 @@
+import numpy as np
+import dataloaders.transforms as transforms
+from dataloaders.dataloader import MyDataloader
+
+color_jitter = transforms.ColorJitter(0.4, 0.4, 0.4)
+
+class KITTIDataset(MyDataloader):
+    def __init__(self, root, type, sparsifier=None, modality='rgb'):
+        super(KITTIDataset, self).__init__(root, type, sparsifier=None, modality='rgb')
+        self.output_size = (228, 912)
+
+    def train_transform(self, rgb, depth):
+        s = np.random.uniform(1.0, 1.5)  # random scaling
+        depth_np = depth / s
+        angle = np.random.uniform(-5.0, 5.0)  # random rotation degrees
+        do_flip = np.random.uniform(0.0, 1.0) < 0.5  # random horizontal flip
+
+        # perform 1st step of data augmentation
+        transform = transforms.Compose([
+            transforms.Crop(130, 10, 240, 1200),
+            transforms.Rotate(angle),
+            transforms.Resize(s),
+            transforms.CenterCrop(self.output_size),
+            transforms.HorizontalFlip(do_flip)
+        ])
+        rgb_np = transform(rgb)
+        rgb_np = color_jitter(rgb_np) # random color jittering
+        rgb_np = np.asfarray(rgb_np, dtype='float') / 255
+        # Scipy affine_transform produced RuntimeError when the depth map was
+        # given as a 'numpy.ndarray'
+        depth_np = np.asfarray(depth_np, dtype='float32')
+        depth_np = transform(depth_np)
+
+        return rgb_np, depth_np
+
+    def val_transform(self, rgb, depth):
+        depth_np = depth
+        transform = transforms.Compose([
+            transforms.Crop(130, 10, 240, 1200),
+            transforms.CenterCrop(self.output_size),
+        ])
+        rgb_np = transform(rgb)
+        rgb_np = np.asfarray(rgb_np, dtype='float') / 255
+        depth_np = np.asfarray(depth_np, dtype='float32')
+        depth_np = transform(depth_np)
+
+        return rgb_np, depth_np
+

dataloaders/kitti_dataloader.py

@@ -0,0 +1,44 @@
+import numpy as np
+import dataloaders.transforms as transforms
+from dataloaders.dataloader import MyDataloader
+
+iheight, iwidth = 480, 640 # raw image size
+color_jitter = transforms.ColorJitter(0.4, 0.4, 0.4)
+
+class NYUDataset(MyDataloader):
+    def __init__(self, root, type, sparsifier=None, modality='rgb'):
+        super(NYUDataset, self).__init__(root, type, sparsifier=None, modality='rgb')
+        self.output_size = (228, 304)
+
+    def train_transform(self, rgb, depth):
+        s = np.random.uniform(1.0, 1.5) # random scaling
+        depth_np = depth / s
+        angle = np.random.uniform(-5.0, 5.0) # random rotation degrees
+        do_flip = np.random.uniform(0.0, 1.0) < 0.5 # random horizontal flip
+
+        # perform 1st step of data augmentation
+        transform = transforms.Compose([
+            transforms.Resize(250.0 / iheight), # this is for computational efficiency, since rotation can be slow
+            transforms.Rotate(angle),
+            transforms.Resize(s),
+            transforms.CenterCrop(self.output_size),
+            transforms.HorizontalFlip(do_flip)
+        ])
+        rgb_np = transform(rgb)
+        rgb_np = color_jitter(rgb_np) # random color jittering
+        rgb_np = np.asfarray(rgb_np, dtype='float') / 255
+        depth_np = transform(depth_np)
+
+        return rgb_np, depth_np
+
+    def val_transform(self, rgb, depth):
+        depth_np = depth
+        transform = transforms.Compose([
+            transforms.Resize(240.0 / iheight),
+            transforms.CenterCrop(self.output_size),
+        ])
+        rgb_np = transform(rgb)
+        rgb_np = np.asfarray(rgb_np, dtype='float') / 255
+        depth_np = transform(depth_np)
+
+        return rgb_np, depth_np