Initial commit

metric_depth/dataset/hypersim.py

@@ -0,0 +1,74 @@
+import cv2
+import h5py
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+from torchvision.transforms import Compose
+
+from dataset.transform import Resize, NormalizeImage, PrepareForNet, Crop
+
+
+def hypersim_distance_to_depth(npyDistance):
+    intWidth, intHeight, fltFocal = 1024, 768, 886.81
+
+    npyImageplaneX = np.linspace((-0.5 * intWidth) + 0.5, (0.5 * intWidth) - 0.5, intWidth).reshape(
+        1, intWidth).repeat(intHeight, 0).astype(np.float32)[:, :, None]
+    npyImageplaneY = np.linspace((-0.5 * intHeight) + 0.5, (0.5 * intHeight) - 0.5,
+                                 intHeight).reshape(intHeight, 1).repeat(intWidth, 1).astype(np.float32)[:, :, None]
+    npyImageplaneZ = np.full([intHeight, intWidth, 1], fltFocal, np.float32)
+    npyImageplane = np.concatenate(
+        [npyImageplaneX, npyImageplaneY, npyImageplaneZ], 2)
+
+    npyDepth = npyDistance / np.linalg.norm(npyImageplane, 2, 2) * fltFocal
+    return npyDepth
+
+
+class Hypersim(Dataset):
+    def __init__(self, filelist_path, mode, size=(518, 518)):
+        
+        self.mode = mode
+        self.size = size
+        
+        with open(filelist_path, 'r') as f:
+            self.filelist = f.read().splitlines()
+        
+        net_w, net_h = size
+        self.transform = Compose([
+            Resize(
+                width=net_w,
+                height=net_h,
+                resize_target=True if mode == 'train' else False,
+                keep_aspect_ratio=True,
+                ensure_multiple_of=14,
+                resize_method='lower_bound',
+                image_interpolation_method=cv2.INTER_CUBIC,
+            ),
+            NormalizeImage(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+            PrepareForNet(),
+        ] + ([Crop(size[0])] if self.mode == 'train' else []))
+        
+    def __getitem__(self, item):
+        img_path = self.filelist[item].split(' ')[0]
+        depth_path = self.filelist[item].split(' ')[1]
+        
+        image = cv2.imread(img_path)
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) / 255.0
+        
+        depth_fd = h5py.File(depth_path, "r")
+        distance_meters = np.array(depth_fd['dataset'])
+        depth = hypersim_distance_to_depth(distance_meters)
+        
+        sample = self.transform({'image': image, 'depth': depth})
+
+        sample['image'] = torch.from_numpy(sample['image'])
+        sample['depth'] = torch.from_numpy(sample['depth'])
+        
+        sample['valid_mask'] = (torch.isnan(sample['depth']) == 0)
+        sample['depth'][sample['valid_mask'] == 0] = 0
+        
+        sample['image_path'] = self.filelist[item].split(' ')[0]
+        
+        return sample
+
+    def __len__(self):
+        return len(self.filelist)