trending/spirula-studio
1
0
Fork 0
mirror of https://github.com/harry7557558/spirula-studio.git synced 2026-10-02 02:44:54 +08:00
Code Issues Packages Projects Releases Wiki Activity
Files
master
spirula-studio/reference/scripts/process_data_colmap.py
T
Harry Chen 3157c78647 move scripts to reference
2026-08-07 17:46:01 -04:00

910 lines
36 KiB
Python
Raw Permalink Blame History

#!/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()
Reference in New Issue View Git Blame Copy Permalink