(color, depth, sem_feature, intrinsics, w2c, transform_pts=True,
mask=None, compute_mean_sq_dist=False, mean_sq_dist_method="projective")
| 44 | |
| 45 | |
| 46 | def get_pointcloud(color, depth, sem_feature, intrinsics, w2c, transform_pts=True, |
| 47 | mask=None, compute_mean_sq_dist=False, mean_sq_dist_method="projective"): |
| 48 | width, height = color.shape[2], color.shape[1] |
| 49 | CX = intrinsics[0][2] |
| 50 | CY = intrinsics[1][2] |
| 51 | FX = intrinsics[0][0] |
| 52 | FY = intrinsics[1][1] |
| 53 | |
| 54 | # Compute indices of pixels |
| 55 | x_grid, y_grid = torch.meshgrid(torch.arange(width).cuda().float(), |
| 56 | torch.arange(height).cuda().float(), |
| 57 | indexing='xy') |
| 58 | xx = (x_grid - CX)/FX |
| 59 | yy = (y_grid - CY)/FY |
| 60 | xx = xx.reshape(-1) |
| 61 | yy = yy.reshape(-1) |
| 62 | depth_z = depth[0].reshape(-1) |
| 63 | |
| 64 | # Initialize point cloud |
| 65 | pts_cam = torch.stack((xx * depth_z, yy * depth_z, depth_z), dim=-1) |
| 66 | if transform_pts: |
| 67 | pix_ones = torch.ones(height * width, 1).cuda().float() |
| 68 | pts4 = torch.cat((pts_cam, pix_ones), dim=1) |
| 69 | c2w = torch.inverse(w2c) |
| 70 | pts = (c2w @ pts4.T).T[:, :3] |
| 71 | else: |
| 72 | pts = pts_cam |
| 73 | |
| 74 | # Compute mean squared distance for initializing the scale of the Gaussians |
| 75 | if compute_mean_sq_dist: |
| 76 | if mean_sq_dist_method == "projective": |
| 77 | # Projective Geometry (this is fast, farther -> larger radius) |
| 78 | scale_gaussian = depth_z / ((FX + FY)/2) |
| 79 | mean3_sq_dist = scale_gaussian**2 |
| 80 | else: |
| 81 | raise ValueError(f"Unknown mean_sq_dist_method {mean_sq_dist_method}") |
| 82 | |
| 83 | cols = torch.permute(color, (1, 2, 0)).reshape(-1, 3) # (C, H, W) -> (H, W, C) -> (H * W, C) |
| 84 | feat_dim = sem_feature.shape[0] |
| 85 | semantic = torch.permute(sem_feature, (1, 2, 0)).reshape(-1, feat_dim) |
| 86 | point_cld = torch.cat((pts, cols, semantic), -1) |
| 87 | |
| 88 | # Select points based on mask |
| 89 | if mask is not None: |
| 90 | point_cld = point_cld[mask] |
| 91 | if compute_mean_sq_dist: |
| 92 | mean3_sq_dist = mean3_sq_dist[mask] |
| 93 | |
| 94 | if compute_mean_sq_dist: |
| 95 | return point_cld, mean3_sq_dist |
| 96 | else: |
| 97 | return point_cld |
| 98 | |
| 99 | |
| 100 | def initialize_params(init_pt_cld, num_frames, mean3_sq_dist): |
no outgoing calls
no test coverage detected