def _get_projected_positions_of_sphere_points(self, sphere_points, rotation, translation): """ For the given points on unit sphere calculates the 3D coordinates on the mesh template and projects them back to image plane :param sphere_points: A (B X 3 X H X W) tensor containing the predicted points on the sphere :param rotation: A (B X CP X 3 X 3) camera rotation tensor :param translation: A (B X CP X 3) camera translation tensor :return: A tuple(xy, z, uv, uv_3d) - xy - (B X CP X 2 X H X W) x,y values of the 3D points after projecting onto image plane - z - (B X CP X 1 X H X W) z value of the projection - uv - (B X 2 X H X W) UV values of the sphere coordinates - uv_3d - (B X H X W X 3) tensor with the 3D coordinates on the mesh template for the given sphere coordinates """ uv = convert_3d_to_uv_coordinates(sphere_points.permute(0, 2, 3, 1)) batch_size = uv.size(0) height = uv.size(1) width = uv.size(2) num_poses = rotation.size(1) uv_flatten = uv.view(-1, 2) uv_3d = self.uv_to_3d(uv_flatten).view(batch_size, 1, -1, 3) uv_3d = uv_3d.repeat(1, num_poses, 1, 1).view(batch_size*num_poses, -1, 3) cameras = OpenGLOrthographicCameras(device=sphere_points.device, R=rotation.view(-1, 3, 3), T=translation.view(-1, 3)) xyz_cam = cameras.get_world_to_view_transform().transform_points(uv_3d) z = xyz_cam[:, :, 2:].view(batch_size, num_poses, height, width, 1) xy = cameras.transform_points(uv_3d)[:, :, :2].view(batch_size, num_poses, height, width, 2) xy = xy.permute(0, 1, 4, 2, 3).flip(2) z = z.permute(0, 1, 4, 2, 3) uv = uv.permute(0, 3, 1, 2) uv_3d = uv_3d.view(batch_size, num_poses, height, width, 3)[:, 0, :, :, :].squeeze() return xy, z, uv, uv_3d
scale, trans, quat, True) camera = OpenGLOrthographicCameras(device=device, R=rotation, T=translation) kps = (((data['kp'].to(device, dtype=torch.float) + 1) / 2) * 255).to( torch.int32) kp = draw_key_points(img, kps, key_point_colors) fig = plt.figure() plt.subplot(2, 2, 1) plt.imshow(kp[0].permute(1, 2, 0).cpu()) # Plot the key points directly using the 3D keypoints and projecting them onto image plane kp_3d = torch.from_numpy(dataset.kp_3d).to( device, dtype=torch.float32).unsqueeze(0) xyz = camera.transform_points(kp_3d) xy = (((xyz[:, :, :2] + 1) / 2) * 255).to(torch.int32) kp_xy = torch.cat((xy, kps[:, :, 2:]), dim=2) kp_pred_direct = draw_key_points(img, kp_xy, key_point_colors) plt.subplot(2, 2, 2) plt.imshow(kp_pred_direct[0].permute(1, 2, 0).cpu()) # Draw key points by converting 3D kps to uv values and then back to 3D kp_uv = convert_3d_to_uv_coordinates(kp_3d) uv_flatten = kp_uv.view(-1, 2) uv_3d = uv_to_3d(uv_flatten).view(1, -1, 3) xyz = camera.transform_points(uv_3d) xy = (((xyz[:, :, :2] + 1) / 2) * 255).to(torch.int32) kp_xy = torch.cat((xy, kps[:, :, 2:]), dim=2) kp_pred = draw_key_points(img, kp_xy, key_point_colors) plt.subplot(2, 2, 3)