diff --git a/.github/workflows/core_code_checks.yml b/.github/workflows/core_code_checks.yml index 0f1fcf0ca..556432bfc 100644 --- a/.github/workflows/core_code_checks.yml +++ b/.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 diff --git a/nerfstudio/cameras/camera_utils.py b/nerfstudio/cameras/camera_utils.py index 7c98ef080..9c1ee0220 100644 --- a/nerfstudio/cameras/camera_utils.py +++ b/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: diff --git a/nerfstudio/cameras/cameras.py b/nerfstudio/cameras/cameras.py index 4202c8c27..e390360b5 100644 --- a/nerfstudio/cameras/cameras.py +++ b/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( [ diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 55ba44eb4..f60450163 100644 --- a/nerfstudio/configs/method_configs.py +++ b/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(), ), diff --git a/nerfstudio/data/datamanagers/full_images_datamanager.py b/nerfstudio/data/datamanagers/full_images_datamanager.py index ec90b4b34..ea0df36e3 100644 --- a/nerfstudio/data/datamanagers/full_images_datamanager.py +++ b/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 diff --git a/nerfstudio/data/dataparsers/colmap_dataparser.py b/nerfstudio/data/dataparsers/colmap_dataparser.py index ce3cf8919..5f5917e70 100644 --- a/nerfstudio/data/dataparsers/colmap_dataparser.py +++ b/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. diff --git a/nerfstudio/data/dataparsers/nerfstudio_dataparser.py b/nerfstudio/data/dataparsers/nerfstudio_dataparser.py index ff8bc0177..892ed0684 100644 --- a/nerfstudio/data/dataparsers/nerfstudio_dataparser.py +++ b/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( diff --git a/nerfstudio/models/gaussian_splatting.py b/nerfstudio/models/gaussian_splatting.py index 154224c36..2365e43f7 100644 --- a/nerfstudio/models/gaussian_splatting.py +++ b/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 = ( diff --git a/nerfstudio/process_data/colmap_utils.py b/nerfstudio/process_data/colmap_utils.py index 27d6e8615..2f2ac3021 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 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") diff --git a/nerfstudio/scripts/datasets/process_project_aria.py b/nerfstudio/scripts/datasets/process_project_aria.py index 10f26653f..23d304cf6 100644 --- a/nerfstudio/scripts/datasets/process_project_aria.py +++ b/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" diff --git a/nerfstudio/scripts/github/run_actions.py b/nerfstudio/scripts/github/run_actions.py index 5a966338c..c758d23b6 100644 --- a/nerfstudio/scripts/github/run_actions.py +++ b/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}") diff --git a/nerfstudio/viewer/viewer.py b/nerfstudio/viewer/viewer.py index c7a881b80..579dbb5cc 100644 --- a/nerfstudio/viewer/viewer.py +++ b/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: diff --git a/tests/process_data/test_process_images.py b/tests/process_data/test_process_images.py index eed454b11..8482676c5 100644 --- a/tests/process_data/test_process_images.py +++ b/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) diff --git a/tests/utils/test_tensor_dataclass.py b/tests/utils/test_tensor_dataclass.py index 8971c83f8..10e21adc0 100644 --- a/tests/utils/test_tensor_dataclass.py +++ b/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()