#!/usr/bin/env python3 """Script to convert COLMAP data into Nerfstudio format.""" # Largely copied from: # - https://github.com/nerfstudio-project/nerfstudio/blob/main/nerfstudio/process_data/colmap_utils.py # - https://github.com/nerfstudio-project/nerfstudio/blob/main/nerfstudio/process_data/colmap_converter_to_nerfstudio_dataset.py # - https://github.com/nerfstudio-project/nerfstudio/blob/main/nerfstudio/data/utils/colmap_parsing_utils.py # Added support for THIN_PRISM_FISHEYE import json from pathlib import Path from typing import Any, Dict, List, Tuple, Optional, Union import numpy as np import torch from dataclasses import dataclass import collections import struct # Copyright (c) 2023, ETH Zurich and UNC Chapel Hill. # All rights reserved. # # Redistribution and use in source and binary forms, with or without # modification, are permitted provided that the following conditions are met: # # * Redistributions of source code must retain the above copyright # notice, this list of conditions and the following disclaimer. # # * Redistributions in binary form must reproduce the above copyright # notice, this list of conditions and the following disclaimer in the # documentation and/or other materials provided with the distribution. # # * Neither the name of ETH Zurich and UNC Chapel Hill nor the names of # its contributors may be used to endorse or promote products derived # from this software without specific prior written permission. # # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" # AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE # ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS BE # LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR # CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF # SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS # INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN # CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) # ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE # POSSIBILITY OF SUCH DAMAGE. # # Author: Johannes L. Schoenberger (jsch-at-demuc-dot-de) CameraModel = collections.namedtuple("CameraModel", ["model_id", "model_name", "num_params"]) Camera = collections.namedtuple("Camera", ["id", "model", "width", "height", "params"]) BaseImage = collections.namedtuple("Image", ["id", "qvec", "tvec", "camera_id", "name", "xys", "point3D_ids"]) Point3D = collections.namedtuple("Point3D", ["id", "xyz", "rgb", "error", "image_ids", "point2D_idxs"]) def qvec2rotmat(qvec): return np.array( [ [ 1 - 2 * qvec[2] ** 2 - 2 * qvec[3] ** 2, 2 * qvec[1] * qvec[2] - 2 * qvec[0] * qvec[3], 2 * qvec[3] * qvec[1] + 2 * qvec[0] * qvec[2], ], [ 2 * qvec[1] * qvec[2] + 2 * qvec[0] * qvec[3], 1 - 2 * qvec[1] ** 2 - 2 * qvec[3] ** 2, 2 * qvec[2] * qvec[3] - 2 * qvec[0] * qvec[1], ], [ 2 * qvec[3] * qvec[1] - 2 * qvec[0] * qvec[2], 2 * qvec[2] * qvec[3] + 2 * qvec[0] * qvec[1], 1 - 2 * qvec[1] ** 2 - 2 * qvec[2] ** 2, ], ] ) class Image(BaseImage): def qvec2rotmat(self): return qvec2rotmat(self.qvec) CAMERA_MODELS = { CameraModel(model_id=0, model_name="SIMPLE_PINHOLE", num_params=3), CameraModel(model_id=1, model_name="PINHOLE", num_params=4), CameraModel(model_id=2, model_name="SIMPLE_RADIAL", num_params=4), CameraModel(model_id=3, model_name="RADIAL", num_params=5), CameraModel(model_id=4, model_name="OPENCV", num_params=8), CameraModel(model_id=5, model_name="OPENCV_FISHEYE", num_params=8), CameraModel(model_id=6, model_name="FULL_OPENCV", num_params=12), CameraModel(model_id=7, model_name="FOV", num_params=5), CameraModel(model_id=8, model_name="SIMPLE_RADIAL_FISHEYE", num_params=4), CameraModel(model_id=9, model_name="RADIAL_FISHEYE", num_params=5), CameraModel(model_id=10, model_name="THIN_PRISM_FISHEYE", num_params=12), } CAMERA_MODEL_IDS = dict([(camera_model.model_id, camera_model) for camera_model in CAMERA_MODELS]) CAMERA_MODEL_NAMES = dict([(camera_model.model_name, camera_model) for camera_model in CAMERA_MODELS]) def read_next_bytes(fid, num_bytes, format_char_sequence, endian_character="<"): """Read and unpack the next bytes from a binary file. :param fid: :param num_bytes: Sum of combination of {2, 4, 8}, e.g. 2, 6, 16, 30, etc. :param format_char_sequence: List of {c, e, f, d, h, H, i, I, l, L, q, Q}. :param endian_character: Any of {@, =, <, >, !} :return: Tuple of read and unpacked values. """ data = fid.read(num_bytes) return struct.unpack(endian_character + format_char_sequence, data) def write_next_bytes(fid, data, format_char_sequence, endian_character="<"): """pack and write to a binary file. :param fid: :param data: data to send, if multiple elements are sent at the same time, they should be encapsuled either in a list or a tuple :param format_char_sequence: List of {c, e, f, d, h, H, i, I, l, L, q, Q}. should be the same length as the data list or tuple :param endian_character: Any of {@, =, <, >, !} """ if isinstance(data, (list, tuple)): bytes = struct.pack(endian_character + format_char_sequence, *data) else: bytes = struct.pack(endian_character + format_char_sequence, data) fid.write(bytes) def read_cameras_text(path): """ see: src/base/reconstruction.cc void Reconstruction::WriteCamerasText(const std::string& path) void Reconstruction::ReadCamerasText(const std::string& path) """ cameras = {} with open(path, "r") as fid: while True: line = fid.readline() if not line: break line = line.strip() if len(line) > 0 and line[0] != "#": elems = line.split() camera_id = int(elems[0]) model = elems[1] width = int(elems[2]) height = int(elems[3]) params = np.array(tuple(map(float, elems[4:]))) cameras[camera_id] = Camera(id=camera_id, model=model, width=width, height=height, params=params) return cameras def read_cameras_binary(path_to_model_file): """ see: src/base/reconstruction.cc void Reconstruction::WriteCamerasBinary(const std::string& path) void Reconstruction::ReadCamerasBinary(const std::string& path) """ cameras = {} with open(path_to_model_file, "rb") as fid: num_cameras = read_next_bytes(fid, 8, "Q")[0] for _ in range(num_cameras): camera_properties = read_next_bytes(fid, num_bytes=24, format_char_sequence="iiQQ") camera_id = camera_properties[0] model_id = camera_properties[1] model_name = CAMERA_MODEL_IDS[camera_properties[1]].model_name width = camera_properties[2] height = camera_properties[3] num_params = CAMERA_MODEL_IDS[model_id].num_params params = read_next_bytes(fid, num_bytes=8 * num_params, format_char_sequence="d" * num_params) cameras[camera_id] = Camera( id=camera_id, model=model_name, width=width, height=height, params=np.array(params) ) assert len(cameras) == num_cameras return cameras def write_cameras_text(cameras, path): """ see: src/base/reconstruction.cc void Reconstruction::WriteCamerasText(const std::string& path) void Reconstruction::ReadCamerasText(const std::string& path) """ HEADER = ( "# Camera list with one line of data per camera:\n" + "# CAMERA_ID, MODEL, WIDTH, HEIGHT, PARAMS[]\n" + "# Number of cameras: {}\n".format(len(cameras)) ) with open(path, "w") as fid: fid.write(HEADER) for _, cam in cameras.items(): to_write = [cam.id, cam.model, cam.width, cam.height, *cam.params] line = " ".join([str(elem) for elem in to_write]) fid.write(line + "\n") def write_cameras_binary(cameras, path_to_model_file): """ see: src/base/reconstruction.cc void Reconstruction::WriteCamerasBinary(const std::string& path) void Reconstruction::ReadCamerasBinary(const std::string& path) """ with open(path_to_model_file, "wb") as fid: write_next_bytes(fid, len(cameras), "Q") for _, cam in cameras.items(): model_id = CAMERA_MODEL_NAMES[cam.model].model_id camera_properties = [cam.id, model_id, cam.width, cam.height] write_next_bytes(fid, camera_properties, "iiQQ") for p in cam.params: write_next_bytes(fid, float(p), "d") return cameras def read_images_text(path): """ see: src/base/reconstruction.cc void Reconstruction::ReadImagesText(const std::string& path) void Reconstruction::WriteImagesText(const std::string& path) """ images = {} with open(path, "r") as fid: while True: line = fid.readline() if not line: break line = line.strip() if len(line) > 0 and line[0] != "#": elems = line.split() image_id = int(elems[0]) qvec = np.array(tuple(map(float, elems[1:5]))) tvec = np.array(tuple(map(float, elems[5:8]))) camera_id = int(elems[8]) image_name = elems[9] elems = fid.readline().split() xys = np.column_stack([tuple(map(float, elems[0::3])), tuple(map(float, elems[1::3]))]) point3D_ids = np.array(tuple(map(int, elems[2::3]))) images[image_id] = Image( id=image_id, qvec=qvec, tvec=tvec, camera_id=camera_id, name=image_name, xys=xys, point3D_ids=point3D_ids, ) return images def read_images_binary(path_to_model_file): """ see: src/base/reconstruction.cc void Reconstruction::ReadImagesBinary(const std::string& path) void Reconstruction::WriteImagesBinary(const std::string& path) """ images = {} with open(path_to_model_file, "rb") as fid: num_reg_images = read_next_bytes(fid, 8, "Q")[0] for _ in range(num_reg_images): binary_image_properties = read_next_bytes(fid, num_bytes=64, format_char_sequence="idddddddi") image_id = binary_image_properties[0] qvec = np.array(binary_image_properties[1:5]) tvec = np.array(binary_image_properties[5:8]) camera_id = binary_image_properties[8] image_name = b"" current_char = read_next_bytes(fid, 1, "c")[0] while current_char != b"\x00": # look for the ASCII 0 entry image_name += current_char current_char = read_next_bytes(fid, 1, "c")[0] image_name = image_name.decode("utf-8") num_points2D = read_next_bytes(fid, num_bytes=8, format_char_sequence="Q")[0] x_y_id_s = read_next_bytes(fid, num_bytes=24 * num_points2D, format_char_sequence="ddq" * num_points2D) xys = np.column_stack([tuple(map(float, x_y_id_s[0::3])), tuple(map(float, x_y_id_s[1::3]))]) point3D_ids = np.array(tuple(map(int, x_y_id_s[2::3]))) images[image_id] = Image( id=image_id, qvec=qvec, tvec=tvec, camera_id=camera_id, name=image_name, xys=xys, point3D_ids=point3D_ids, ) return images def write_images_text(images, path): """ see: src/base/reconstruction.cc void Reconstruction::ReadImagesText(const std::string& path) void Reconstruction::WriteImagesText(const std::string& path) """ if len(images) == 0: mean_observations = 0 else: mean_observations = sum((len(img.point3D_ids) for _, img in images.items())) / len(images) HEADER = ( "# Image list with two lines of data per image:\n" + "# IMAGE_ID, QW, QX, QY, QZ, TX, TY, TZ, CAMERA_ID, NAME\n" + "# POINTS2D[] as (X, Y, POINT3D_ID)\n" + "# Number of images: {}, mean observations per image: {}\n".format(len(images), mean_observations) ) with open(path, "w") as fid: fid.write(HEADER) for _, img in images.items(): image_header = [img.id, *img.qvec, *img.tvec, img.camera_id, img.name] first_line = " ".join(map(str, image_header)) fid.write(first_line + "\n") points_strings = [] for xy, point3D_id in zip(img.xys, img.point3D_ids): points_strings.append(" ".join(map(str, [*xy, point3D_id]))) fid.write(" ".join(points_strings) + "\n") def write_images_binary(images, path_to_model_file): """ see: src/base/reconstruction.cc void Reconstruction::ReadImagesBinary(const std::string& path) void Reconstruction::WriteImagesBinary(const std::string& path) """ with open(path_to_model_file, "wb") as fid: write_next_bytes(fid, len(images), "Q") for _, img in images.items(): write_next_bytes(fid, img.id, "i") write_next_bytes(fid, img.qvec.tolist(), "dddd") write_next_bytes(fid, img.tvec.tolist(), "ddd") write_next_bytes(fid, img.camera_id, "i") for char in img.name: write_next_bytes(fid, char.encode("utf-8"), "c") write_next_bytes(fid, b"\x00", "c") write_next_bytes(fid, len(img.point3D_ids), "Q") for xy, p3d_id in zip(img.xys, img.point3D_ids): write_next_bytes(fid, [*xy, p3d_id], "ddq") def read_points3D_text(path): """ see: src/base/reconstruction.cc void Reconstruction::ReadPoints3DText(const std::string& path) void Reconstruction::WritePoints3DText(const std::string& path) """ points3D = {} with open(path, "r") as fid: while True: line = fid.readline() if not line: break line = line.strip() if len(line) > 0 and line[0] != "#": elems = line.split() point3D_id = int(elems[0]) xyz = np.array(tuple(map(float, elems[1:4]))) rgb = np.array(tuple(map(int, elems[4:7]))) error = float(elems[7]) image_ids = np.array(tuple(map(int, elems[8::2]))) point2D_idxs = np.array(tuple(map(int, elems[9::2]))) points3D[point3D_id] = Point3D( id=point3D_id, xyz=xyz, rgb=rgb, error=error, image_ids=image_ids, point2D_idxs=point2D_idxs ) return points3D def read_points3D_binary(path_to_model_file): """ see: src/base/reconstruction.cc void Reconstruction::ReadPoints3DBinary(const std::string& path) void Reconstruction::WritePoints3DBinary(const std::string& path) """ points3D = {} with open(path_to_model_file, "rb") as fid: num_points = read_next_bytes(fid, 8, "Q")[0] for _ in range(num_points): binary_point_line_properties = read_next_bytes(fid, num_bytes=43, format_char_sequence="QdddBBBd") point3D_id = binary_point_line_properties[0] xyz = np.array(binary_point_line_properties[1:4]) rgb = np.array(binary_point_line_properties[4:7]) error = np.array(binary_point_line_properties[7]) track_length = read_next_bytes(fid, num_bytes=8, format_char_sequence="Q")[0] track_elems = read_next_bytes(fid, num_bytes=8 * track_length, format_char_sequence="ii" * track_length) image_ids = np.array(tuple(map(int, track_elems[0::2]))) point2D_idxs = np.array(tuple(map(int, track_elems[1::2]))) points3D[point3D_id] = Point3D( id=point3D_id, xyz=xyz, rgb=rgb, error=error, image_ids=image_ids, point2D_idxs=point2D_idxs ) return points3D def load_points3D(recon_dir: Path): print("Loading points3D") if (recon_dir / "points3D.bin").exists(): return read_points3D_binary(recon_dir / "points3D.bin") elif (recon_dir / "points3D.txt").exists(): return read_points3D_text(recon_dir / "points3D.txt") else: raise ValueError(f"Could not find points3D.txt or points3D.bin in {recon_dir}") def load_cameras(recon_dir: Path): print("Loading cameras") if (recon_dir / "cameras.bin").exists(): return read_cameras_binary(recon_dir / "cameras.bin") elif (recon_dir / "cameras.txt").exists(): return read_cameras_text(recon_dir / "cameras.txt") else: raise ValueError(f"Could not find cameras.txt or cameras.bin in {recon_dir}") def load_images(recon_dir: Path): print("Loading images") if (recon_dir / "images.bin").exists(): return read_images_binary(recon_dir / "images.bin") elif (recon_dir / "images.txt").exists(): return read_images_text(recon_dir / "images.txt") else: raise ValueError(f"Could not find images.txt or images.bin in {recon_dir}") # Copyright 2022 the Regents of the University of California, Nerfstudio Team and contributors. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. def parse_colmap_camera_params(camera) -> Dict[str, Any]: """ Parses all currently supported COLMAP cameras into the transforms.json metadata Args: camera: COLMAP camera Returns: transforms.json metadata containing camera's intrinsics and distortion parameters """ out: Dict[str, Any] = { "w": camera.width, "h": camera.height, } # Parameters match https://github.com/colmap/colmap/blob/dev/src/base/camera_models.h camera_params = camera.params if camera.model == "SIMPLE_PINHOLE": # du = 0 # dv = 0 out["fl_x"] = float(camera_params[0]) out["fl_y"] = float(camera_params[0]) out["cx"] = float(camera_params[1]) out["cy"] = float(camera_params[2]) out["k1"] = 0.0 out["k2"] = 0.0 out["p1"] = 0.0 out["p2"] = 0.0 camera_model = "OPENCV" elif camera.model == "PINHOLE": # f, cx, cy, k # du = 0 # dv = 0 out["fl_x"] = float(camera_params[0]) out["fl_y"] = float(camera_params[1]) out["cx"] = float(camera_params[2]) out["cy"] = float(camera_params[3]) out["k1"] = 0.0 out["k2"] = 0.0 out["p1"] = 0.0 out["p2"] = 0.0 camera_model = "OPENCV" elif camera.model == "SIMPLE_RADIAL": # f, cx, cy, k # r2 = u**2 + v**2; # radial = k * r2 # du = u * radial # dv = u * radial out["fl_x"] = float(camera_params[0]) out["fl_y"] = float(camera_params[0]) out["cx"] = float(camera_params[1]) out["cy"] = float(camera_params[2]) out["k1"] = float(camera_params[3]) out["k2"] = 0.0 out["p1"] = 0.0 out["p2"] = 0.0 camera_model = "OPENCV" elif camera.model == "RADIAL": # f, cx, cy, k1, k2 # r2 = u**2 + v**2; # radial = k1 * r2 + k2 * r2 ** 2 # du = u * radial # dv = v * radial out["fl_x"] = float(camera_params[0]) out["fl_y"] = float(camera_params[0]) out["cx"] = float(camera_params[1]) out["cy"] = float(camera_params[2]) out["k1"] = float(camera_params[3]) out["k2"] = float(camera_params[4]) out["p1"] = 0.0 out["p2"] = 0.0 camera_model = "OPENCV" elif camera.model == "OPENCV": # fx, fy, cx, cy, k1, k2, p1, p2 # uv = u * v; # r2 = u**2 + v**2 # radial = k1 * r2 + k2 * r2 ** 2 # du = u * radial + 2 * p1 * u*v + p2 * (r2 + 2 * u**2) # dv = v * radial + 2 * p2 * u*v + p1 * (r2 + 2 * v**2) out["fl_x"] = float(camera_params[0]) out["fl_y"] = float(camera_params[1]) out["cx"] = float(camera_params[2]) out["cy"] = float(camera_params[3]) out["k1"] = float(camera_params[4]) out["k2"] = float(camera_params[5]) out["p1"] = float(camera_params[6]) out["p2"] = float(camera_params[7]) camera_model = "OPENCV" elif camera.model == "OPENCV_FISHEYE": # fx, fy, cx, cy, k1, k2, k3, k4 # r = sqrt(u**2 + v**2) # if r > eps: # theta = atan(r) # theta2 = theta ** 2 # theta4 = theta2 ** 2 # theta6 = theta4 * theta2 # theta8 = theta4 ** 2 # thetad = theta * (1 + k1 * theta2 + k2 * theta4 + k3 * theta6 + k4 * theta8) # du = u * thetad / r - u; # dv = v * thetad / r - v; # else: # du = dv = 0 out["fl_x"] = float(camera_params[0]) out["fl_y"] = float(camera_params[1]) out["cx"] = float(camera_params[2]) out["cy"] = float(camera_params[3]) out["k1"] = float(camera_params[4]) out["k2"] = float(camera_params[5]) out["k3"] = float(camera_params[6]) out["k4"] = float(camera_params[7]) camera_model = "OPENCV_FISHEYE" elif camera.model == "FULL_OPENCV": # fx, fy, cx, cy, k1, k2, p1, p2, k3, k4, k5, k6 # u2 = u ** 2 # uv = u * v # v2 = v ** 2 # r2 = u2 + v2 # r4 = r2 * r2 # r6 = r4 * r2 # radial = (1 + k1 * r2 + k2 * r4 + k3 * r6) / # (1 + k4 * r2 + k5 * r4 + k6 * r6) # du = u * radial + 2 * p1 * uv + p2 * (r2 + 2 * u2) - u # dv = v * radial + 2 * p2 * uv + p1 * (r2 + 2 * v2) - v out["fl_x"] = float(camera_params[0]) out["fl_y"] = float(camera_params[1]) out["cx"] = float(camera_params[2]) out["cy"] = float(camera_params[3]) out["k1"] = float(camera_params[4]) out["k2"] = float(camera_params[5]) out["p1"] = float(camera_params[6]) out["p2"] = float(camera_params[7]) out["k3"] = float(camera_params[8]) k4 = float(camera_params[9]) k5 = float(camera_params[10]) k6 = float(camera_params[11]) if k4 != 0.0 or k5 != 0.0 or k6 != 0.0: raise NotImplementedError(f"{camera.model} camera model is not supported yet!") camera_model = "OPENCV" elif camera.model == "FOV": # fx, fy, cx, cy, omega out["fl_x"] = float(camera_params[0]) out["fl_y"] = float(camera_params[1]) out["cx"] = float(camera_params[2]) out["cy"] = float(camera_params[3]) out["omega"] = float(camera_params[4]) raise NotImplementedError(f"{camera.model} camera model is not supported yet!") elif camera.model == "SIMPLE_RADIAL_FISHEYE": # f, cx, cy, k # r = sqrt(u ** 2 + v ** 2) # if r > eps: # theta = atan(r) # theta2 = theta ** 2 # thetad = theta * (1 + k * theta2) # du = u * thetad / r - u; # dv = v * thetad / r - v; # else: # du = dv = 0 out["fl_x"] = float(camera_params[0]) out["fl_y"] = float(camera_params[0]) out["cx"] = float(camera_params[1]) out["cy"] = float(camera_params[2]) out["k1"] = float(camera_params[3]) out["k2"] = 0.0 out["k3"] = 0.0 out["k4"] = 0.0 camera_model = "OPENCV_FISHEYE" elif camera.model == "RADIAL_FISHEYE": # f, cx, cy, k1, k2 # r = sqrt(u ** 2 + v ** 2) # if r > eps: # theta = atan(r) # theta2 = theta ** 2 # theta4 = theta2 ** 2 # thetad = theta * (1 + k * theta2) # thetad = theta * (1 + k1 * theta2 + k2 * theta4) # du = u * thetad / r - u; # dv = v * thetad / r - v; # else: # du = dv = 0 out["fl_x"] = float(camera_params[0]) out["fl_y"] = float(camera_params[0]) out["cx"] = float(camera_params[1]) out["cy"] = float(camera_params[2]) out["k1"] = float(camera_params[3]) out["k2"] = float(camera_params[4]) out["k3"] = 0 out["k4"] = 0 camera_model = "OPENCV_FISHEYE" elif camera.model == "THIN_PRISM_FISHEYE": # fx, fy, cx, cy, k1, k2, p1, p2, k3, k4, sx1, sy1 out["fl_x"] = float(camera_params[0]) out["fl_y"] = float(camera_params[1]) out["cx"] = float(camera_params[2]) out["cy"] = float(camera_params[3]) out["k1"] = float(camera_params[4]) out["k2"] = float(camera_params[5]) out["p1"] = float(camera_params[6]) out["p2"] = float(camera_params[7]) out["k3"] = float(camera_params[8]) out["k4"] = float(camera_params[9]) out["sx1"] = float(camera_params[10]) out["sy1"] = float(camera_params[11]) camera_model = "OPENCV_FISHEYE" else: # RAD_TAN_THIN_PRISM_FISHEYE not supported! raise NotImplementedError(f"{camera.model} camera model is not supported yet!") out["camera_model"] = camera_model return out def colmap_to_json( recon_dir: Path, output_dir: Path, camera_mask_path: Optional[Path] = None, image_id_to_depth_path: Optional[Dict[int, Path]] = None, image_rename_map: Optional[Dict[str, str]] = None, ply_filename="sparse_pc.ply", keep_original_world_coordinate: bool = False, use_single_camera_mode: bool = True, ) -> Tuple[int, int]: """Converts COLMAP's cameras.bin and images.bin to a JSON file. Args: recon_dir: Path to the reconstruction directory, e.g. "sparse/0" output_dir: Path to the output directory. camera_model: Camera model used. camera_mask_path: Path to the camera mask. image_id_to_depth_path: When including sfm-based depth, embed these depth file paths in the exported json image_rename_map: Use these image names instead of the names embedded in the COLMAP db keep_original_world_coordinate: If True, no extra transform will be applied to world coordinate. Colmap optimized world often have y direction of the first camera pointing towards down direction, while nerfstudio world set z direction to be up direction for viewer. Returns: The number of total images. The number of registered images. """ # TODO(1480) use pycolmap # recon = pycolmap.Reconstruction(recon_dir) # cam_id_to_camera = recon.cameras # im_id_to_image = recon.images cam_id_to_camera = load_cameras(recon_dir) im_id_to_image = load_images(recon_dir) if len(set(cam_id_to_camera.keys())) > 1: print(f"Warning: More than one camera is found in {recon_dir}") use_single_camera_mode = False # update bool: one camera per frame out = {} # out = {"camera_model": parse_colmap_camera_params(cam_id_to_camera[1])["camera_model"]} else: # one camera for all frames camera_id = [*cam_id_to_camera.keys()][0] out = parse_colmap_camera_params(cam_id_to_camera[camera_id]) frames = [] for im_id, im_data in im_id_to_image.items(): # NB: COLMAP uses Eigen / scalar-first quaternions # * https://colmap.github.io/format.html # * https://github.com/colmap/colmap/blob/bf3e19140f491c3042bfd85b7192ef7d249808ec/src/base/pose.cc#L75 # the `rotation_matrix()` handles that format for us. # TODO(1480) BEGIN use pycolmap API # rotation = im_data.rotation_matrix() rotation = qvec2rotmat(im_data.qvec) translation = im_data.tvec.reshape(3, 1) w2c = np.concatenate([rotation, translation], 1) w2c = np.concatenate([w2c, np.array([[0, 0, 0, 1]])], 0) c2w = np.linalg.inv(w2c) # Convert from COLMAP's camera coordinate system (OpenCV) to ours (OpenGL) c2w[0:3, 1:3] *= -1 if not keep_original_world_coordinate: c2w = c2w[np.array([0, 2, 1, 3]), :] c2w[2, :] *= -1 name = im_data.name if image_rename_map is not None: name = image_rename_map[name] frame = { "file_path": (Path("images") / name).as_posix(), "transform_matrix": c2w.tolist(), "colmap_im_id": im_id, } if camera_mask_path is not None: mask_path = camera_mask_path / (name+".png") if mask_path.exists(): frame["mask_path"] = mask_path.as_posix() else: print("Warning:", mask_path, "does not exist") if image_id_to_depth_path is not None: depth_path = image_id_to_depth_path[im_id] frame["depth_file_path"] = str(depth_path.relative_to(depth_path.parent.parent)) if not use_single_camera_mode: # add the camera parameters for this frame frame.update(parse_colmap_camera_params(cam_id_to_camera[im_data.camera_id])) frames.append(frame) out["frames"] = frames applied_transform = None if not keep_original_world_coordinate: applied_transform = np.eye(4)[:3, :] applied_transform = applied_transform[np.array([0, 2, 1]), :] applied_transform[2, :] *= -1 out["applied_transform"] = applied_transform.tolist() # create ply from colmap assert ply_filename.endswith(".ply"), f"ply_filename: {ply_filename} does not end with '.ply'" create_ply_from_colmap( ply_filename, recon_dir, output_dir, torch.from_numpy(applied_transform).float() if applied_transform is not None else None, ) out["ply_file_path"] = ply_filename with open(output_dir / "transforms.json", "w", encoding="utf-8") as f: json.dump(out, f, indent=4) return len(im_id_to_image), len(frames) def create_ply_from_colmap( filename: str, recon_dir: Path, output_dir: Path, applied_transform: Union[torch.Tensor, None] ) -> None: """Writes a ply file from colmap. Args: filename: file name for .ply recon_dir: Directory to grab colmap points output_dir: Directory to output .ply """ colmap_points = load_points3D(recon_dir) # Load point Positions points3D = torch.from_numpy(np.array([p.xyz for p in colmap_points.values()], dtype=np.float32)) if applied_transform is not None: assert applied_transform.shape == (3, 4) points3D = torch.einsum("ij,bj->bi", applied_transform[:3, :3], points3D) + applied_transform[:3, 3] # Load point colours points3D_rgb = torch.from_numpy(np.array([p.rgb for p in colmap_points.values()], dtype=np.uint8)) # write ply with open(output_dir / filename, "w") as f: # Header f.write("ply\n") f.write("format ascii 1.0\n") f.write(f"element vertex {len(points3D)}\n") f.write("property float x\n") f.write("property float y\n") f.write("property float z\n") f.write("property uint8 red\n") f.write("property uint8 green\n") f.write("property uint8 blue\n") f.write("end_header\n") for coord, color in zip(points3D, points3D_rgb): x, y, z = coord r, g, b = color f.write(f"{x:8f} {y:8f} {z:8f} {r} {g} {b}\n") def get_matching_summary(num_initial_frames: int, num_matched_frames: int) -> str: """Returns a summary of the matching results. Args: num_initial_frames: The number of initial frames. num_matched_frames: The number of matched frames. Returns: A summary of the matching results. """ match_ratio = num_matched_frames / num_initial_frames if match_ratio == 1: return "[bold green]COLMAP found poses for all images, CONGRATS!" if match_ratio < 0.4: result = f"[bold red]COLMAP only found poses for {num_matched_frames / num_initial_frames * 100:.2f}%" result += " of the images. This is low.\nThis can be caused by a variety of reasons," result += " such poor scene coverage, blurry images, or large exposure changes." return result if match_ratio < 0.8: result = f"[bold yellow]COLMAP only found poses for {num_matched_frames / num_initial_frames * 100:.2f}%" result += " of the images.\nThis isn't great, but may be ok." result += "\nMissing poses can be caused by a variety of reasons, such poor scene coverage, blurry images," result += " or large exposure changes." return result return f"[bold green]COLMAP found poses for {num_matched_frames / num_initial_frames * 100:.2f}% of the images." @dataclass class ColmapConverterToNerfstudioDataset(): """Base class to process images or video into a nerfstudio dataset using colmap""" work_dir: Path """Path to the dataset directory.""" colmap_model_path: Path = Path("sparse/0") """Path of the colmap model, relative to the dataset directory. """ camera_mask_path: Path = Path("masks") """Path of the colmap model, relative to the dataset directory. """ @property def absolute_colmap_model_path(self) -> Path: return self.work_dir / self.colmap_model_path @property def absolute_camera_mask_path(self) -> Optional[Path]: path = self.work_dir / self.camera_mask_path return path if path.exists() else None @property def absolute_colmap_path(self) -> Path: return self.work_dir / "colmap" def main(self) -> List[str]: """Save colmap transforms into the output folder Args: image_id_to_depth_path: When including sfm-based depth, embed these depth file paths in the exported json image_rename_map: Use these image names instead of the names embedded in the COLMAP db """ summary_log = [] if (self.absolute_colmap_model_path / "cameras.bin").exists() \ or (self.absolute_colmap_model_path / "cameras.txt").exists(): print("Saving results to transforms.json") num_frames, num_matched_frames = colmap_to_json( recon_dir=self.absolute_colmap_model_path, output_dir=self.work_dir, image_id_to_depth_path=None, camera_mask_path=self.absolute_camera_mask_path, image_rename_map=None, use_single_camera_mode=True, ) summary_log.append(f"Colmap matched {num_matched_frames} images") summary_log.append(get_matching_summary(num_frames, num_matched_frames)) else: print( "Warning: Could not find existing COLMAP results. " "Not generating transforms.json" + \ "If needed, specify `--colmap_model_path` (default is `sparse/0`)." ) for log in summary_log: print(log) def entrypoint(): """Entrypoint for use with pyproject scripts.""" import tyro tyro.extras.set_accent_color("bright_yellow") try: tyro.cli(ColmapConverterToNerfstudioDataset).main() except (RuntimeError, ValueError) as e: print(e.args[0]) if __name__ == "__main__": entrypoint()