| 7 | |
| 8 | |
| 9 | class MeshExtractor: |
| 10 | def __init__(self, args): |
| 11 | self.voxel_size = args.mapper_specs["voxel_size"] # 0.2 |
| 12 | self.rays_d = None |
| 13 | self.depth_points = None |
| 14 | self.model = args.decoder |
| 15 | |
| 16 | @torch.no_grad() |
| 17 | def linearize_id(self, xyz, n_xyz): |
| 18 | return xyz[:, 2] + n_xyz[-1] * xyz[:, 1] + (n_xyz[-1] * n_xyz[-2]) * xyz[:, 0] |
| 19 | |
| 20 | @torch.no_grad() |
| 21 | def downsample_points(self, points, voxel_size=0.01): |
| 22 | pcd = o3d.geometry.PointCloud() |
| 23 | pcd.points = o3d.utility.Vector3dVector(points) |
| 24 | pcd = pcd.voxel_down_sample(voxel_size) |
| 25 | return np.asarray(pcd.points) |
| 26 | |
| 27 | @torch.no_grad() |
| 28 | def get_rays(self, w=None, h=None, K=None): |
| 29 | w = self.w if w == None else w |
| 30 | h = self.h if h == None else h |
| 31 | if K is None: |
| 32 | K = np.eye(3) |
| 33 | K[0, 0] = self.K[0, 0] * w / self.w |
| 34 | K[1, 1] = self.K[1, 1] * h / self.h |
| 35 | K[0, 2] = self.K[0, 2] * w / self.w |
| 36 | K[1, 2] = self.K[1, 2] * h / self.h |
| 37 | ix, iy = torch.meshgrid( |
| 38 | torch.arange(w), torch.arange(h), indexing='xy') |
| 39 | rays_d = torch.stack( |
| 40 | [(ix - K[0, 2]) / K[0, 0], |
| 41 | (iy - K[1, 2]) / K[1, 1], |
| 42 | torch.ones_like(ix)], -1).float() |
| 43 | return rays_d |
| 44 | |
| 45 | @torch.no_grad() |
| 46 | def get_valid_points(self, frame_poses, depth_maps): |
| 47 | if isinstance(frame_poses, list): |
| 48 | all_points = [] |
| 49 | print("extracting all points") |
| 50 | for i in range(0, len(frame_poses), 5): |
| 51 | pose = frame_poses[i] |
| 52 | depth = depth_maps[i] |
| 53 | points = self.rays_d * depth.unsqueeze(-1) |
| 54 | points = points.reshape(-1, 3) |
| 55 | points = points @ pose[:3, :3].transpose(-1, -2) + pose[:3, 3] |
| 56 | if len(all_points) == 0: |
| 57 | all_points = points.detach().cpu().numpy() |
| 58 | else: |
| 59 | all_points = np.concatenate( |
| 60 | [all_points, points.detach().cpu().numpy()], 0) |
| 61 | print("downsample all points") |
| 62 | all_points = self.downsample_points(all_points) |
| 63 | return all_points |
| 64 | else: |
| 65 | pose = frame_poses |
| 66 | depth = depth_maps |