(
exports: ModelExports,
frame_paths: list[Path],
uncropped_exports_shape: tuple[int, int],
uncropped_videos: Float[Tensor, "batch frame 3 uncropped_height uncropped_width"],
path: Path,
)
| 54 | |
| 55 | |
| 56 | def export_to_colmap( |
| 57 | exports: ModelExports, |
| 58 | frame_paths: list[Path], |
| 59 | uncropped_exports_shape: tuple[int, int], |
| 60 | uncropped_videos: Float[Tensor, "batch frame 3 uncropped_height uncropped_width"], |
| 61 | path: Path, |
| 62 | ) -> None: |
| 63 | # Account for the cropping that FlowMap does during optimization. |
| 64 | _, _, h_cropped, w_cropped = exports.depths.shape |
| 65 | h_uncropped, w_uncropped = uncropped_exports_shape |
| 66 | intrinsics = center_crop_intrinsics( |
| 67 | exports.intrinsics, |
| 68 | (h_cropped, w_cropped), |
| 69 | (h_uncropped, w_uncropped), |
| 70 | ) |
| 71 | |
| 72 | # Write out the camera parameters. |
| 73 | sparse_path = path / "sparse/0" |
| 74 | _, _, _, h_full, w_full = uncropped_videos.shape |
| 75 | write_colmap_model( |
| 76 | sparse_path, |
| 77 | exports.extrinsics[0], |
| 78 | intrinsics[0], |
| 79 | [path.name for path in frame_paths], |
| 80 | (h_full, w_full), |
| 81 | ) |
| 82 | |
| 83 | # Define the point cloud. For compatibility with 3D Gaussian Splatting, this is |
| 84 | # stored as a .ply instead of Points3D, which seems to be intended for a much |
| 85 | # smaller number of points. |
| 86 | _, _, dh, dw = exports.depths.shape |
| 87 | xy, _ = sample_image_grid((dh, dw), exports.extrinsics.device) |
| 88 | bundle = zip( |
| 89 | exports.extrinsics[0], |
| 90 | exports.intrinsics[0], |
| 91 | exports.depths[0], |
| 92 | exports.colors[0], |
| 93 | ) |
| 94 | points = [] |
| 95 | colors = [] |
| 96 | for extrinsics, intrinsics, depths, rgb in bundle: |
| 97 | xyz = unproject(xy, depths, intrinsics) |
| 98 | xyz = homogenize_points(xyz) |
| 99 | xyz = einsum(extrinsics, xyz, "i j, ... j -> ... i")[..., :3] |
| 100 | points.append(rearrange(xyz, "h w xyz -> (h w) xyz").detach().cpu().numpy()) |
| 101 | colors.append(rearrange(rgb, "c h w -> (h w) c").detach().cpu().numpy()) |
| 102 | points = np.concatenate(points) |
| 103 | colors = np.concatenate(colors) |
| 104 | |
| 105 | sparse_path.mkdir(parents=True, exist_ok=True) |
| 106 | write_ply(sparse_path / "points3D.ply", points, colors) |
| 107 | |
| 108 | # Write out the images. |
| 109 | (path / "images").mkdir(exist_ok=True, parents=True) |
| 110 | for frame_path in frame_paths: |
| 111 | shutil.copy(frame_path, path / "images" / frame_path.name) |
| 112 | |
| 113 |
no test coverage detected