Method get_data

eval_code/recons/datasets/nrgbd.py:115–194

(
            self,
            index: Optional[int] = None,
            sequence_name: Optional[str] = None,
            ids: Union[List[int], np.ndarray, None] = None,
        )

        return self.get_data(index=index, ids=ids)

    def get_data(
            self,
            index: Optional[int] = None,
            sequence_name: Optional[str] = None,
            ids: Union[List[int], np.ndarray, None] = None,
        ):
        if sequence_name is None:
            if index is None:
                raise ValueError("Please specify either index or sequence_name")
            sequence_name: str = self.sequence_list[index]
        seq_extrinsics: np.ndarray = self.metadata[sequence_name]  # (N, 3, 4)
        seq_len: int = seq_extrinsics.shape[0]

        if ids is None:
            ids = np.arange(seq_len).tolist()
        elif isinstance(ids, np.ndarray):
            assert ids.ndim == 1, f"ids should be a 1D array, but got {ids.ndim}D"
            ids = ids.tolist()

        fx, fy, cx, cy = 554.2562584220408, 554.2562584220408, 320, 240  # hard code

        image_paths: list      = [""] * len(ids)
        images: list           = [0]  * len(ids)
        depths: list           = [0]  * len(ids)
        extrinsics: np.ndarray = seq_extrinsics[ids]  # (N, 3, 4) -> (S, 3, 4)
        intrinsics: np.ndarray = np.tile(
            np.array(
                [
                    [fx, 0,  cx],
                    [0,  fy, cy],
                    [0,  0,  1 ]
                ],
                dtype=np.float32
            ),
            reps=(len(ids), 1, 1)
        )  # (S, 3, 3)

        for id_index, id in enumerate(ids):
            
            impath    = osp.join(self.NRGBD_DIR, sequence_name, "images", f"img{id}.png")
            depthpath = osp.join(self.NRGBD_DIR, sequence_name, "depth", f"depth{id}.png")

            rgb_image: Image.Image = Image.open(impath)
            depthmap: np.ndarray   = imageio.v2.imread(depthpath)
            assert depthmap.shape == (480, 640), f"Depth map shape {depthmap.shape} does not match expected (480, 640)"
            rgb_image: Image.Image = resize_image(rgb_image, (depthmap.shape[1], depthmap.shape[0]))

            depthmap = np.nan_to_num(depthmap.astype(np.float32), 0.0) / 1000.0
            depthmap[depthmap > 10] = 0    # to far, invalid
            depthmap[depthmap < 1e-3] = 0  # to near, invalid

            rgb_image, depthmap, intrinsics[id_index] = resize_image_depth_and_intrinsic(
                image=rgb_image,
                depth_map=depthmap,
                intrinsic=intrinsics[id_index],
                output_width=self.load_img_size, # finally width = 518, height = 388
            )