def point_cloud_3d()

in pytouch/tasks/surface_3d.py [0:0]

• 28 lines of code
• 4 McCabe index (conditional complexity)

    def point_cloud_3d(self, img_color, img_normal_gt=None, img_depth_gt=None):
        normal_pred = self.normals(img_color)

        color = copy.deepcopy(img_color)
        normal = copy.deepcopy(normal_pred)

        # TODO (psodhi): Background gt normals nx, ny are non-zero for sim but correctly zero for real.
        # Hence, we mask out background in sim relying on gt depth. Once background is fixed,
        # we can remove the code snippet below.
        bg_mask = None
        if self.sensor == "sim":
            depth_gt = copy.deepcopy(img_depth_gt)
            bg_mask = (depth_gt > self.gel_depth).squeeze()

        img_grad_depth, grad_x, grad_y = self.normal_to_grad_depth(
            normal, self.sensor_params.gel_width, self.sensor_params.gel_height, bg_mask
        )

        img_depth = self.grad_depth_to_depth(
            img_grad_depth,
            grad_x,
            grad_y,
            bg_mask,
            remove_bg_depth=self.sensor_params.remove_background_depth,
            max_depth=self.sensor_params.max_depth,
        )

        img_points3d = self.depth_to_points3d(
            img_depth, self.sensor_params.T_cam_offset_sim, self.sensor_params.P
        )

        # numpy
        img_color_np = (color.permute(1, 2, 0)).cpu().detach().numpy()
        img_normal_np = (normal.permute(1, 2, 0)).cpu().detach().numpy()
        img_depth_np = img_depth.cpu().detach().numpy()

        return self.Surface3DReturn(
            img_points3d, img_color_np, img_normal_np, img_depth_np
        )