Merge 539b55ffce into 50e0e3c70c

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])

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