From 7309273d9241b681e404d7f6a7f13bddc477e473 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jon=C3=A1=C5=A1=20Kulh=C3=A1nek?= Date: Sun, 26 Feb 2023 18:53:08 +0100 Subject: [PATCH] Add support for more colmap camera models in ns-process-data (#1500) * Add support for more colmap camera models in ns-process-data * Fix bugs introduced in #1457 --- nerfstudio/process_data/colmap_utils.py | 242 ++++++++++++++++++++---- scripts/process_data.py | 7 +- 2 files changed, 209 insertions(+), 40 deletions(-) diff --git a/nerfstudio/process_data/colmap_utils.py b/nerfstudio/process_data/colmap_utils.py index 064c2110a..f8254f0a8 100644 --- a/nerfstudio/process_data/colmap_utils.py +++ b/nerfstudio/process_data/colmap_utils.py @@ -18,7 +18,7 @@ Tools supporting the execution of COLMAP and preparation of COLMAP-based dataset import json from pathlib import Path -from typing import Dict, Optional +from typing import Any, Dict, Optional import appdirs import cv2 @@ -182,10 +182,212 @@ def run_colmap( CONSOLE.log("[bold green]:tada: Done refining intrinsics.") +def parse_colmap_camera_params(camera) -> Dict[str, Any]: # pylint: disable=too-many-statements + """ + 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 = CameraModel.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 = CameraModel.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 = CameraModel.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 = CameraModel.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 = CameraModel.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 = CameraModel.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]) + out["k4"] = float(camera_params[9]) + out["k5"] = float(camera_params[10]) + out["k6"] = float(camera_params[11]) + raise NotImplementedError(f"{camera.model} camera model is not supported yet!") + 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 = CameraModel.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 = CameraModel.OPENCV_FISHEYE + else: + # THIN_PRISM_FISHEYE not supported! + raise NotImplementedError(f"{camera.model} camera model is not supported yet!") + + out["camera_model"] = camera_model.value + return out + + def colmap_to_json( recon_dir: Path, output_dir: Path, - camera_model: CameraModel, camera_mask_path: Optional[Path] = None, image_id_to_depth_path: Optional[Dict[int, Path]] = None, image_rename_map: Optional[Dict[str, str]] = None, @@ -211,10 +413,6 @@ def colmap_to_json( cam_id_to_camera = read_cameras_binary(recon_dir / "cameras.bin") im_id_to_image = read_images_binary(recon_dir / "images.bin") - # Only support first camera - CAMERA_ID = 1 # pylint: disable=invalid-name - camera_params = cam_id_to_camera[CAMERA_ID].params - frames = [] for im_id, im_data in im_id_to_image.items(): # NB: COLMAP uses Eigen / scalar-first quaternions @@ -252,35 +450,9 @@ def colmap_to_json( frame["depth_file_path"] = str(depth_path) frames.append(frame) - out = { - "fl_x": float(camera_params[0]), - "fl_y": float(camera_params[1]), - "cx": float(camera_params[2]), - "cy": float(camera_params[3]), - "w": cam_id_to_camera[CAMERA_ID].width, - "h": cam_id_to_camera[CAMERA_ID].height, - "camera_model": camera_model.value, - } - - if camera_model == CameraModel.OPENCV: - out.update( - { - "k1": float(camera_params[4]), - "k2": float(camera_params[5]), - "p1": float(camera_params[6]), - "p2": float(camera_params[7]), - } - ) - if camera_model == CameraModel.OPENCV_FISHEYE: - out.update( - { - "k1": float(camera_params[4]), - "k2": float(camera_params[5]), - "k3": float(camera_params[6]), - "k4": float(camera_params[7]), - } - ) - + if set(cam_id_to_camera.keys()) != {1}: + raise RuntimeError("Only single camera shared for all images is supported.") + out = parse_colmap_camera_params(cam_id_to_camera[1]) out["frames"] = frames with open(output_dir / "transforms.json", "w", encoding="utf-8") as f: diff --git a/scripts/process_data.py b/scripts/process_data.py index 53ac9c0da..dfb875bd0 100644 --- a/scripts/process_data.py +++ b/scripts/process_data.py @@ -182,15 +182,14 @@ class ProcessImages: ) ) else: - image_id_to_depth_path = {} + image_id_to_depth_path = None # Save transforms.json if (colmap_model_path / "cameras.bin").exists(): with CONSOLE.status("[bold yellow]Saving results to transforms.json", spinner="balloon"): num_matched_frames = colmap_utils.colmap_to_json( - recon_dir=colmap_dir / "sparse" / "0", + recon_dir=colmap_model_path, output_dir=self.output_dir, - camera_model=CAMERA_MODELS[self.camera_type], image_id_to_depth_path=image_id_to_depth_path, image_rename_map=image_rename_map, ) @@ -424,7 +423,6 @@ class ProcessVideo: num_matched_frames = colmap_utils.colmap_to_json( recon_dir=colmap_dir / "sparse" / "0", output_dir=self.output_dir, - camera_model=CAMERA_MODELS[self.camera_type], image_id_to_depth_path=image_id_to_depth_path, camera_mask_path=mask_path, image_rename_map=None, @@ -568,7 +566,6 @@ class ProcessInsta360: num_matched_frames = colmap_utils.colmap_to_json( recon_dir=colmap_dir / "sparse" / "0", output_dir=self.output_dir, - camera_model=CAMERA_MODELS["fisheye"], image_id_to_depth_path=image_id_to_depth_path, image_rename_map=None, )