def __getitem__()

in threedod/benchmark_scripts/utils/tenFpsDataLoader.py [0:0]

• 43 lines of code
• 7 McCabe index (conditional complexity)

    def __getitem__(self, idx):
        """
        Returns:
            frame: a dict
                {frame_id}: str
                {depth}: (h, w)
                {image}: (h, w)
                {image_path}: str
                {intrinsics}: np.array 3x3
                {pose}: np.array 4x4
                {pcd}: np.array (n, 3)
                    in world coordinate
                {color}: (n, 3)
        """
        frame_id = self.frame_ids[idx]
        frame = {}
        frame["frame_id"] = frame_id
        fname = "{}_{}.png".format(self.video_id, frame_id)
        # fname = "{}.png".format(frame_id)
        depth_image_path = os.path.join(self.root_path, "lowres_depth", fname)
        if not os.path.exists(depth_image_path):
            print(depth_image_path)

        image_path = os.path.join(self.root_path, "lowres_wide", fname)

        if not os.path.exists(depth_image_path):
            print(depth_image_path, "does not exist")
        frame["depth"] = cv2.imread(depth_image_path, -1)
        frame["image"] = cv2.imread(image_path)
        frame["image_path"] = image_path
        depth_height, depth_width = frame["depth"].shape
        im_height, im_width, im_channels = frame["image"].shape

        frame["intrinsics"] = copy.deepcopy(self.intrinsics[frame_id])
        if str(frame_id) in self.poses.keys():
            frame_pose = np.array(self.poses[str(frame_id)])
        else:
            for my_key in list(self.poses.keys()):
                if abs(float(frame_id) - float(my_key)) < 0.005:
                    frame_pose = np.array(self.poses[str(my_key)])
        frame["pose"] = copy.deepcopy(frame_pose)

        im_height_scale = np.float(depth_height) / im_height
        im_width_scale = np.float(depth_width) / im_width

        if depth_height != im_height:
            frame["image"] = np.zeros([depth_height, depth_width, 3])  # 288, 384, 3
            frame["image"][48 : 48 + 192, 64 : 64 + 256, :] = cv2.imread(image_path)

        (m, n, _) = frame["image"].shape
        depth_image = frame["depth"] / 1000.0
        rgb_image = frame["image"] / 255.0

        pcd, rgb_feat = generate_point(
            rgb_image,
            depth_image,
            frame["intrinsics"],
            self.subsample,
            self.world_coordinate,
            frame_pose,
        )

        frame["pcd"] = pcd
        frame["color"] = rgb_feat
        return frame