support more evaluation datasets

Author: jingheya

datasets/eval/depth/configs/data_diode_all.yaml

@@ -0,0 +1,6 @@
+name: diode
+disp_name: diode_val_all
+# dir: diode
+dir: diode/diode_val.tar
+filenames: datasets/eval/depth/data_split/diode/diode_val_all_filename_list.txt 
+processing_res: 640

datasets/eval/depth/data_split/diode/diode_val_all_filename_list.txt

@@ -207,6 +207,7 @@ def gen_normal(img, pipe, prompt="", num_inference_steps=1):
                 "kitti": "configs/data_kitti_eigen_test.yaml",
                 "scannet": "configs/data_scannet_val.yaml",
                 "eth3d": "configs/data_eth3d.yaml",
+                "diode": "configs/data_diode_all.yaml",
             }
             for dataset_name, config_path in test_depth_dataset_configs.items():
                 eval_dir = os.path.join(args.output_dir, args.task_name, dataset_name)
@@ -218,7 +219,7 @@ def gen_normal(img, pipe, prompt="", num_inference_steps=1):
         elif args.task_name == 'normal':
             test_data_dir = os.path.join(args.base_test_data_dir, args.task_name)
             dataset_split_path = "evaluation/dataset_normal"
-            eval_datasets = [('nyuv2', 'test'), ('scannet', 'test'), ('ibims', 'ibims'), ('sintel', 'sintel')]
+            eval_datasets = [ ('nyuv2', 'test'), ('scannet', 'test'), ('ibims', 'ibims'), ('sintel', 'sintel'),  ('oasis', 'val')]
             eval_dir = os.path.join(args.output_dir, args.task_name)
             evaluation_normal(eval_dir, test_data_dir, dataset_split_path, eval_mode="generate_prediction", 
                               gen_prediction=gen_normal, pipeline=pipeline, eval_datasets=eval_datasets,

datasets/eval/depth/data_split/diode/diode_val_indoor_filename_list.txt

@@ -8,13 +8,14 @@
 from .kitti_dataset import KITTIDataset
 from .nyu_dataset import NYUDataset
 from .scannet_dataset import ScanNetDataset
-
+from .diode_dataset import DIODEDataset
 
 dataset_name_class_dict = {
     "nyu_v2": NYUDataset,
     "kitti": KITTIDataset,
     "eth3d": ETH3DDataset,
     "scannet": ScanNetDataset,
+    "diode": DIODEDataset,
 }
 
 

datasets/eval/depth/data_split/diode/diode_val_outdoor_filename_list.txt

@@ -0,0 +1,72 @@
+# Author: Bingxin Ke
+# Last modified: 2024-02-26
+
+import os
+import tarfile
+from io import BytesIO
+
+import numpy as np
+import torch
+
+from .base_depth_dataset import BaseDepthDataset, DepthFileNameMode, DatasetMode
+
+
+class DIODEDataset(BaseDepthDataset):
+    def __init__(
+        self,
+        **kwargs,
+    ) -> None:
+        super().__init__(
+            # DIODE data parameter
+            min_depth=0.6,
+            max_depth=350,
+            has_filled_depth=False,
+            name_mode=DepthFileNameMode.id,
+            **kwargs,
+        )
+
+    def _read_npy_file(self, rel_path):
+        if self.is_tar:
+            if self.tar_obj is None:
+                self.tar_obj = tarfile.open(self.dataset_dir)
+            fileobj = self.tar_obj.extractfile("./" + rel_path)
+            npy_path_or_content = BytesIO(fileobj.read())
+        else:
+            npy_path_or_content = os.path.join(self.dataset_dir, rel_path)
+        data = np.load(npy_path_or_content).squeeze()[np.newaxis, :, :]
+        return data
+
+    def _read_depth_file(self, rel_path):
+        depth = self._read_npy_file(rel_path)
+        return depth
+
+    def _get_data_path(self, index):
+        return self.filenames[index]
+
+    def _get_data_item(self, index):
+        # Special: depth mask is read from data
+
+        rgb_rel_path, depth_rel_path, mask_rel_path = self._get_data_path(index=index)
+
+        rasters = {}
+
+        # RGB data
+        rasters.update(self._load_rgb_data(rgb_rel_path=rgb_rel_path))
+
+        # Depth data
+        if DatasetMode.RGB_ONLY != self.mode:
+            # load data
+            depth_data = self._load_depth_data(
+                depth_rel_path=depth_rel_path, filled_rel_path=None
+            )
+            rasters.update(depth_data)
+
+            # valid mask
+            mask = self._read_npy_file(mask_rel_path).astype(bool)
+            mask = torch.from_numpy(mask).bool()
+            rasters["valid_mask_raw"] = mask.clone()
+            rasters["valid_mask_filled"] = mask.clone()
+
+        other = {"index": index, "rgb_relative_path": rgb_rel_path}
+
+        return rasters, other

eval.py

@@ -0,0 +1,76 @@
+""" Get samples from OASIS validation set (https://pvl.cs.princeton.edu/OASIS/)
+"""
+import os
+import cv2
+import numpy as np
+import pickle
+
+from evaluation.dataset_normal import Sample
+
+
+def read_normal(path, h, w):
+    normal_dict = pickle.load(open(path, 'rb'))
+
+    mask = np.zeros((h,w))
+    normal = np.zeros((h,w,3))
+
+    # Stuff ROI normal into bounding box
+    min_y = normal_dict['min_y']
+    max_y = normal_dict['max_y']
+    min_x = normal_dict['min_x']
+    max_x = normal_dict['max_x']
+    roi_normal = normal_dict['normal']
+
+    # to LUB
+    normal[min_y:max_y+1, min_x:max_x+1, :] = roi_normal
+    normal = normal.astype(np.float32)
+    normal[:,:,0] *= -1
+    normal[:,:,1] *= -1
+
+    # Make mask
+    roi_mask = np.logical_or(np.logical_or(roi_normal[:,:,0] != 0, roi_normal[:,:,1] != 0), roi_normal[:,:,2] != 0).astype(np.float32)
+    mask[min_y:max_y+1, min_x:max_x+1] = roi_mask
+    mask = mask[:, :, None]
+    mask = mask > 0.5
+
+    return normal, mask
+
+
+def get_sample(base_data_dir, sample_path, info):
+    # e.g. sample_path = "val/100277_DT_img.png"
+    scene_name = sample_path.split('/')[0]
+    img_name, img_ext = sample_path.split('/')[-1].split('_img')
+
+    dataset_path = os.path.join(base_data_dir, 'dsine_eval', 'oasis')
+    img_path = '%s/%s' % (dataset_path, sample_path)
+    normal_path = img_path.replace('_img'+img_ext, '_normal.pkl')
+    intrins_path = img_path.replace('_img'+img_ext, '_intrins.npy')
+    assert os.path.exists(img_path)
+    assert os.path.exists(normal_path)
+    assert os.path.exists(intrins_path)
+
+    # read image (H, W, 3)
+    img = cv2.cvtColor(cv2.imread(img_path, cv2.IMREAD_UNCHANGED), cv2.COLOR_BGR2RGB)
+    img = img.astype(np.float32) / 255.0
+
+    # read normal (H, W, 3)
+    h = img.shape[0]
+    w = img.shape[1]
+    normal, normal_mask = read_normal(normal_path, h, w)
+
+    # read intrins (3, 3)
+    intrins = np.load(intrins_path)
+
+    sample = Sample(
+        img=img,
+        normal=normal,
+        normal_mask=normal_mask,
+        intrins=intrins,
+
+        dataset_name='oasis',
+        scene_name=scene_name,
+        img_name=img_name,
+        info=info
+    )
+
+    return sample