diff --git a/nerfstudio/data/datamanagers/full_images_datamanager.py b/nerfstudio/data/datamanagers/full_images_datamanager.py index 6e041b33d..772139ce5 100644 --- a/nerfstudio/data/datamanagers/full_images_datamanager.py +++ b/nerfstudio/data/datamanagers/full_images_datamanager.py @@ -30,10 +30,12 @@ from pathlib import Path from typing import Dict, ForwardRef, Generic, List, Literal, Optional, Tuple, Type, Union, cast, get_args, get_origin import cv2 +import fpsample import numpy as np import torch from rich.progress import track from torch.nn import Parameter +from typing_extensions import assert_never from nerfstudio.cameras.camera_utils import fisheye624_project, fisheye624_unproject_helper from nerfstudio.cameras.cameras import Cameras, CameraType @@ -61,12 +63,22 @@ class FullImageDatamanagerConfig(DataManagerConfig): new images. If -1, never pick new images.""" eval_image_indices: Optional[Tuple[int, ...]] = (0,) """Specifies the image indices to use during eval; if None, uses all.""" - cache_images: Literal["cpu", "gpu"] = "cpu" + cache_images: Literal["cpu", "gpu"] = "gpu" """Whether to cache images in memory. If "cpu", caches on cpu. If "gpu", caches on device.""" cache_images_type: Literal["uint8", "float32"] = "float32" """The image type returned from manager, caching images in uint8 saves memory""" max_thread_workers: Optional[int] = None """The maximum number of threads to use for caching images. If None, uses all available threads.""" + train_cameras_sampling_strategy: Literal["random", "fps"] = "random" + """Specifies which sampling strategy is used to generate train cameras, 'random' means sampling + uniformly random without replacement, 'fps' means farthest point sampling which is helpful to reduce the artifacts + due to oversampling subsets of cameras that are very close to each other.""" + train_cameras_sampling_seed: int = 42 + """Random seed for sampling train cameras. Fixing seed may help reduce variance of trained models across + different runs.""" + fps_reset_every: int = 100 + """The number of iterations before one resets fps sampler repeatly, which is essentially drawing fps_reset_every + samples from the pool of all training cameras without replacement before a new round of sampling starts.""" class FullImageDatamanager(DataManager, Generic[TDataset]): @@ -115,7 +127,6 @@ class FullImageDatamanager(DataManager, Generic[TDataset]): style="bold yellow", ) self.config.cache_images = "cpu" - self.cached_train, self.cached_eval = self.cache_images(self.config.cache_images) self.exclude_batch_keys_from_device = self.train_dataset.exclude_batch_keys_from_device if self.config.masks_on_gpu is True: self.exclude_batch_keys_from_device.remove("mask") @@ -123,110 +134,132 @@ class FullImageDatamanager(DataManager, Generic[TDataset]): self.exclude_batch_keys_from_device.remove("image") # Some logic to make sure we sample every camera in equal amounts - self.train_unseen_cameras = [i for i in range(len(self.train_dataset))] + self.train_unseen_cameras = self.sample_train_cameras() self.eval_unseen_cameras = [i for i in range(len(self.eval_dataset))] assert len(self.train_unseen_cameras) > 0, "No data found in dataset" super().__init__() - def cache_images(self, cache_images_option): - cached_train = [] - cached_eval = [] - - def process_train_data(idx): - # cv2.undistort the images / cameras - data = self.train_dataset.get_data(idx, image_type=self.config.cache_images_type) - camera = self.train_dataset.cameras[idx].reshape(()) - K = camera.get_intrinsics_matrices().numpy() - if camera.distortion_params is None: - return data - distortion_params = camera.distortion_params.numpy() - image = data["image"].numpy() - - 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 - - self.train_dataset.cameras.fx[idx] = float(K[0, 0]) - self.train_dataset.cameras.fy[idx] = float(K[1, 1]) - self.train_dataset.cameras.cx[idx] = float(K[0, 2]) - self.train_dataset.cameras.cy[idx] = float(K[1, 2]) - self.train_dataset.cameras.width[idx] = image.shape[1] - self.train_dataset.cameras.height[idx] = image.shape[0] - return data - - def process_eval_data(idx): - # cv2.undistort the images / cameras - data = self.eval_dataset.get_data(idx, image_type=self.config.cache_images_type) - camera = self.eval_dataset.cameras[idx].reshape(()) - K = camera.get_intrinsics_matrices().numpy() - if camera.distortion_params is None: - return data - distortion_params = camera.distortion_params.numpy() - image = data["image"].numpy() - - 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 - - self.eval_dataset.cameras.fx[idx] = float(K[0, 0]) - self.eval_dataset.cameras.fy[idx] = float(K[1, 1]) - self.eval_dataset.cameras.cx[idx] = float(K[0, 2]) - self.eval_dataset.cameras.cy[idx] = float(K[1, 2]) - self.eval_dataset.cameras.width[idx] = image.shape[1] - self.eval_dataset.cameras.height[idx] = image.shape[0] - return data - - CONSOLE.log("Caching / undistorting train images") - with ThreadPoolExecutor(max_workers=2) as executor: - cached_train = list( - track( - executor.map( - process_train_data, - range(len(self.train_dataset)), - ), - description="Caching / undistorting train images", - transient=True, - total=len(self.train_dataset), - ) + def sample_train_cameras(self): + """Return a list of camera indices sampled using the strategy specified by + self.config.train_cameras_sampling_strategy""" + num_train_cameras = len(self.train_dataset) + if self.config.train_cameras_sampling_strategy == "random": + if not hasattr(self, "random_generator"): + self.random_generator = random.Random(self.config.train_cameras_sampling_seed) + indices = list(range(num_train_cameras)) + self.random_generator.shuffle(indices) + return indices + elif self.config.train_cameras_sampling_strategy == "fps": + if not hasattr(self, "train_unsampled_epoch_count"): + np.random.seed(self.config.train_cameras_sampling_seed) # fix random seed of fpsample + self.train_unsampled_epoch_count = np.zeros(num_train_cameras) + camera_origins = self.train_dataset.cameras.camera_to_worlds[..., 3].numpy() + # We concatenate camera origins with weighted train_unsampled_epoch_count because we want to + # increase the chance to sample camera that hasn't been sampled in consecutive epochs previously. + # We assume the camera origins are also rescaled, so the weight 0.1 is relative to the scale of scene + data = np.concatenate( + (camera_origins, 0.1 * np.expand_dims(self.train_unsampled_epoch_count, axis=-1)), axis=-1 ) - - CONSOLE.log("Caching / undistorting eval images") - with ThreadPoolExecutor(max_workers=2) as executor: - cached_eval = list( - track( - executor.map( - process_eval_data, - range(len(self.eval_dataset)), - ), - description="Caching / undistorting eval images", - transient=True, - total=len(self.eval_dataset), + n = self.config.fps_reset_every + if num_train_cameras < n: + CONSOLE.log( + f"num_train_cameras={num_train_cameras} is smaller than fps_reset_ever={n}, the behavior of " + "camera sampler will be very similar to sampling random without replacement (default setting)." ) - ) + n = num_train_cameras + kdline_fps_samples_idx = fpsample.bucket_fps_kdline_sampling(data, n, h=3) - if cache_images_option == "gpu": - for cache in cached_train: - cache["image"] = cache["image"].to(self.device) - if "mask" in cache: - cache["mask"] = cache["mask"].to(self.device) - for cache in cached_eval: - cache["image"] = cache["image"].to(self.device) - if "mask" in cache: - cache["mask"] = cache["mask"].to(self.device) + self.train_unsampled_epoch_count += 1 + self.train_unsampled_epoch_count[kdline_fps_samples_idx] = 0 + return kdline_fps_samples_idx.tolist() else: - for cache in cached_train: - cache["image"] = cache["image"].pin_memory() - if "mask" in cache: - cache["mask"] = cache["mask"].pin_memory() - for cache in cached_eval: - cache["image"] = cache["image"].pin_memory() - if "mask" in cache: - cache["mask"] = cache["mask"].pin_memory() + raise ValueError(f"Unknown train camera sampling strategy: {self.config.train_cameras_sampling_strategy}") - return cached_train, cached_eval + @cached_property + def cached_train(self) -> List[Dict[str, torch.Tensor]]: + """Get the training images. Will load and undistort the images the + first time this (cached) property is accessed.""" + return self._load_images("train", cache_images_device=self.config.cache_images) + + @cached_property + def cached_eval(self) -> List[Dict[str, torch.Tensor]]: + """Get the eval images. Will load and undistort the images the + first time this (cached) property is accessed.""" + return self._load_images("eval", cache_images_device=self.config.cache_images) + + def _load_images( + self, split: Literal["train", "eval"], cache_images_device: Literal["cpu", "gpu"] + ) -> List[Dict[str, torch.Tensor]]: + undistorted_images: List[Dict[str, torch.Tensor]] = [] + + # Which dataset? + if split == "train": + dataset = self.train_dataset + elif split == "eval": + dataset = self.eval_dataset + else: + assert_never(split) + + def undistort_idx(idx: int) -> Dict[str, torch.Tensor]: + data = dataset.get_data(idx, image_type=self.config.cache_images_type) + camera = dataset.cameras[idx].reshape(()) + assert data["image"].shape[1] == camera.width.item() and data["image"].shape[0] == camera.height.item(), ( + f'The size of image ({data["image"].shape[1]}, {data["image"].shape[0]}) loaded ' + f'does not match the camera parameters ({camera.width.item(), camera.height.item()})' + ) + if camera.distortion_params is None or torch.all(camera.distortion_params == 0): + return data + K = camera.get_intrinsics_matrices().numpy() + distortion_params = camera.distortion_params.numpy() + image = data["image"].numpy() + + 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 + + dataset.cameras.fx[idx] = float(K[0, 0]) + dataset.cameras.fy[idx] = float(K[1, 1]) + dataset.cameras.cx[idx] = float(K[0, 2]) + dataset.cameras.cy[idx] = float(K[1, 2]) + dataset.cameras.width[idx] = image.shape[1] + dataset.cameras.height[idx] = image.shape[0] + return data + + CONSOLE.log(f"Caching / undistorting {split} images") + with ThreadPoolExecutor(max_workers=2) as executor: + undistorted_images = list( + track( + executor.map( + undistort_idx, + range(len(dataset)), + ), + description=f"Caching / undistorting {split} images", + transient=True, + total=len(dataset), + ) + ) + + # Move to device. + if cache_images_device == "gpu": + for cache in undistorted_images: + cache["image"] = cache["image"].to(self.device) + if "mask" in cache: + cache["mask"] = cache["mask"].to(self.device) + if "depth" in cache: + cache["depth"] = cache["depth"].to(self.device) + self.train_cameras = self.train_dataset.cameras.to(self.device) + elif cache_images_device == "cpu": + for cache in undistorted_images: + cache["image"] = cache["image"].pin_memory() + if "mask" in cache: + cache["mask"] = cache["mask"].pin_memory() + self.train_cameras = self.train_dataset.cameras + else: + assert_never(cache_images_device) + + return undistorted_images def create_train_dataset(self) -> TDataset: """Sets up the data loaders for training""" @@ -282,7 +315,7 @@ class FullImageDatamanager(DataManager, Generic[TDataset]): Pretends to be the dataloader for evaluation, it returns a list of (camera, data) tuples """ image_indices = [i for i in range(len(self.eval_dataset))] - data = deepcopy(self.cached_eval) + data = [d.copy() for d in self.cached_eval] _cameras = deepcopy(self.eval_dataset.cameras).to(self.device) cameras = [] for i in image_indices: @@ -299,23 +332,31 @@ class FullImageDatamanager(DataManager, Generic[TDataset]): return {} def get_train_rays_per_batch(self): - # TODO: fix this to be the resolution of the last image rendered - return 800 * 800 + """Returns resolution of the image returned from datamanager.""" + if len(self.cached_train) != 0: + h = self.cached_train[0]["image"].shape[0] + w = self.cached_train[0]["image"].shape[1] + return h * w + else: + return 800 * 800 def next_train(self, step: int) -> Tuple[Cameras, Dict]: """Returns the next training batch Returns a Camera instead of raybundle""" - image_idx = self.train_unseen_cameras.pop(random.randint(0, len(self.train_unseen_cameras) - 1)) + image_idx = self.train_unseen_cameras.pop(0) # Make sure to re-populate the unseen cameras list if we have exhausted it if len(self.train_unseen_cameras) == 0: - self.train_unseen_cameras = [i for i in range(len(self.train_dataset))] + self.train_unseen_cameras = self.sample_train_cameras() - data = deepcopy(self.cached_train[image_idx]) + data = self.cached_train[image_idx] + # We're going to copy to make sure we don't mutate the cached dictionary. + # This can cause a memory leak: https://github.com/nerfstudio-project/nerfstudio/issues/3335 + data = data.copy() data["image"] = data["image"].to(self.device) - assert len(self.train_dataset.cameras.shape) == 1, "Assumes single batch dimension" - camera = self.train_dataset.cameras[image_idx : image_idx + 1].to(self.device) + assert len(self.train_cameras.shape) == 1, "Assumes single batch dimension" + camera = self.train_cameras[image_idx : image_idx + 1].to(self.device) if camera.metadata is None: camera.metadata = {} camera.metadata["cam_idx"] = image_idx @@ -325,15 +366,7 @@ class FullImageDatamanager(DataManager, Generic[TDataset]): """Returns the next evaluation batch Returns a Camera instead of raybundle""" - image_idx = self.eval_unseen_cameras.pop(random.randint(0, len(self.eval_unseen_cameras) - 1)) - # Make sure to re-populate the unseen cameras list if we have exhausted it - if len(self.eval_unseen_cameras) == 0: - self.eval_unseen_cameras = [i for i in range(len(self.eval_dataset))] - data = deepcopy(self.cached_eval[image_idx]) - data["image"] = data["image"].to(self.device) - 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 + return self.next_eval_image(step=step) def next_eval_image(self, step: int) -> Tuple[Cameras, Dict]: """Returns the next evaluation batch @@ -345,7 +378,8 @@ class FullImageDatamanager(DataManager, Generic[TDataset]): # Make sure to re-populate the unseen cameras list if we have exhausted it if len(self.eval_unseen_cameras) == 0: self.eval_unseen_cameras = [i for i in range(len(self.eval_dataset))] - data = deepcopy(self.cached_eval[image_idx]) + data = self.cached_eval[image_idx] + data = data.copy() data["image"] = data["image"].to(self.device) 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) @@ -357,6 +391,12 @@ def _undistort_image( ) -> Tuple[np.ndarray, np.ndarray, Optional[torch.Tensor]]: mask = None if camera.camera_type.item() == CameraType.PERSPECTIVE.value: + assert distortion_params[3] == 0, ( + "We doesn't support the 4th Brown parameter for image undistortion, " + "Only k1, k2, k3, p1, p2 can be non-zero." + ) + # because OpenCV expects the order of distortion parameters to be (k1, k2, p1, p2, k3), we need to reorder them + # see https://docs.opencv.org/4.x/dc/dbb/tutorial_py_calibration.html distortion_params = np.array( [ distortion_params[0], @@ -369,6 +409,10 @@ def _undistort_image( 0, ] ) + # because OpenCV expects the pixel coord to be top-left, we need to shift the principal point by 0.5 + # see https://github.com/nerfstudio-project/nerfstudio/issues/3048 + K[0, 2] = K[0, 2] - 0.5 + K[1, 2] = K[1, 2] - 0.5 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 @@ -378,6 +422,9 @@ def _undistort_image( # crop the image and update the intrinsics accordingly x, y, w, h = roi image = image[y : y + h, x : x + w] + # update the principal point based on our cropped region of interest (ROI) + newK[0, 2] -= x + newK[1, 2] -= y if "depth_image" in data: data["depth_image"] = data["depth_image"][y : y + h, x : x + w] if "mask" in data: @@ -387,9 +434,15 @@ def _undistort_image( 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() + if len(mask.shape) == 2: + mask = mask[:, :, None] + newK[0, 2] = newK[0, 2] + 0.5 + newK[1, 2] = newK[1, 2] + 0.5 K = newK elif camera.camera_type.item() == CameraType.FISHEYE.value: + K[0, 2] = K[0, 2] - 0.5 + K[1, 2] = K[1, 2] - 0.5 distortion_params = np.array( [distortion_params[0], distortion_params[1], distortion_params[2], distortion_params[3]] ) @@ -406,6 +459,10 @@ def _undistort_image( mask = mask.astype(np.uint8) * 255 mask = cv2.fisheye.undistortImage(mask, K, distortion_params, None, newK) mask = torch.from_numpy(mask).bool() + if len(mask.shape) == 2: + mask = mask[:, :, None] + newK[0, 2] = newK[0, 2] + 0.5 + newK[1, 2] = newK[1, 2] + 0.5 K = newK elif camera.camera_type.item() == CameraType.FISHEYE624.value: fisheye624_params = torch.cat( @@ -498,6 +555,8 @@ def _undistort_image( ) / 255.0 ).bool()[..., None] + if len(mask.shape) == 2: + mask = mask[:, :, None] assert mask.shape == (undist_h, undist_w, 1) K = undist_K.numpy() else: