diff --git a/nerfstudio/process_data/colmap_utils.py b/nerfstudio/process_data/colmap_utils.py index 1d9405c81..220caa088 100644 --- a/nerfstudio/process_data/colmap_utils.py +++ b/nerfstudio/process_data/colmap_utils.py @@ -547,11 +547,6 @@ def create_sfm_depth( 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 - W = cam_id_to_camera[CAMERA_ID].width - H = cam_id_to_camera[CAMERA_ID].height - if verbose: iter_images = track( im_id_to_image.items(), total=len(im_id_to_image.items()), description="Creating depth maps ..." @@ -561,6 +556,8 @@ def create_sfm_depth( image_id_to_depth_path = {} for im_id, im_data in iter_images: + camera = cam_id_to_camera[im_data.camera_id] + W, H = camera.width, camera.height # TODO(1480) BEGIN delete when abandoning colmap_parsing_utils pids = [pid for pid in im_data.point3D_ids if pid != -1] xyz_world = np.array([ptid_to_info[pid].xyz for pid in pids]) diff --git a/tests/process_data/test_sfm_depth_camera_size.py b/tests/process_data/test_sfm_depth_camera_size.py new file mode 100644 index 000000000..1c9c059c4 --- /dev/null +++ b/tests/process_data/test_sfm_depth_camera_size.py @@ -0,0 +1,34 @@ +from pathlib import Path + +import cv2 +import numpy as np + +from nerfstudio.data.utils.colmap_parsing_utils import ( + Camera, + Image, + Point3D, + write_cameras_binary, + write_images_binary, + write_points3D_binary, +) +from nerfstudio.process_data.colmap_utils import create_sfm_depth + + +def test_sfm_depth_uses_each_images_camera(tmp_path: Path): + cameras = {key: Camera(key, "PINHOLE", w, h, [1, 1, 0, 0]) for key, w, h in [(3, 4, 2), (8, 9, 5)]} + images = { + i: Image(i, np.array([1, 0, 0, 0]), np.zeros(3), camera_id, f"{i}.png", np.array([xy]), np.array([1])) + for i, camera_id, xy in [(1, 3, [2, 1]), (2, 8, [7, 4])] + } + point = Point3D(1, np.array([0, 0, 2]), np.zeros(3, dtype=int), 0.0, np.array([1, 2]), np.array([0, 0])) + write_cameras_binary(cameras, tmp_path / "cameras.bin") + write_images_binary(images, tmp_path / "images.bin") + write_points3D_binary({1: point}, tmp_path / "points3D.bin") + + paths = create_sfm_depth(tmp_path, tmp_path, verbose=False) + for i, image in images.items(): + depth = cv2.imread(str(paths[i]), cv2.IMREAD_UNCHANGED) + camera = cameras[image.camera_id] + assert depth.shape == (camera.height, camera.width) + x, y = image.xys[0] + assert depth[y, x] == 2000