Method __getitem__

pointersect/data/mesh_dataset_v2.py:237–335

Returns: point_cloud: PointCloud, (b=1, n), n: number of points ray: Ray, (b=1, q=n_target_img, ho, wo) target rays ray_gt_dict: ray_rgbs: (b=1, q=n_target_img, ho, wo, 3) ray_ts: (b=1, q=n

(self, i)

        return self.total

    def __getitem__(self, i) -> T.Dict[str, T.Any]:
        """
        Returns:
            point_cloud:
                PointCloud, (b=1, n),  n: number of points
            ray:
                Ray, (b=1, q=n_target_img, ho, wo)  target rays
            ray_gt_dict:
                ray_rgbs: (b=1, q=n_target_img, ho, wo, 3)
                ray_ts: (b=1, q=n_target_img, ho, wo)
                surface_normals_w: (b=1, q=n_target_img, ho, wo, 3)
                hit_map: (b=1, q=n_target_img, ho, wo)  1 if hit a surface, 0 otherwise
        """

        # read texture images
        if self.texture_idxs is not None:
            texture_imgs = []
            for tidx in range(len(self.texture_idxs[i])):
                timeout_sec = 5
                try:
                    # time out for 5 seconds passed
                    img = func_timeout(
                        timeout_sec, imageio.v3.imread, args=(self.texture_filenames[self.texture_idxs[i, tidx]],))
                except FunctionTimedOut:
                    print(
                        f"reading image {self.texture_filenames[self.texture_idxs[i, tidx]]}, time out after {timeout_sec} secs")
                    img = np.ones((500, 500, 3))
                except Exception as e:
                    print(
                        f"reading image {self.texture_filenames[self.texture_idxs[i, tidx]]}, failed")
                    img = np.ones((500, 500, 3))

                if len(img.shape) == 2:
                    img = np.stack((img, img, img), axis=-1)

                if img.shape[-1] > 3:
                    img = img[..., :3]

                # crop image
                h, w = img.shape[:2]
                if self.texture_crop_method == 'full':
                    s = min(h, w)
                elif isinstance(self.texture_crop_method, str):
                    # {min_s}-{max_s}
                    min_s, max_s = self.texture_crop_method.split('-', 1)
                    max_s = min(w, min(h, int(max_s)))
                    min_s = int(min_s)
                    min_s = max(1, min_s)
                    max_s = max(1, max_s)
                    min_s = min(min_s, max_s)
                    s = self.rng.randint(low=min_s, high=max_s + 1)
                elif isinstance(self.texture_crop_method, int):
                    max_s = min(h, w)
                    min_s = min(max_s, self.texture_crop_method)
                    s = self.rng.randint(low=min_s, high=max_s + 1)
                else:
                    raise NotImplementedError