COLMAP visualizer#
Visualize COLMAP sparse reconstruction outputs. To get demo data, see ./assets/download_colmap_garden.sh
.
1import random
2import time
3from pathlib import Path
4
5import imageio.v3 as iio
6import numpy as onp
7import tyro
8import viser
9import viser.transforms as tf
10from tqdm.auto import tqdm
11from viser.extras.colmap import (
12 read_cameras_binary,
13 read_images_binary,
14 read_points3d_binary,
15)
16
17
18def main(
19 colmap_path: Path = Path(__file__).parent / "assets/colmap_garden/sparse/0",
20 images_path: Path = Path(__file__).parent / "assets/colmap_garden/images_8",
21 downsample_factor: int = 2,
22) -> None:
23 """Visualize COLMAP sparse reconstruction outputs.
24
25 Args:
26 colmap_path: Path to the COLMAP reconstruction directory.
27 images_path: Path to the COLMAP images directory.
28 downsample_factor: Downsample factor for the images.
29 """
30 server = viser.ViserServer()
31 server.gui.configure_theme(titlebar_content=None, control_layout="collapsible")
32
33 # Load the colmap info.
34 cameras = read_cameras_binary(colmap_path / "cameras.bin")
35 images = read_images_binary(colmap_path / "images.bin")
36 points3d = read_points3d_binary(colmap_path / "points3D.bin")
37 gui_reset_up = server.gui.add_button(
38 "Reset up direction",
39 hint="Set the camera control 'up' direction to the current camera's 'up'.",
40 )
41
42 @gui_reset_up.on_click
43 def _(event: viser.GuiEvent) -> None:
44 client = event.client
45 assert client is not None
46 client.camera.up_direction = tf.SO3(client.camera.wxyz) @ onp.array(
47 [0.0, -1.0, 0.0]
48 )
49
50 gui_points = server.gui.add_slider(
51 "Max points",
52 min=1,
53 max=len(points3d),
54 step=1,
55 initial_value=min(len(points3d), 50_000),
56 )
57 gui_frames = server.gui.add_slider(
58 "Max frames",
59 min=1,
60 max=len(images),
61 step=1,
62 initial_value=min(len(images), 100),
63 )
64 gui_point_size = server.gui.add_number("Point size", initial_value=0.05)
65
66 def visualize_colmap() -> None:
67 """Send all COLMAP elements to viser for visualization. This could be optimized
68 a ton!"""
69 # Set the point cloud.
70 points = onp.array([points3d[p_id].xyz for p_id in points3d])
71 colors = onp.array([points3d[p_id].rgb for p_id in points3d])
72 points_selection = onp.random.choice(
73 points.shape[0], gui_points.value, replace=False
74 )
75 points = points[points_selection]
76 colors = colors[points_selection]
77
78 server.scene.add_point_cloud(
79 name="/colmap/pcd",
80 points=points,
81 colors=colors,
82 point_size=gui_point_size.value,
83 )
84
85 # Interpret the images and cameras.
86 img_ids = [im.id for im in images.values()]
87 random.shuffle(img_ids)
88 img_ids = sorted(img_ids[: gui_frames.value])
89
90 def attach_callback(
91 frustum: viser.CameraFrustumHandle, frame: viser.FrameHandle
92 ) -> None:
93 @frustum.on_click
94 def _(_) -> None:
95 for client in server.get_clients().values():
96 client.camera.wxyz = frame.wxyz
97 client.camera.position = frame.position
98
99 for img_id in tqdm(img_ids):
100 img = images[img_id]
101 cam = cameras[img.camera_id]
102
103 # Skip images that don't exist.
104 image_filename = images_path / img.name
105 if not image_filename.exists():
106 continue
107
108 T_world_camera = tf.SE3.from_rotation_and_translation(
109 tf.SO3(img.qvec), img.tvec
110 ).inverse()
111 frame = server.scene.add_frame(
112 f"/colmap/frame_{img_id}",
113 wxyz=T_world_camera.rotation().wxyz,
114 position=T_world_camera.translation(),
115 axes_length=0.1,
116 axes_radius=0.005,
117 )
118
119 # For pinhole cameras, cam.params will be (fx, fy, cx, cy).
120 if cam.model != "PINHOLE":
121 print(f"Expected pinhole camera, but got {cam.model}")
122
123 H, W = cam.height, cam.width
124 fy = cam.params[1]
125 image = iio.imread(image_filename)
126 image = image[::downsample_factor, ::downsample_factor]
127 frustum = server.scene.add_camera_frustum(
128 f"/colmap/frame_{img_id}/frustum",
129 fov=2 * onp.arctan2(H / 2, fy),
130 aspect=W / H,
131 scale=0.15,
132 image=image,
133 )
134 attach_callback(frustum, frame)
135
136 need_update = True
137
138 @gui_points.on_update
139 def _(_) -> None:
140 nonlocal need_update
141 need_update = True
142
143 @gui_frames.on_update
144 def _(_) -> None:
145 nonlocal need_update
146 need_update = True
147
148 @gui_point_size.on_update
149 def _(_) -> None:
150 nonlocal need_update
151 need_update = True
152
153 while True:
154 if need_update:
155 need_update = False
156
157 server.scene.reset()
158 visualize_colmap()
159
160 time.sleep(1e-3)
161
162
163if __name__ == "__main__":
164 tyro.cli(main)