(
self,
index: Optional[int] = None,
sequence_name: Optional[str] = None,
ids: Union[List[int], np.ndarray, None] = None,
)
| 76 | return self.get_data(index=index, ids=ids) |
| 77 | |
| 78 | def get_data( |
| 79 | self, |
| 80 | index: Optional[int] = None, |
| 81 | sequence_name: Optional[str] = None, |
| 82 | ids: Union[List[int], np.ndarray, None] = None, |
| 83 | ): |
| 84 | if sequence_name is None: |
| 85 | if index is None: |
| 86 | raise ValueError("Please specify either index or sequence_name") |
| 87 | sequence_name: str = self.sequence_list[index] |
| 88 | image_list: list = self.metadata[sequence_name] |
| 89 | seq_len: int = len(image_list) |
| 90 | |
| 91 | if ids is None: |
| 92 | ids = np.arange(seq_len).tolist() |
| 93 | elif isinstance(ids, np.ndarray): |
| 94 | assert ids.ndim == 1, f"ids should be a 1D array, but got {ids.ndim}D" |
| 95 | ids = ids.tolist() |
| 96 | |
| 97 | image_paths: list = [""] * len(ids) |
| 98 | images: list = [0] * len(ids) |
| 99 | depths: list = [0] * len(ids) |
| 100 | extrinsics: np.ndarray = np.zeros((len(ids), 3, 4)) |
| 101 | intrinsics: np.ndarray = np.zeros((len(ids), 3, 3)) |
| 102 | |
| 103 | for id_index, id in enumerate(ids): |
| 104 | img_name = image_list[id] |
| 105 | impath = os.path.join(self.ETH3D_DIR, sequence_name, 'images', 'custom_undistorted', img_name) |
| 106 | depthpath = os.path.join(self.ETH3D_DIR, sequence_name, 'ground_truth_depth', 'custom_undistorted', img_name) |
| 107 | cam_path = os.path.join(self.ETH3D_DIR, sequence_name, 'custom_undistorted_cam', img_name.replace('JPG', 'npz')) |
| 108 | |
| 109 | cam = np.load(cam_path) |
| 110 | intrinsic = cam['intrinsics'] |
| 111 | extrinsic = cam['extrinsics'] |
| 112 | |
| 113 | # load image and depth |
| 114 | rgb_image: Image.Image = Image.open(impath) |
| 115 | width, height = rgb_image.size |
| 116 | depthmap: np.ndarray = np.fromfile(depthpath, dtype=np.float32).reshape(height, width) |
| 117 | depthmap[~np.isfinite(depthmap)] = -1 |
| 118 | |
| 119 | rgb_image, depthmap, intrinsic = resize_image_depth_and_intrinsic( |
| 120 | image=rgb_image, |
| 121 | depth_map=depthmap, |
| 122 | intrinsic=intrinsic, |
| 123 | output_width=self.load_img_size, # finally width = 518, height = 388 |
| 124 | ) |
| 125 | |
| 126 | image_paths[id_index] = impath |
| 127 | images[id_index] = to_tensor(rgb_image) |
| 128 | depths[id_index] = depthmap |
| 129 | intrinsics[id_index] = intrinsic |
| 130 | extrinsics[id_index] = extrinsic[:3, :] |
| 131 | |
| 132 | depths = np.array(depths) # (S, H, W) |
| 133 | pointclouds = unproject_depth_map_to_point_map( |
| 134 | depth_map=depths[..., None], |
| 135 | intrinsics_cam=intrinsics, |
no test coverage detected