merged main

.github/workflows/core_code_checks.yml

@@ -33,9 +33,9 @@ jobs:
         run: |
           python ./nerfstudio/scripts/docs/add_nb_tags.py --check
       - name: Run Ruff Linter
-        run: ruff check docs/ nerfstudio/ tests/
+        run: ruff check docs/ nerfstudio/ tests/ --output-format=github
       - name: Run Ruff Formatter
-        run: ruff format docs/ nerfstudio/ tests/ --check
+        run: ruff format docs/ nerfstudio/ tests/ --diff
       - name: Run Pyright
         run: |
           pyright

nerfstudio/cameras/camera_utils.py

@@ -720,7 +720,7 @@ def fisheye624_unproject_helper(uv, params, max_iters: int = 5):
     function so this solves an optimization problem using Newton's method to get
     the inverse.
     Inputs:
-        uv: BxNx3 tensor of 2D pixels to be projected
+        uv: BxNx2 tensor of 2D pixels to be unprojected
         params: Bx16 tensor of Fisheye624 parameters formatted like this:
                 [f_u f_v c_u c_v {k_0 ... k_5} {p_0 p_1} {s_0 s_1 s_2 s_3}]
                 or Bx15 tensor of Fisheye624 parameters formatted like this:

nerfstudio/cameras/cameras.py

@@ -864,7 +864,7 @@ class Cameras(TensorDataclass):
 
                 assert distortion_params is not None
                 masked_coords = pcoord_stack[coord_mask, :]
-                # The fisheye unprojection does not rely on planar/pinhold unprojection, thus the method needs
+                # The fisheye unprojection does not rely on planar/pinhole unprojection, thus the method needs
                 # to access the focal length and principle points directly.
                 camera_params = torch.cat(
                     [

nerfstudio/configs/method_configs.py

@@ -31,7 +31,6 @@ from nerfstudio.data.datamanagers.full_images_datamanager import FullImageDatama
 from nerfstudio.data.datamanagers.parallel_datamanager import ParallelDataManagerConfig
 from nerfstudio.data.datamanagers.random_cameras_datamanager import RandomCamerasDataManagerConfig
 from nerfstudio.data.dataparsers.blender_dataparser import BlenderDataParserConfig
-from nerfstudio.data.dataparsers.colmap_dataparser import ColmapDataParserConfig
 from nerfstudio.data.dataparsers.dnerf_dataparser import DNeRFDataParserConfig
 from nerfstudio.data.dataparsers.instant_ngp_dataparser import InstantNGPDataParserConfig
 from nerfstudio.data.dataparsers.nerfstudio_dataparser import NerfstudioDataParserConfig
@@ -600,7 +599,7 @@ method_configs["gaussian-splatting"] = TrainerConfig(
     gradient_accumulation_steps={"camera_opt": 100},
     pipeline=VanillaPipelineConfig(
         datamanager=FullImageDatamanagerConfig(
-            dataparser=ColmapDataParserConfig(load_3D_points=True),
+            dataparser=NerfstudioDataParserConfig(load_3D_points=True),
         ),
         model=GaussianSplattingModelConfig(),
     ),

nerfstudio/data/datamanagers/full_images_datamanager.py

@@ -34,6 +34,7 @@ import torch
 from torch.nn import Parameter
 from tqdm import tqdm
 
+from nerfstudio.cameras.camera_utils import fisheye624_project, fisheye624_unproject_helper
 from nerfstudio.cameras.cameras import Cameras, CameraType
 from nerfstudio.configs.dataparser_configs import AnnotatedDataParserUnion
 from nerfstudio.data.datamanagers.base_datamanager import DataManager, DataManagerConfig, TDataset
@@ -135,63 +136,11 @@ class FullImageDatamanager(DataManager, Generic[TDataset]):
                 continue
             distortion_params = camera.distortion_params.numpy()
             image = data["image"].numpy()
-            if camera.camera_type.item() == CameraType.PERSPECTIVE.value:
-                distortion_params = np.array(
-                    [
-                        distortion_params[0],
-                        distortion_params[1],
-                        distortion_params[4],
-                        distortion_params[5],
-                        distortion_params[2],
-                        distortion_params[3],
-                        0,
-                        0,
-                    ]
-                )
-                if np.any(distortion_params):
-                    newK, roi = cv2.getOptimalNewCameraMatrix(K, distortion_params, (image.shape[1], image.shape[0]), 0)
-                    image = cv2.undistort(image, K, distortion_params, None, newK)  # type: ignore
-                else:
-                    newK = K
-                    roi = 0, 0, image.shape[1], image.shape[0]
-                # crop the image and update the intrinsics accordingly
-                x, y, w, h = roi
-                image = image[y : y + h, x : x + w]
-                if "depth_image" in data:
-                    data["depth_image"] = data["depth_image"][y : y + h, x : x + w]
-                # update the width, height
-                self.train_dataset.cameras.width[i] = w
-                self.train_dataset.cameras.height[i] = h
-                if "mask" in data:
-                    mask = data["mask"].numpy()
-                    mask = mask.astype(np.uint8) * 255
-                    if np.any(distortion_params):
-                        mask = cv2.undistort(mask, K, distortion_params, None, newK)  # type: ignore
-                    mask = mask[y : y + h, x : x + w]
-                    data["mask"] = torch.from_numpy(mask).bool()
-                K = newK
 
-            elif camera.camera_type.item() == CameraType.FISHEYE.value:
-                distortion_params = np.array(
-                    [distortion_params[0], distortion_params[1], distortion_params[2], distortion_params[3]]
-                )
-                newK = cv2.fisheye.estimateNewCameraMatrixForUndistortRectify(
-                    K, distortion_params, (image.shape[1], image.shape[0]), np.eye(3), balance=0
-                )
-                map1, map2 = cv2.fisheye.initUndistortRectifyMap(
-                    K, distortion_params, np.eye(3), newK, (image.shape[1], image.shape[0]), cv2.CV_32FC1
-                )
-                # and then remap:
-                image = cv2.remap(image, map1, map2, interpolation=cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT)
-                if "mask" in data:
-                    mask = data["mask"].numpy()
-                    mask = mask.astype(np.uint8) * 255
-                    mask = cv2.fisheye.undistortImage(mask, K, distortion_params, None, newK)
-                    data["mask"] = torch.from_numpy(mask).bool()
-                K = newK
-            else:
-                raise NotImplementedError("Only perspective and fisheye cameras are supported")
+            K, image, mask = _undistort_image(camera, distortion_params, data, image, K)
             data["image"] = torch.from_numpy(image)
+            if mask is not None:
+                data["mask"] = mask
 
             cached_train.append(data)
 
@@ -199,6 +148,8 @@ class FullImageDatamanager(DataManager, Generic[TDataset]):
             self.train_dataset.cameras.fy[i] = float(K[1, 1])
             self.train_dataset.cameras.cx[i] = float(K[0, 2])
             self.train_dataset.cameras.cy[i] = float(K[1, 2])
+            self.train_dataset.cameras.width[i] = image.shape[1]
+            self.train_dataset.cameras.height[i] = image.shape[0]
 
         CONSOLE.log("Caching / undistorting eval images")
         for i in tqdm(range(len(self.eval_dataset)), leave=False):
@@ -210,61 +161,11 @@ class FullImageDatamanager(DataManager, Generic[TDataset]):
                 continue
             distortion_params = camera.distortion_params.numpy()
             image = data["image"].numpy()
-            if camera.camera_type.item() == CameraType.PERSPECTIVE.value:
-                distortion_params = np.array(
-                    [
-                        distortion_params[0],
-                        distortion_params[1],
-                        distortion_params[4],
-                        distortion_params[5],
-                        distortion_params[2],
-                        distortion_params[3],
-                        0,
-                        0,
-                    ]
-                )
-                if np.any(distortion_params):
-                    newK, roi = cv2.getOptimalNewCameraMatrix(K, distortion_params, (image.shape[1], image.shape[0]), 0)
-                    image = cv2.undistort(image, K, distortion_params, None, newK)  # type: ignore
-                else:
-                    newK = K
-                    roi = 0, 0, image.shape[1], image.shape[0]
-                # crop the image and update the intrinsics accordingly
-                x, y, w, h = roi
-                image = image[y : y + h, x : x + w]
-                # update the width, height
-                self.eval_dataset.cameras.width[i] = w
-                self.eval_dataset.cameras.height[i] = h
-                if "mask" in data:
-                    mask = data["mask"].numpy()
-                    mask = mask.astype(np.uint8) * 255
-                    if np.any(distortion_params):
-                        mask = cv2.undistort(mask, K, distortion_params, None, newK)  # type: ignore
-                    mask = mask[y : y + h, x : x + w]
-                    data["mask"] = torch.from_numpy(mask).bool()
-                K = newK
 
-            elif camera.camera_type.item() == CameraType.FISHEYE.value:
-                distortion_params = np.array(
-                    [distortion_params[0], distortion_params[1], distortion_params[2], distortion_params[3]]
-                )
-                newK = cv2.fisheye.estimateNewCameraMatrixForUndistortRectify(
-                    K, distortion_params, (image.shape[1], image.shape[0]), np.eye(3), balance=0
-                )
-                map1, map2 = cv2.fisheye.initUndistortRectifyMap(
-                    K, distortion_params, np.eye(3), newK, (image.shape[1], image.shape[0]), cv2.CV_32FC1
-                )
-                # and then remap:
-                image = cv2.remap(image, map1, map2, interpolation=cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT)
-                if "mask" in data:
-                    mask = data["mask"].numpy()
-                    mask = mask.astype(np.uint8) * 255
-                    mask = cv2.fisheye.undistortImage(mask, K, distortion_params, None, newK)
-                    data["mask"] = torch.from_numpy(mask).bool()
-                K = newK
-            else:
-                raise NotImplementedError("Only perspective and fisheye cameras are supported")
+            K, image, mask = _undistort_image(camera, distortion_params, data, image, K)
             data["image"] = torch.from_numpy(image)
+            if mask is not None:
+                data["mask"] = mask
 
             cached_eval.append(data)
 
@@ -272,6 +173,8 @@ class FullImageDatamanager(DataManager, Generic[TDataset]):
             self.eval_dataset.cameras.fy[i] = float(K[1, 1])
             self.eval_dataset.cameras.cx[i] = float(K[0, 2])
             self.eval_dataset.cameras.cy[i] = float(K[1, 2])
+            self.eval_dataset.cameras.width[i] = image.shape[1]
+            self.eval_dataset.cameras.height[i] = image.shape[0]
 
         if cache_images_option == "gpu":
             for cache in cached_train:
@@ -416,3 +319,156 @@ class FullImageDatamanager(DataManager, Generic[TDataset]):
         assert len(self.eval_dataset.cameras.shape) == 1, "Assumes single batch dimension"
         camera = self.eval_dataset.cameras[image_idx : image_idx + 1].to(self.device)
         return camera, data
+
+
+def _undistort_image(
+    camera: Cameras, distortion_params: np.ndarray, data: dict, image: np.ndarray, K: np.ndarray
+) -> Tuple[np.ndarray, np.ndarray, Optional[torch.Tensor]]:
+    mask = None
+    if camera.camera_type.item() == CameraType.PERSPECTIVE.value:
+        distortion_params = np.array(
+            [
+                distortion_params[0],
+                distortion_params[1],
+                distortion_params[4],
+                distortion_params[5],
+                distortion_params[2],
+                distortion_params[3],
+                0,
+                0,
+            ]
+        )
+        if np.any(distortion_params):
+            newK, roi = cv2.getOptimalNewCameraMatrix(K, distortion_params, (image.shape[1], image.shape[0]), 0)
+            image = cv2.undistort(image, K, distortion_params, None, newK)  # type: ignore
+        else:
+            newK = K
+            roi = 0, 0, image.shape[1], image.shape[0]
+        # crop the image and update the intrinsics accordingly
+        x, y, w, h = roi
+        image = image[y : y + h, x : x + w]
+        if "depth_image" in data:
+            data["depth_image"] = data["depth_image"][y : y + h, x : x + w]
+        if "mask" in data:
+            mask = data["mask"].numpy()
+            mask = mask.astype(np.uint8) * 255
+            if np.any(distortion_params):
+                mask = cv2.undistort(mask, K, distortion_params, None, newK)  # type: ignore
+            mask = mask[y : y + h, x : x + w]
+            mask = torch.from_numpy(mask).bool()
+        K = newK
+
+    elif camera.camera_type.item() == CameraType.FISHEYE.value:
+        distortion_params = np.array(
+            [distortion_params[0], distortion_params[1], distortion_params[2], distortion_params[3]]
+        )
+        newK = cv2.fisheye.estimateNewCameraMatrixForUndistortRectify(
+            K, distortion_params, (image.shape[1], image.shape[0]), np.eye(3), balance=0
+        )
+        map1, map2 = cv2.fisheye.initUndistortRectifyMap(
+            K, distortion_params, np.eye(3), newK, (image.shape[1], image.shape[0]), cv2.CV_32FC1
+        )
+        # and then remap:
+        image = cv2.remap(image, map1, map2, interpolation=cv2.INTER_LINEAR)
+        if "mask" in data:
+            mask = data["mask"].numpy()
+            mask = mask.astype(np.uint8) * 255
+            mask = cv2.fisheye.undistortImage(mask, K, distortion_params, None, newK)
+            mask = torch.from_numpy(mask).bool()
+        K = newK
+    elif camera.camera_type.item() == CameraType.FISHEYE624.value:
+        fisheye624_params = torch.cat(
+            [camera.fx, camera.fy, camera.cx, camera.cy, torch.from_numpy(distortion_params)], dim=0
+        )
+        assert fisheye624_params.shape == (16,)
+        assert (
+            "mask" not in data
+            and camera.metadata is not None
+            and "fisheye_crop_radius" in camera.metadata
+            and isinstance(camera.metadata["fisheye_crop_radius"], float)
+        )
+        fisheye_crop_radius = camera.metadata["fisheye_crop_radius"]
+
+        # Approximate the FOV of the unmasked region of the camera.
+        upper, lower, left, right = fisheye624_unproject_helper(
+            torch.tensor(
+                [
+                    [camera.cx, camera.cy - fisheye_crop_radius],
+                    [camera.cx, camera.cy + fisheye_crop_radius],
+                    [camera.cx - fisheye_crop_radius, camera.cy],
+                    [camera.cx + fisheye_crop_radius, camera.cy],
+                ],
+                dtype=torch.float32,
+            )[None],
+            params=fisheye624_params[None],
+        ).squeeze(dim=0)
+        fov_radians = torch.max(
+            torch.acos(torch.sum(upper * lower / torch.linalg.norm(upper) / torch.linalg.norm(lower))),
+            torch.acos(torch.sum(left * right / torch.linalg.norm(left) / torch.linalg.norm(right))),
+        )
+
+        # Heuristics to determine parameters of an undistorted image.
+        undist_h = int(fisheye_crop_radius * 2)
+        undist_w = int(fisheye_crop_radius * 2)
+        undistort_focal = undist_h / (2 * torch.tan(fov_radians / 2.0))
+        undist_K = torch.eye(3)
+        undist_K[0, 0] = undistort_focal  # fx
+        undist_K[1, 1] = undistort_focal  # fy
+        undist_K[0, 2] = (undist_w - 1) / 2.0  # cx; for a 1x1 image, center should be at (0, 0).
+        undist_K[1, 2] = (undist_h - 1) / 2.0  # cy
+
+        # Undistorted 2D coordinates -> rays -> reproject to distorted UV coordinates.
+        undist_uv_homog = torch.stack(
+            [
+                *torch.meshgrid(
+                    torch.arange(undist_w, dtype=torch.float32),
+                    torch.arange(undist_h, dtype=torch.float32),
+                ),
+                torch.ones((undist_w, undist_h), dtype=torch.float32),
+            ],
+            dim=-1,
+        )
+        assert undist_uv_homog.shape == (undist_w, undist_h, 3)
+        dist_uv = (
+            fisheye624_project(
+                xyz=(
+                    torch.einsum(
+                        "ij,bj->bi",
+                        torch.linalg.inv(undist_K),
+                        undist_uv_homog.reshape((undist_w * undist_h, 3)),
+                    )[None]
+                ),
+                params=fisheye624_params[None, :],
+            )
+            .reshape((undist_w, undist_h, 2))
+            .numpy()
+        )
+        map1 = dist_uv[..., 1]
+        map2 = dist_uv[..., 0]
+
+        # Use correspondence to undistort image.
+        image = cv2.remap(image, map1, map2, interpolation=cv2.INTER_LINEAR)
+
+        # Compute undistorted mask as well.
+        dist_h = camera.height.item()
+        dist_w = camera.width.item()
+        mask = np.mgrid[:dist_h, :dist_w]
+        mask[0, ...] -= dist_h // 2
+        mask[1, ...] -= dist_w // 2
+        mask = np.linalg.norm(mask, axis=0) < fisheye_crop_radius
+        mask = torch.from_numpy(
+            cv2.remap(
+                mask.astype(np.uint8) * 255,
+                map1,
+                map2,
+                interpolation=cv2.INTER_LINEAR,
+                borderMode=cv2.BORDER_CONSTANT,
+                borderValue=0,
+            )
+            / 255.0
+        ).bool()
+        K = undist_K.numpy()
+    else:
+        raise NotImplementedError("Only perspective and fisheye cameras are supported")
+
+    return K, image, mask

nerfstudio/data/dataparsers/colmap_dataparser.py

@@ -64,6 +64,12 @@ class ColmapDataParserConfig(DataParserConfig):
     """The method to use to center the poses."""
     auto_scale_poses: bool = True
     """Whether to automatically scale the poses to fit in +/- 1 bounding box."""
+    assume_colmap_world_coordinate_convention: bool = True
+    """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. Therefore, we usually need to apply an extra
+    transform when orientation_method=none. This parameter has no effects if orientation_method is set other than none.
+    When this parameter is set to False, no extra transform is applied when reading data from colmap.
+    """
     eval_mode: Literal["fraction", "filename", "interval", "all"] = "interval"
     """
     The method to use for splitting the dataset into train and eval.
@@ -87,8 +93,9 @@ class ColmapDataParserConfig(DataParserConfig):
     """Path to depth maps directory. If not set, depths are not loaded."""
     colmap_path: Path = Path("colmap/sparse/0")
     """Path to the colmap reconstruction directory relative to the data path."""
-    load_3D_points: bool = False
-    """Whether to load the 3D points from the colmap reconstruction."""
+    load_3D_points: bool = True
+    """Whether to load the 3D points from the colmap reconstruction. This is helpful for Gaussian splatting and
+    generally unused otherwise, but it's typically harmless so we default to True."""
     max_2D_matches_per_3D_point: int = 0
     """Maximum number of 2D matches per 3D point. If set to -1, all 2D matches are loaded. If set to 0, no 2D matches are loaded."""
 
@@ -154,10 +161,10 @@ class ColmapDataParser(DataParser):
             c2w = np.linalg.inv(w2c)
             # Convert from COLMAP's camera coordinate system (OpenCV) to ours (OpenGL)
             c2w[0:3, 1:3] *= -1
-            # Why do we want to flip Z with a handedness transform?
-            # See https://github.com/nerfstudio-project/nerfstudio/issues/1504
-            c2w = c2w[np.array([1, 0, 2, 3]), :]
-            c2w[2, :] *= -1
+            if self.config.assume_colmap_world_coordinate_convention:
+                # world coordinate transform: map colmap gravity guess (-y) to nerfstudio convention (+z)
+                c2w = c2w[np.array([0, 2, 1, 3]), :]
+                c2w[2, :] *= -1
 
             frame = {
                 "file_path": (self.config.data / self.config.images_path / im_data.name).as_posix(),
@@ -181,12 +188,12 @@ class ColmapDataParser(DataParser):
 
         out = {}
         out["frames"] = frames
-        # Why do we want to flip Z with a handedness transform?
-        # See https://github.com/nerfstudio-project/nerfstudio/issues/1504
-        applied_transform = np.eye(4)[:3, :]
-        applied_transform = applied_transform[np.array([1, 0, 2]), :]
-        applied_transform[2, :] *= -1
-        out["applied_transform"] = applied_transform.tolist()
+        if self.config.assume_colmap_world_coordinate_convention:
+            # world coordinate transform: map colmap gravity guess (-y) to nerfstudio convention (+z)
+            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()
         out["camera_model"] = camera_model
         assert len(frames) > 0, "No images found in the colmap model"
         return out
@@ -304,7 +311,6 @@ class ColmapDataParser(DataParser):
         if self.config.auto_scale_poses:
             scale_factor /= float(torch.max(torch.abs(poses[:, :3, 3])))
         scale_factor *= self.config.scale_factor
-
         poses[:, :3, 3] *= scale_factor
 
         # Choose image_filenames and poses based on split, but after auto orient and scaling the poses.

nerfstudio/data/dataparsers/nerfstudio_dataparser.py

@@ -73,6 +73,8 @@ class NerfstudioDataParserConfig(DataParserConfig):
     """The interval between frames to use for eval. Only used when eval_mode is eval-interval."""
     depth_unit_scale_factor: float = 1e-3
     """Scales the depth values to meters. Default value is 0.001 for a millimeter to meter conversion."""
+    load_3D_points: bool = False
+    """Whether to load the 3D points from the colmap reconstruction."""
 
 
 @dataclass
@@ -310,20 +312,94 @@ class Nerfstudio(DataParser):
         assert self.downscale_factor is not None
         cameras.rescale_output_resolution(scaling_factor=1.0 / self.downscale_factor)
 
+        # The naming is somewhat confusing, but:
+        # - transform_matrix contains the transformation to dataparser output coordinates from saved coordinates.
+        # - dataparser_transform_matrix contains the transformation to dataparser output coordinates from original data coordinates.
+        # - applied_transform contains the transformation to saved coordinates from original data coordinates.
+        applied_transform = None
+        colmap_path = self.config.data / "colmap/sparse/0"
         if "applied_transform" in meta:
             applied_transform = torch.tensor(meta["applied_transform"], dtype=transform_matrix.dtype)
-            transform_matrix = transform_matrix @ torch.cat(
+        elif colmap_path.exists():
+            # For converting from colmap, this was the effective value of applied_transform that was being
+            # used before we added the applied_transform field to the output dataformat.
+            meta["applied_transform"] = [[0, 1, 0, 0], [1, 0, 0, 0], [0, 0, -1, 0]]
+            applied_transform = torch.tensor(meta["applied_transform"], dtype=transform_matrix.dtype)
+
+        if applied_transform is not None:
+            dataparser_transform_matrix = transform_matrix @ torch.cat(
                 [applied_transform, torch.tensor([[0, 0, 0, 1]], dtype=transform_matrix.dtype)], 0
             )
+        else:
+            dataparser_transform_matrix = transform_matrix
+
         if "applied_scale" in meta:
             applied_scale = float(meta["applied_scale"])
             scale_factor *= applied_scale
 
-        # Load 3D points
+        # reinitialize metadata for dataparser_outputs
         metadata = {}
-        if "ply_file_path" in meta:
-            ply_file_path = data_dir / meta["ply_file_path"]
-            metadata.update(self._load_3D_points(ply_file_path, transform_matrix, scale_factor))
+
+        # _generate_dataparser_outputs might be called more than once so we check if we already loaded the point cloud
+        try:
+            self.prompted_user
+        except AttributeError:
+            self.prompted_user = False
+
+        # Load 3D points
+        if self.config.load_3D_points:
+            if "ply_file_path" in meta:
+                ply_file_path = data_dir / meta["ply_file_path"]
+
+            elif colmap_path.exists():
+                from rich.prompt import Confirm
+
+                # check if user wants to make a point cloud from colmap points
+                if not self.prompted_user:
+                    self.create_pc = Confirm.ask(
+                        "load_3D_points is true, but the dataset was processed with an outdated ns-process-data that didn't convert colmap points to .ply! Update the colmap dataset automatically?"
+                    )
+
+                if self.create_pc:
+                    import json
+
+                    from nerfstudio.process_data.colmap_utils import create_ply_from_colmap
+
+                    with open(self.config.data / "transforms.json") as f:
+                        transforms = json.load(f)
+
+                    # Update dataset if missing the applied_transform field.
+                    if "applied_transform" not in transforms:
+                        transforms["applied_transform"] = meta["applied_transform"]
+
+                    ply_filename = "sparse_pc.ply"
+                    create_ply_from_colmap(
+                        filename=ply_filename,
+                        recon_dir=colmap_path,
+                        output_dir=self.config.data,
+                        applied_transform=applied_transform,
+                    )
+                    ply_file_path = data_dir / ply_filename
+                    transforms["ply_file_path"] = ply_filename
+
+                    # This was the applied_transform value
+
+                    with open(self.config.data / "transforms.json", "w", encoding="utf-8") as f:
+                        json.dump(transforms, f, indent=4)
+                else:
+                    ply_file_path = None
+            else:
+                if not self.prompted_user:
+                    CONSOLE.print(
+                        "[bold yellow]Warning: load_3D_points set to true but no point cloud found. gaussian-splatting models will use random point cloud initialization."
+                    )
+                ply_file_path = None
+
+            if ply_file_path:
+                sparse_points = self._load_3D_points(ply_file_path, transform_matrix, scale_factor)
+                if sparse_points is not None:
+                    metadata.update(sparse_points)
+            self.prompted_user = True
 
         dataparser_outputs = DataparserOutputs(
             image_filenames=image_filenames,
@@ -331,7 +407,7 @@ class Nerfstudio(DataParser):
             scene_box=scene_box,
             mask_filenames=mask_filenames if len(mask_filenames) > 0 else None,
             dataparser_scale=scale_factor,
-            dataparser_transform=transform_matrix,
+            dataparser_transform=dataparser_transform_matrix,
             metadata={
                 "depth_filenames": depth_filenames if len(depth_filenames) > 0 else None,
                 "depth_unit_scale_factor": self.config.depth_unit_scale_factor,
@@ -341,10 +417,24 @@ class Nerfstudio(DataParser):
         return dataparser_outputs
 
     def _load_3D_points(self, ply_file_path: Path, transform_matrix: torch.Tensor, scale_factor: float):
+        """Loads point clouds positions and colors from .ply
+
+        Args:
+            ply_file_path: Path to .ply file
+            transform_matrix: Matrix to transform world coordinates
+            scale_factor: How much to scale the camera origins by.
+
+        Returns:
+            A dictionary of points: points3D_xyz and colors: points3D_rgb
+        """
         import open3d as o3d  # Importing open3d is slow, so we only do it if we need it.
 
         pcd = o3d.io.read_point_cloud(str(ply_file_path))
 
+        # if no points found don't read in an initial point cloud
+        if len(pcd.points) == 0:
+            return None
+
         points3D = torch.from_numpy(np.asarray(pcd.points, dtype=np.float32))
         points3D = (
             torch.cat(

nerfstudio/models/gaussian_splatting.py

@@ -158,16 +158,18 @@ class GaussianSplattingModel(Model):
 
     config: GaussianSplattingModelConfig
 
-    def __init__(self, *args, **kwargs):
-        if "seed_points" in kwargs:
-            self.seed_pts = kwargs["seed_points"]
-        else:
-            self.seed_pts = None
+    def __init__(
+        self,
+        *args,
+        seed_points: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        **kwargs,
+    ):
+        self.seed_points = seed_points
         super().__init__(*args, **kwargs)
 
     def populate_modules(self):
-        if self.seed_pts is not None and not self.config.random_init:
-            self.means = torch.nn.Parameter(self.seed_pts[0])  # (Location, Color)
+        if self.seed_points is not None and not self.config.random_init:
+            self.means = torch.nn.Parameter(self.seed_points[0])  # (Location, Color)
         else:
             self.means = torch.nn.Parameter((torch.rand((500000, 3)) - 0.5) * 10)
         self.xys_grad_norm = None
@@ -180,14 +182,19 @@ class GaussianSplattingModel(Model):
         self.quats = torch.nn.Parameter(random_quat_tensor(self.num_points))
         dim_sh = num_sh_bases(self.config.sh_degree)
 
-        if self.seed_pts is not None and not self.config.random_init:
-            shs = torch.zeros((self.seed_pts[1].shape[0], dim_sh, 3)).float().cuda()
+        if (
+            self.seed_points is not None
+            and not self.config.random_init
+            # We can have colors without points.
+            and self.seed_points[1].shape[0] > 0
+        ):
+            shs = torch.zeros((self.seed_points[1].shape[0], dim_sh, 3)).float().cuda()
             if self.config.sh_degree > 0:
-                shs[:, 0, :3] = RGB2SH(self.seed_pts[1] / 255)
+                shs[:, 0, :3] = RGB2SH(self.seed_points[1] / 255)
                 shs[:, 1:, 3:] = 0.0
             else:
                 CONSOLE.log("use color only optimization with sigmoid activation")
-                shs[:, 0, :3] = torch.logit(self.seed_pts[1] / 255, eps=1e-10)
+                shs[:, 0, :3] = torch.logit(self.seed_points[1] / 255, eps=1e-10)
             self.features_dc = torch.nn.Parameter(shs[:, 0, :])
             self.features_rest = torch.nn.Parameter(shs[:, 1:, :])
         else:
@@ -754,8 +761,18 @@ class GaussianSplattingModel(Model):
             metrics_dict: dictionary of metrics, some of which we can use for loss
         """
         gt_img = self.get_gt_img(batch["image"])
-        Ll1 = torch.abs(gt_img - outputs["rgb"]).mean()
-        simloss = 1 - self.ssim(gt_img.permute(2, 0, 1)[None, ...], outputs["rgb"].permute(2, 0, 1)[None, ...])
+        pred_img = outputs["rgb"]
+
+        # Set masked part of both ground-truth and rendered image to black.
+        # This is a little bit sketchy for the SSIM loss.
+        if "mask" in batch:
+            assert batch["mask"].shape == gt_img.shape[:2] == pred_img.shape[:2]
+            mask = batch["mask"][..., None].to(self.device)
+            gt_img = gt_img * mask
+            pred_img = pred_img * mask
+
+        Ll1 = torch.abs(gt_img - pred_img).mean()
+        simloss = 1 - self.ssim(gt_img.permute(2, 0, 1)[None, ...], pred_img.permute(2, 0, 1)[None, ...])
         if self.config.use_scale_regularization and self.step % 10 == 0:
             scale_exp = torch.exp(self.scales)
             scale_reg = (

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 Any, Dict, Literal, Optional
+from typing import Any, Dict, Literal, Optional, Union
 
 import appdirs
 import cv2
@@ -34,6 +34,7 @@ from nerfstudio.data.utils.colmap_parsing_utils import (
     read_cameras_binary,
     read_images_binary,
     read_points3D_binary,
+    read_points3D_text,
 )
 from nerfstudio.process_data.process_data_utils import CameraModel
 from nerfstudio.utils import colormaps
@@ -391,6 +392,8 @@ def colmap_to_json(
     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,
 ) -> int:
     """Converts COLMAP's cameras.bin and images.bin to a JSON file.
 
@@ -401,7 +404,9 @@ def colmap_to_json(
         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 registered images.
     """
@@ -430,8 +435,9 @@ def colmap_to_json(
         c2w = np.linalg.inv(w2c)
         # Convert from COLMAP's camera coordinate system (OpenCV) to ours (OpenGL)
         c2w[0:3, 1:3] *= -1
-        c2w = c2w[np.array([1, 0, 2, 3]), :]
-        c2w[2, :] *= -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:
@@ -455,10 +461,22 @@ def colmap_to_json(
     out = parse_colmap_camera_params(cam_id_to_camera[1])
     out["frames"] = frames
 
-    applied_transform = np.eye(4)[:3, :]
-    applied_transform = applied_transform[np.array([1, 0, 2]), :]
-    applied_transform[2, :] *= -1
-    out["applied_transform"] = applied_transform.tolist()
+    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)
@@ -638,3 +656,49 @@ def get_matching_summary(num_initial_frames: int, num_matched_frames: int) -> st
         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."
+
+
+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
+    """
+    if (recon_dir / "points3D.bin").exists():
+        colmap_points = read_points3D_binary(recon_dir / "points3D.bin")
+    elif (recon_dir / "points3D.txt").exists():
+        colmap_points = read_points3D_text(recon_dir / "points3D.txt")
+    else:
+        raise ValueError(f"Could not find points3D.txt or points3D.bin in {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")

nerfstudio/scripts/datasets/process_project_aria.py

@@ -17,15 +17,17 @@ import sys
 import threading
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Dict, List
+from typing import Any, Dict, List, cast
 
 import numpy as np
+import open3d as o3d
 import tyro
 from PIL import Image
 
 try:
     from projectaria_tools.core import mps
     from projectaria_tools.core.data_provider import VrsDataProvider, create_vrs_data_provider
+    from projectaria_tools.core.mps.utils import filter_points_from_confidence
     from projectaria_tools.core.sophus import SE3
 except ImportError:
     print("projectaria_tools import failed, please install with pip3 install projectaria-tools'[all]'")
@@ -219,6 +221,21 @@ class ProcessProjectAria:
             "fisheye_crop_radius": rgb_valid_radius,
         }
 
+        # save global point cloud, which is useful for Gaussian Splatting.
+        points_path = self.mps_data_dir / "global_points.csv.gz"
+        if points_path.exists():
+            print("Found global points, saving to PLY...")
+            points_data = mps.read_global_point_cloud(str(points_path))  # type: ignore
+            points_data = filter_points_from_confidence(points_data)
+            pcd = o3d.geometry.PointCloud()
+            pcd.points = o3d.utility.Vector3dVector(np.array([cast(Any, it).position_world for it in points_data]))
+            ply_file_path = self.output_dir / "global_points.ply"
+            o3d.io.write_point_cloud(str(ply_file_path), pcd)
+
+            nerfstudio_frames["ply_file_path"] = "global_points.ply"
+        else:
+            print("No global points found!")
+
         # write the json out to disk as transforms.json
         print("Writing transforms.json")
         transform_file = self.output_dir / "transforms.json"

nerfstudio/scripts/github/run_actions.py

@@ -56,11 +56,14 @@ def run_github_actions_file(filename: str, continue_on_fail: bool = False):
 
     for step in steps:
         if "name" in step and step["name"] in LOCAL_TESTS:
-            compressed = step["run"].replace("\n", ";").replace("\\", "")
-            if "ruff check" in compressed:
-                curr_command = f"{compressed} --fix"
-            else:
-                curr_command = compressed.replace("--check", "")
+            curr_command = step["run"].replace("\n", ";").replace("\\", "")
+            if curr_command.startswith("ruff"):
+                if "ruff check" in curr_command:
+                    curr_command = f"{curr_command} --fix"
+
+                curr_command = curr_command.replace(" --check", "")
+                curr_command = curr_command.replace(" --diff", "")
+                curr_command = curr_command.replace(" --output-format=github", "")
 
             CONSOLE.line()
             CONSOLE.rule(f"[bold green]Running: {curr_command}")

nerfstudio/viewer/viewer.py

@@ -32,6 +32,7 @@ from nerfstudio.cameras.cameras import CameraType
 from nerfstudio.configs import base_config as cfg
 from nerfstudio.data.datasets.base_dataset import InputDataset
 from nerfstudio.models.base_model import Model
+from nerfstudio.models.gaussian_splatting import GaussianSplattingModel
 from nerfstudio.pipelines.base_pipeline import Pipeline
 from nerfstudio.utils.decorators import check_main_thread, decorate_all
 from nerfstudio.utils.writer import GLOBAL_BUFFER, EventName
@@ -247,6 +248,17 @@ class Viewer:
         for c in self.viewer_controls:
             c._setup(self)
 
+        # Diagnostics for Gaussian Splatting: where the points are at the start of training.
+        # This is hidden by default, it can be shown from the Viser UI's scene tree table.
+        if isinstance(pipeline.model, GaussianSplattingModel):
+            self.viser_server.add_point_cloud(
+                "/gaussian_splatting_initial_points",
+                points=pipeline.model.means.numpy(force=True) * VISER_NERFSTUDIO_SCALE_RATIO,
+                colors=(255, 0, 0),
+                point_size=0.01,
+                point_shape="circle",
+                visible=False,  # Hidden by default.
+            )
         self.ready = True
 
     def toggle_pause_button(self) -> None:

tests/process_data/test_process_images.py

@@ -12,9 +12,11 @@ from nerfstudio.data.dataparsers.nerfstudio_dataparser import NerfstudioDataPars
 from nerfstudio.data.utils.colmap_parsing_utils import (
     Camera,
     Image as ColmapImage,
+    Point3D,
     qvec2rotmat,
     write_cameras_binary,
     write_images_binary,
+    write_points3D_binary,
 )
 from nerfstudio.process_data.images_to_nerfstudio_dataset import ImagesToNerfstudioDataset
 
@@ -50,6 +52,19 @@ def test_process_images_skip_colmap(tmp_path: Path):
         {1: Camera(1, "OPENCV", width, height, [110, 110, 50, 75, 0, 0, 0, 0, 0, 0])},
         sparse_path / "cameras.bin",
     )
+    write_points3D_binary(
+        {
+            1: Point3D(
+                id=1,
+                xyz=np.array([0, 0, 0]),
+                rgb=np.array([0, 0, 0]),
+                error=np.array([0]),
+                image_ids=np.array([1]),
+                point2D_idxs=np.array([0]),
+            ),
+        },
+        sparse_path / "points3D.bin",
+    )
     frames = {}
     num_frames = 10
     qvecs = random_quaternion(num_frames)

tests/utils/test_tensor_dataclass.py

@@ -179,6 +179,12 @@ def test_iter():
         assert batch.b.shape == (4, 5)
 
 
+def test_non_tensor():
+    """Test iterating over tensor dataclass"""
+    # We shouldn't throw away non-dataclass values.
+    assert DummyTensorDataclass(a=torch.ones((3, 10)), b={"k": 2}, c=None).b == {"k": 2}  # type: ignore
+
+
 if __name__ == "__main__":
     test_init()
     test_broadcasting()