beverley-gorry/colmap-vslamlab
0
1# Copyright (c), ETH Zurich and UNC Chapel Hill.2# All rights reserved.3#4# Redistribution and use in source and binary forms, with or without5# modification, are permitted provided that the following conditions are met:6#7# * Redistributions of source code must retain the above copyright8# notice, this list of conditions and the following disclaimer.9#10# * Redistributions in binary form must reproduce the above copyright11# notice, this list of conditions and the following disclaimer in the12# documentation and/or other materials provided with the distribution.13#14# * Neither the name of ETH Zurich and UNC Chapel Hill nor the names of15# its contributors may be used to endorse or promote products derived16# from this software without specific prior written permission.17#18# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"19# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE20# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE21# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS BE22# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR23# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF24# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS25# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN26# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)27# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE28# POSSIBILITY OF SUCH DAMAGE.29 30 31import argparse32import collections33import os34import struct35 36import numpy as np37 38CameraModel = collections.namedtuple(39 "CameraModel", ["model_id", "model_name", "num_params"]40)41Camera = collections.namedtuple(42 "Camera", ["id", "model", "width", "height", "params"]43)44BaseImage = collections.namedtuple(45 "Image", ["id", "qvec", "tvec", "camera_id", "name", "xys", "point3D_ids"]46)47Point3D = collections.namedtuple(48 "Point3D", ["id", "xyz", "rgb", "error", "image_ids", "point2D_idxs"]49)50 51 52class Image(BaseImage):53 def qvec2rotmat(self):54 return qvec2rotmat(self.qvec)55 56 57CAMERA_MODELS = {58 CameraModel(model_id=0, model_name="SIMPLE_PINHOLE", num_params=3),59 CameraModel(model_id=1, model_name="PINHOLE", num_params=4),60 CameraModel(model_id=2, model_name="SIMPLE_RADIAL", num_params=4),61 CameraModel(model_id=3, model_name="RADIAL", num_params=5),62 CameraModel(model_id=4, model_name="OPENCV", num_params=8),63 CameraModel(model_id=5, model_name="OPENCV_FISHEYE", num_params=8),64 CameraModel(model_id=6, model_name="FULL_OPENCV", num_params=12),65 CameraModel(model_id=7, model_name="FOV", num_params=5),66 CameraModel(model_id=8, model_name="SIMPLE_RADIAL_FISHEYE", num_params=4),67 CameraModel(model_id=9, model_name="RADIAL_FISHEYE", num_params=5),68 CameraModel(model_id=10, model_name="THIN_PRISM_FISHEYE", num_params=12),69}70CAMERA_MODEL_IDS = dict(71 [(camera_model.model_id, camera_model) for camera_model in CAMERA_MODELS]72)73CAMERA_MODEL_NAMES = dict(74 [(camera_model.model_name, camera_model) for camera_model in CAMERA_MODELS]75)76 77 78def read_next_bytes(fid, num_bytes, format_char_sequence, endian_character="<"):79 """Read and unpack the next bytes from a binary file.80 :param fid:81 :param num_bytes: Sum of combination of {2, 4, 8}, e.g. 2, 6, 16, 30, etc.82 :param format_char_sequence: List of {c, e, f, d, h, H, i, I, l, L, q, Q}.83 :param endian_character: Any of {@, =, <, >, !}84 :return: Tuple of read and unpacked values.85 """86 data = fid.read(num_bytes)87 return struct.unpack(endian_character + format_char_sequence, data)88 89 90def write_next_bytes(fid, data, format_char_sequence, endian_character="<"):91 """pack and write to a binary file.92 :param fid:93 :param data: data to send, if multiple elements are sent at the same time,94 they should be encapsuled either in a list or a tuple95 :param format_char_sequence: List of {c, e, f, d, h, H, i, I, l, L, q, Q}.96 should be the same length as the data list or tuple97 :param endian_character: Any of {@, =, <, >, !}98 """99 if isinstance(data, (list, tuple)):100 bytes = struct.pack(endian_character + format_char_sequence, *data)101 else:102 bytes = struct.pack(endian_character + format_char_sequence, data)103 fid.write(bytes)104 105 106def read_cameras_text(path):107 """108 see: src/colmap/scene/reconstruction.cc109 void Reconstruction::WriteCamerasText(const std::string& path)110 void Reconstruction::ReadCamerasText(const std::string& path)111 """112 cameras = {}113 with open(path, "r") as fid:114 while True:115 line = fid.readline()116 if not line:117 break118 line = line.strip()119 if len(line) > 0 and line[0] != "#":120 elems = line.split()121 camera_id = int(elems[0])122 model = elems[1]123 width = int(elems[2])124 height = int(elems[3])125 params = np.array(tuple(map(float, elems[4:])))126 cameras[camera_id] = Camera(127 id=camera_id,128 model=model,129 width=width,130 height=height,131 params=params,132 )133 return cameras134 135 136def read_cameras_binary(path_to_model_file):137 """138 see: src/colmap/scene/reconstruction.cc139 void Reconstruction::WriteCamerasBinary(const std::string& path)140 void Reconstruction::ReadCamerasBinary(const std::string& path)141 """142 cameras = {}143 with open(path_to_model_file, "rb") as fid:144 num_cameras = read_next_bytes(fid, 8, "Q")[0]145 for _ in range(num_cameras):146 camera_properties = read_next_bytes(147 fid, num_bytes=24, format_char_sequence="iiQQ"148 )149 camera_id = camera_properties[0]150 model_id = camera_properties[1]151 model_name = CAMERA_MODEL_IDS[camera_properties[1]].model_name152 width = camera_properties[2]153 height = camera_properties[3]154 num_params = CAMERA_MODEL_IDS[model_id].num_params155 params = read_next_bytes(156 fid,157 num_bytes=8 * num_params,158 format_char_sequence="d" * num_params,159 )160 cameras[camera_id] = Camera(161 id=camera_id,162 model=model_name,163 width=width,164 height=height,165 params=np.array(params),166 )167 assert len(cameras) == num_cameras168 return cameras169 170 171def write_cameras_text(cameras, path):172 """173 see: src/colmap/scene/reconstruction.cc174 void Reconstruction::WriteCamerasText(const std::string& path)175 void Reconstruction::ReadCamerasText(const std::string& path)176 """177 HEADER = (178 "# Camera list with one line of data per camera:\n"179 + "# CAMERA_ID, MODEL, WIDTH, HEIGHT, PARAMS[]\n"180 + "# Number of cameras: {}\n".format(len(cameras))181 )182 with open(path, "w") as fid:183 fid.write(HEADER)184 for _, cam in cameras.items():185 to_write = [cam.id, cam.model, cam.width, cam.height, *cam.params]186 line = " ".join([str(elem) for elem in to_write])187 fid.write(line + "\n")188 189 190def write_cameras_binary(cameras, path_to_model_file):191 """192 see: src/colmap/scene/reconstruction.cc193 void Reconstruction::WriteCamerasBinary(const std::string& path)194 void Reconstruction::ReadCamerasBinary(const std::string& path)195 """196 with open(path_to_model_file, "wb") as fid:197 write_next_bytes(fid, len(cameras), "Q")198 for _, cam in cameras.items():199 model_id = CAMERA_MODEL_NAMES[cam.model].model_id200 camera_properties = [cam.id, model_id, cam.width, cam.height]201 write_next_bytes(fid, camera_properties, "iiQQ")202 for p in cam.params:203 write_next_bytes(fid, float(p), "d")204 return cameras205 206 207def read_images_text(path):208 """209 see: src/colmap/scene/reconstruction.cc210 void Reconstruction::ReadImagesText(const std::string& path)211 void Reconstruction::WriteImagesText(const std::string& path)212 """213 images = {}214 with open(path, "r") as fid:215 while True:216 line = fid.readline()217 if not line:218 break219 line = line.strip()220 if len(line) > 0 and line[0] != "#":221 elems = line.split()222 image_id = int(elems[0])223 qvec = np.array(tuple(map(float, elems[1:5])))224 tvec = np.array(tuple(map(float, elems[5:8])))225 camera_id = int(elems[8])226 image_name = elems[9]227 elems = fid.readline().split()228 xys = np.column_stack(229 [230 tuple(map(float, elems[0::3])),231 tuple(map(float, elems[1::3])),232 ]233 )234 point3D_ids = np.array(tuple(map(int, elems[2::3])))235 images[image_id] = Image(236 id=image_id,237 qvec=qvec,238 tvec=tvec,239 camera_id=camera_id,240 name=image_name,241 xys=xys,242 point3D_ids=point3D_ids,243 )244 return images245 246 247def read_images_binary(path_to_model_file):248 """249 see: src/colmap/scene/reconstruction.cc250 void Reconstruction::ReadImagesBinary(const std::string& path)251 void Reconstruction::WriteImagesBinary(const std::string& path)252 """253 images = {}254 with open(path_to_model_file, "rb") as fid:255 num_reg_images = read_next_bytes(fid, 8, "Q")[0]256 for _ in range(num_reg_images):257 binary_image_properties = read_next_bytes(258 fid, num_bytes=64, format_char_sequence="idddddddi"259 )260 image_id = binary_image_properties[0]261 qvec = np.array(binary_image_properties[1:5])262 tvec = np.array(binary_image_properties[5:8])263 camera_id = binary_image_properties[8]264 binary_image_name = b""265 current_char = read_next_bytes(fid, 1, "c")[0]266 while current_char != b"\x00": # look for the ASCII 0 entry267 binary_image_name += current_char268 current_char = read_next_bytes(fid, 1, "c")[0]269 image_name = binary_image_name.decode("utf-8")270 num_points2D = read_next_bytes(271 fid, num_bytes=8, format_char_sequence="Q"272 )[0]273 x_y_id_s = read_next_bytes(274 fid,275 num_bytes=24 * num_points2D,276 format_char_sequence="ddq" * num_points2D,277 )278 xys = np.column_stack(279 [280 tuple(map(float, x_y_id_s[0::3])),281 tuple(map(float, x_y_id_s[1::3])),282 ]283 )284 point3D_ids = np.array(tuple(map(int, x_y_id_s[2::3])))285 images[image_id] = Image(286 id=image_id,287 qvec=qvec,288 tvec=tvec,289 camera_id=camera_id,290 name=image_name,291 xys=xys,292 point3D_ids=point3D_ids,293 )294 return images295 296 297def write_images_text(images, path):298 """299 see: src/colmap/scene/reconstruction.cc300 void Reconstruction::ReadImagesText(const std::string& path)301 void Reconstruction::WriteImagesText(const std::string& path)302 """303 if len(images) == 0:304 mean_observations = 0305 else:306 mean_observations = sum(307 (len(img.point3D_ids) for _, img in images.items())308 ) / len(images)309 HEADER = (310 "# Image list with two lines of data per image:\n"311 + "# IMAGE_ID, QW, QX, QY, QZ, TX, TY, TZ, CAMERA_ID, NAME\n"312 + "# POINTS2D[] as (X, Y, POINT3D_ID)\n"313 + "# Number of images: {}, mean observations per image: {}\n".format(314 len(images), mean_observations315 )316 )317 318 with open(path, "w") as fid:319 fid.write(HEADER)320 for _, img in images.items():321 image_header = [322 img.id,323 *img.qvec,324 *img.tvec,325 img.camera_id,326 img.name,327 ]328 first_line = " ".join(map(str, image_header))329 fid.write(first_line + "\n")330 331 points_strings = []332 for xy, point3D_id in zip(img.xys, img.point3D_ids):333 points_strings.append(" ".join(map(str, [*xy, point3D_id])))334 fid.write(" ".join(points_strings) + "\n")335 336 337def write_images_binary(images, path_to_model_file):338 """339 see: src/colmap/scene/reconstruction.cc340 void Reconstruction::ReadImagesBinary(const std::string& path)341 void Reconstruction::WriteImagesBinary(const std::string& path)342 """343 with open(path_to_model_file, "wb") as fid:344 write_next_bytes(fid, len(images), "Q")345 for _, img in images.items():346 write_next_bytes(fid, img.id, "i")347 write_next_bytes(fid, img.qvec.tolist(), "dddd")348 write_next_bytes(fid, img.tvec.tolist(), "ddd")349 write_next_bytes(fid, img.camera_id, "i")350 for char in img.name:351 write_next_bytes(fid, char.encode("utf-8"), "c")352 write_next_bytes(fid, b"\x00", "c")353 write_next_bytes(fid, len(img.point3D_ids), "Q")354 for xy, p3d_id in zip(img.xys, img.point3D_ids):355 write_next_bytes(fid, [*xy, p3d_id], "ddq")356 357 358def read_points3D_text(path):359 """360 see: src/colmap/scene/reconstruction.cc361 void Reconstruction::ReadPoints3DText(const std::string& path)362 void Reconstruction::WritePoints3DText(const std::string& path)363 """364 points3D = {}365 with open(path, "r") as fid:366 while True:367 line = fid.readline()368 if not line:369 break370 line = line.strip()371 if len(line) > 0 and line[0] != "#":372 elems = line.split()373 point3D_id = int(elems[0])374 xyz = np.array(tuple(map(float, elems[1:4])))375 rgb = np.array(tuple(map(int, elems[4:7])))376 error = float(elems[7])377 image_ids = np.array(tuple(map(int, elems[8::2])))378 point2D_idxs = np.array(tuple(map(int, elems[9::2])))379 points3D[point3D_id] = Point3D(380 id=point3D_id,381 xyz=xyz,382 rgb=rgb,383 error=error,384 image_ids=image_ids,385 point2D_idxs=point2D_idxs,386 )387 return points3D388 389 390def read_points3D_binary(path_to_model_file):391 """392 see: src/colmap/scene/reconstruction.cc393 void Reconstruction::ReadPoints3DBinary(const std::string& path)394 void Reconstruction::WritePoints3DBinary(const std::string& path)395 """396 points3D = {}397 with open(path_to_model_file, "rb") as fid:398 num_points = read_next_bytes(fid, 8, "Q")[0]399 for _ in range(num_points):400 binary_point_line_properties = read_next_bytes(401 fid, num_bytes=43, format_char_sequence="QdddBBBd"402 )403 point3D_id = binary_point_line_properties[0]404 xyz = np.array(binary_point_line_properties[1:4])405 rgb = np.array(binary_point_line_properties[4:7])406 error = np.array(binary_point_line_properties[7])407 track_length = read_next_bytes(408 fid, num_bytes=8, format_char_sequence="Q"409 )[0]410 track_elems = read_next_bytes(411 fid,412 num_bytes=8 * track_length,413 format_char_sequence="ii" * track_length,414 )415 image_ids = np.array(tuple(map(int, track_elems[0::2])))416 point2D_idxs = np.array(tuple(map(int, track_elems[1::2])))417 points3D[point3D_id] = Point3D(418 id=point3D_id,419 xyz=xyz,420 rgb=rgb,421 error=error,422 image_ids=image_ids,423 point2D_idxs=point2D_idxs,424 )425 return points3D426 427 428def write_points3D_text(points3D, path):429 """430 see: src/colmap/scene/reconstruction.cc431 void Reconstruction::ReadPoints3DText(const std::string& path)432 void Reconstruction::WritePoints3DText(const std::string& path)433 """434 if len(points3D) == 0:435 mean_track_length = 0436 else:437 mean_track_length = sum(438 (len(pt.image_ids) for _, pt in points3D.items())439 ) / len(points3D)440 HEADER = (441 "# 3D point list with one line of data per point:\n"442 + "# POINT3D_ID, X, Y, Z, R, G, B, ERROR, TRACK[] as (IMAGE_ID, POINT2D_IDX)\n"443 + "# Number of points: {}, mean track length: {}\n".format(444 len(points3D), mean_track_length445 )446 )447 448 with open(path, "w") as fid:449 fid.write(HEADER)450 for _, pt in points3D.items():451 point_header = [pt.id, *pt.xyz, *pt.rgb, pt.error]452 fid.write(" ".join(map(str, point_header)) + " ")453 track_strings = []454 for image_id, point2D in zip(pt.image_ids, pt.point2D_idxs):455 track_strings.append(" ".join(map(str, [image_id, point2D])))456 fid.write(" ".join(track_strings) + "\n")457 458 459def write_points3D_binary(points3D, path_to_model_file):460 """461 see: src/colmap/scene/reconstruction.cc462 void Reconstruction::ReadPoints3DBinary(const std::string& path)463 void Reconstruction::WritePoints3DBinary(const std::string& path)464 """465 with open(path_to_model_file, "wb") as fid:466 write_next_bytes(fid, len(points3D), "Q")467 for _, pt in points3D.items():468 write_next_bytes(fid, pt.id, "Q")469 write_next_bytes(fid, pt.xyz.tolist(), "ddd")470 write_next_bytes(fid, pt.rgb.tolist(), "BBB")471 write_next_bytes(fid, pt.error, "d")472 track_length = pt.image_ids.shape[0]473 write_next_bytes(fid, track_length, "Q")474 for image_id, point2D_id in zip(pt.image_ids, pt.point2D_idxs):475 write_next_bytes(fid, [image_id, point2D_id], "ii")476 477 478def detect_model_format(path, ext):479 if (480 os.path.isfile(os.path.join(path, "cameras" + ext))481 and os.path.isfile(os.path.join(path, "images" + ext))482 and os.path.isfile(os.path.join(path, "points3D" + ext))483 ):484 print("Detected model format: '" + ext + "'")485 return True486 487 return False488 489 490def read_model(path, ext=""):491 # try to detect the extension automatically492 if ext == "":493 if detect_model_format(path, ".bin"):494 ext = ".bin"495 elif detect_model_format(path, ".txt"):496 ext = ".txt"497 else:498 print("Provide model format: '.bin' or '.txt'")499 return500 501 if ext == ".txt":502 cameras = read_cameras_text(os.path.join(path, "cameras" + ext))503 images = read_images_text(os.path.join(path, "images" + ext))504 points3D = read_points3D_text(os.path.join(path, "points3D") + ext)505 else:506 cameras = read_cameras_binary(os.path.join(path, "cameras" + ext))507 images = read_images_binary(os.path.join(path, "images" + ext))508 points3D = read_points3D_binary(os.path.join(path, "points3D") + ext)509 return cameras, images, points3D510 511 512def write_model(cameras, images, points3D, path, ext=".bin"):513 if ext == ".txt":514 write_cameras_text(cameras, os.path.join(path, "cameras" + ext))515 write_images_text(images, os.path.join(path, "images" + ext))516 write_points3D_text(points3D, os.path.join(path, "points3D") + ext)517 else:518 write_cameras_binary(cameras, os.path.join(path, "cameras" + ext))519 write_images_binary(images, os.path.join(path, "images" + ext))520 write_points3D_binary(points3D, os.path.join(path, "points3D") + ext)521 return cameras, images, points3D522 523 524def qvec2rotmat(qvec):525 return np.array(526 [527 [528 1 - 2 * qvec[2] ** 2 - 2 * qvec[3] ** 2,529 2 * qvec[1] * qvec[2] - 2 * qvec[0] * qvec[3],530 2 * qvec[3] * qvec[1] + 2 * qvec[0] * qvec[2],531 ],532 [533 2 * qvec[1] * qvec[2] + 2 * qvec[0] * qvec[3],534 1 - 2 * qvec[1] ** 2 - 2 * qvec[3] ** 2,535 2 * qvec[2] * qvec[3] - 2 * qvec[0] * qvec[1],536 ],537 [538 2 * qvec[3] * qvec[1] - 2 * qvec[0] * qvec[2],539 2 * qvec[2] * qvec[3] + 2 * qvec[0] * qvec[1],540 1 - 2 * qvec[1] ** 2 - 2 * qvec[2] ** 2,541 ],542 ]543 )544 545 546def rotmat2qvec(R):547 Rxx, Ryx, Rzx, Rxy, Ryy, Rzy, Rxz, Ryz, Rzz = R.flat548 K = (549 np.array(550 [551 [Rxx - Ryy - Rzz, 0, 0, 0],552 [Ryx + Rxy, Ryy - Rxx - Rzz, 0, 0],553 [Rzx + Rxz, Rzy + Ryz, Rzz - Rxx - Ryy, 0],554 [Ryz - Rzy, Rzx - Rxz, Rxy - Ryx, Rxx + Ryy + Rzz],555 ]556 )557 / 3.0558 )559 eigvals, eigvecs = np.linalg.eigh(K)560 qvec = eigvecs[[3, 0, 1, 2], np.argmax(eigvals)]561 if qvec[0] < 0:562 qvec *= -1563 return qvec564 565 566def main():567 parser = argparse.ArgumentParser(568 description="Read and write COLMAP binary and text models"569 )570 parser.add_argument("--input_model", help="path to input model folder")571 parser.add_argument(572 "--input_format",573 choices=[".bin", ".txt"],574 help="input model format",575 default="",576 )577 parser.add_argument("--output_model", help="path to output model folder")578 parser.add_argument(579 "--output_format",580 choices=[".bin", ".txt"],581 help="output model format",582 default=".txt",583 )584 args = parser.parse_args()585 586 cameras, images, points3D = read_model(587 path=args.input_model, ext=args.input_format588 )589 590 print("num_cameras:", len(cameras))591 print("num_images:", len(images))592 print("num_points3D:", len(points3D))593 594 if args.output_model is not None:595 write_model(596 cameras,597 images,598 points3D,599 path=args.output_model,600 ext=args.output_format,601 )602 603 604if __name__ == "__main__":605 main()606