From ed9d3464020b704be591d2e8aa5d6133e0f66db0 Mon Sep 17 00:00:00 2001 From: Riccardo de Lutio Date: Mon, 18 May 2026 15:43:16 -0700 Subject: [PATCH] feat(dataset): support COLMAP OpenCV pinhole distortion --- threedgrut/datasets/dataset_colmap.py | 77 ++++++++++++----- .../datasets/tests/test_dataset_colmap.py | 82 +++++++++++++++++++ 2 files changed, 139 insertions(+), 20 deletions(-) create mode 100644 threedgrut/datasets/tests/test_dataset_colmap.py diff --git a/threedgrut/datasets/dataset_colmap.py b/threedgrut/datasets/dataset_colmap.py index 82e56b1b..67d0ce57 100644 --- a/threedgrut/datasets/dataset_colmap.py +++ b/threedgrut/datasets/dataset_colmap.py @@ -47,6 +47,37 @@ ) +def _opencv_pinhole_intrinsics_from_colmap(intr_model: str, intr_params: np.ndarray, scaling_factor: float): + """Convert COLMAP distorted pinhole parameters to OpenCVPinholeCameraModelParameters fields.""" + intr_params = np.asarray(intr_params, dtype=np.float32) + radial_coeffs = np.zeros((6,), dtype=np.float32) + tangential_coeffs = np.zeros((2,), dtype=np.float32) + thin_prism_coeffs = np.zeros((4,), dtype=np.float32) + + if intr_model == "SIMPLE_RADIAL": + focal_length = np.array([intr_params[0], intr_params[0]], dtype=np.float32) / scaling_factor + principal_point = intr_params[1:3] / scaling_factor + radial_coeffs[0] = intr_params[3] + elif intr_model == "RADIAL": + focal_length = np.array([intr_params[0], intr_params[0]], dtype=np.float32) / scaling_factor + principal_point = intr_params[1:3] / scaling_factor + radial_coeffs[:2] = intr_params[3:5] + elif intr_model == "OPENCV": + focal_length = intr_params[:2] / scaling_factor + principal_point = intr_params[2:4] / scaling_factor + radial_coeffs[:2] = intr_params[4:6] + tangential_coeffs[:] = intr_params[6:8] + elif intr_model == "FULL_OPENCV": + focal_length = intr_params[:2] / scaling_factor + principal_point = intr_params[2:4] / scaling_factor + radial_coeffs[:] = intr_params[[4, 5, 8, 9, 10, 11]] + tangential_coeffs[:] = intr_params[6:8] + else: + raise ValueError(f"Unsupported distorted pinhole camera model: {intr_model}") + + return focal_length, principal_point, radial_coeffs, tangential_coeffs, thin_prism_coeffs + + class ColmapDataset(Dataset, BoundedMultiViewDataset, DatasetVisualization): def __init__( self, @@ -174,9 +205,15 @@ def create_pinhole_camera(focalx, focaly, w, h, cx=None, cy=None): pixel_coords, ) - def create_opencv_pinhole_camera(focalx, focaly, w, h, cx=None, cy=None, radial_coeffs=None): - cx = cx if cx is not None else w / 2.0 - cy = cy if cy is not None else h / 2.0 + def create_opencv_pinhole_camera( + focal_length, + principal_point, + w, + h, + radial_coeffs, + tangential_coeffs, + thin_prism_coeffs, + ): # Generate UV coordinates u = np.tile(np.arange(w), h) v = np.arange(h).repeat(w) @@ -184,15 +221,11 @@ def create_opencv_pinhole_camera(focalx, focaly, w, h, cx=None, cy=None, radial_ params = OpenCVPinholeCameraModelParameters( resolution=np.array([w, h], dtype=np.uint64), shutter_type=ShutterType.GLOBAL, - principal_point=np.array([cx, cy], dtype=np.float32), - focal_length=np.array([focalx, focaly], dtype=np.float32), - radial_coeffs=( - np.zeros((6,), dtype=np.float32) - if radial_coeffs is None - else np.asarray(radial_coeffs, dtype=np.float32) - ), - tangential_coeffs=np.zeros((2,), dtype=np.float32), - thin_prism_coeffs=np.zeros((4,), dtype=np.float32), + principal_point=np.asarray(principal_point, dtype=np.float32), + focal_length=np.asarray(focal_length, dtype=np.float32), + radial_coeffs=np.asarray(radial_coeffs, dtype=np.float32), + tangential_coeffs=np.asarray(tangential_coeffs, dtype=np.float32), + thin_prism_coeffs=np.asarray(thin_prism_coeffs, dtype=np.float32), ) camera_model = ncore.sensors.CameraModel.from_parameters(params, device="cpu", dtype=torch.float32) int_pixel_coords = torch.tensor(np.stack([u, v], axis=1), dtype=torch.int32) @@ -292,14 +325,18 @@ def create_fisheye_camera(params, w, h): focal_length_x, focal_length_y, width, height, cx=cx, cy=cy ) - elif intr.model == "SIMPLE_RADIAL": - focal_length = intr.params[0] / scaling_factor - cx = intr.params[1] / scaling_factor - cy = intr.params[2] / scaling_factor - radial_coeffs = np.zeros((6,), dtype=np.float32) - radial_coeffs[0] = intr.params[3] + elif intr.model in {"SIMPLE_RADIAL", "RADIAL", "OPENCV", "FULL_OPENCV"}: + focal_length, principal_point, radial_coeffs, tangential_coeffs, thin_prism_coeffs = ( + _opencv_pinhole_intrinsics_from_colmap(intr.model, intr.params, scaling_factor) + ) self.intrinsics[intr.id] = create_opencv_pinhole_camera( - focal_length, focal_length, width, height, cx=cx, cy=cy, radial_coeffs=radial_coeffs + focal_length, + principal_point, + width, + height, + radial_coeffs, + tangential_coeffs, + thin_prism_coeffs, ) elif intr.model == "OPENCV_FISHEYE": @@ -310,7 +347,7 @@ def create_fisheye_camera(params, w, h): else: assert False, ( f"Colmap camera model '{intr.model}' not handled: supported camera models are " - "PINHOLE, SIMPLE_PINHOLE, SIMPLE_RADIAL, and OPENCV_FISHEYE." + "PINHOLE, SIMPLE_PINHOLE, SIMPLE_RADIAL, RADIAL, OPENCV, FULL_OPENCV, and OPENCV_FISHEYE." ) # Load poses and paths diff --git a/threedgrut/datasets/tests/test_dataset_colmap.py b/threedgrut/datasets/tests/test_dataset_colmap.py new file mode 100644 index 00000000..b47411cd --- /dev/null +++ b/threedgrut/datasets/tests/test_dataset_colmap.py @@ -0,0 +1,82 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +import numpy as np +import pytest + +pytest.importorskip("ncore") + +from threedgrut.datasets.dataset_colmap import _opencv_pinhole_intrinsics_from_colmap + + +@pytest.mark.parametrize( + "model,params,expected_focal,expected_principal,expected_radial,expected_tangential", + [ + ( + "SIMPLE_RADIAL", + [100.0, 40.0, 45.0, 0.1], + [50.0, 50.0], + [20.0, 22.5], + [0.1, 0.0, 0.0, 0.0, 0.0, 0.0], + [0.0, 0.0], + ), + ( + "RADIAL", + [100.0, 40.0, 45.0, 0.1, 0.2], + [50.0, 50.0], + [20.0, 22.5], + [0.1, 0.2, 0.0, 0.0, 0.0, 0.0], + [0.0, 0.0], + ), + ( + "OPENCV", + [100.0, 110.0, 40.0, 45.0, 0.1, 0.2, 0.01, 0.02], + [50.0, 55.0], + [20.0, 22.5], + [0.1, 0.2, 0.0, 0.0, 0.0, 0.0], + [0.01, 0.02], + ), + ( + "FULL_OPENCV", + [100.0, 110.0, 40.0, 45.0, 0.1, 0.2, 0.01, 0.02, 0.3, 0.4, 0.5, 0.6], + [50.0, 55.0], + [20.0, 22.5], + [0.1, 0.2, 0.3, 0.4, 0.5, 0.6], + [0.01, 0.02], + ), + ], +) +def test_opencv_pinhole_intrinsics_from_colmap( + model, + params, + expected_focal, + expected_principal, + expected_radial, + expected_tangential, +): + focal, principal, radial, tangential, thin_prism = _opencv_pinhole_intrinsics_from_colmap( + model, np.asarray(params), scaling_factor=2 + ) + + np.testing.assert_allclose(focal, expected_focal) + np.testing.assert_allclose(principal, expected_principal) + np.testing.assert_allclose(radial, expected_radial) + np.testing.assert_allclose(tangential, expected_tangential) + np.testing.assert_allclose(thin_prism, np.zeros((4,), dtype=np.float32)) + + +def test_opencv_pinhole_intrinsics_rejects_unsupported_model(): + with pytest.raises(ValueError, match="Unsupported distorted pinhole camera model"): + _opencv_pinhole_intrinsics_from_colmap("FOV", np.zeros((5,)), scaling_factor=1)