diff --git a/docs/nerfology/methods/nerf.md b/docs/nerfology/methods/nerf.md index 90e89cd59..4c56f3154 100644 --- a/docs/nerfology/methods/nerf.md +++ b/docs/nerfology/methods/nerf.md @@ -18,7 +18,7 @@ ns-train vanilla-nerf If you have arrived to this site, it is likely that you have at least heard of NeRFs. This page will discuss the original NeRF paper, _"NeRF: Representing Scenes as Neural Radiance Fields for View Synthesis"_ by Mildenhall, Srinivasan, Tancik et al. (2020). -For most tasks, using the original NeRF model is likely not a good choice and hence we provide implementations of various other NeRF related models. It is however useful to understand how NeRF's work as most follow ups follow a similar structure. +For most tasks, using the original NeRF model is likely not a good choice and hence we provide implementations of various other NeRF related models. It is however useful to understand how NeRF's work as most follow ups follow a similar structure and it doesn't require CUDA to execute (useful for stepping through the code with a debugger if you don't have a GPU at hand). The goal is to optimize a volumetric representation of a scene that can be rendered from novel viewpoints. This representation is optimized from a set of images and associated camera poses. diff --git a/docs/nerfology/model_components/visualize_samplers.ipynb b/docs/nerfology/model_components/visualize_samplers.ipynb index 1ead09832..a23500d58 100644 --- a/docs/nerfology/model_components/visualize_samplers.ipynb +++ b/docs/nerfology/model_components/visualize_samplers.ipynb @@ -1,6 +1,7 @@ { "cells": [ { + "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ @@ -46,7 +47,7 @@ ":class: only-dark\n", "```\n", "\n", - "### Heirarchical Sampling\n", + "### Hierarchical Sampling\n", "\n", "It is important to sample the scene where it has content otherwise the reconstruction quality will be reduced. \n", "\n", @@ -7413,7 +7414,7 @@ ], "metadata": { "kernelspec": { - "display_name": "Python 3.8.12 ('nerfactory')", + "display_name": "Python 3", "language": "python", "name": "python3" }, @@ -7427,11 +7428,11 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.8.12" + "version": "3.10.7" }, "vscode": { "interpreter": { - "hash": "34c28001ff35fb390494047002768a8182dcf55b1b11415165e62ea61557ab83" + "hash": "916dbcbb3f70747c44a77c7bcd40155683ae19c65e1c03b4aa3499c5328201f1" } } }, diff --git a/docs/quickstart/custom_dataset.md b/docs/quickstart/custom_dataset.md index ebc6ce313..009cb287c 100644 --- a/docs/quickstart/custom_dataset.md +++ b/docs/quickstart/custom_dataset.md @@ -290,3 +290,35 @@ ns-process-data metashape --data {data directory} --xml {xml file} --output-dir ```bash ns-train nerfacto --data {output directory} ``` + +## Equirectangular data + +Equirectangular data is data that has been taken by a 360 camera such as Insta360. Both equirectangular image sets and videos can be processed by nerfstudio. + +### Images +For a set of equirectangular images, process the data using the following command: + +```bash +ns-process-data images --camera-type equirectangular --images-per-equirect {8, or 14} --data {data directory} --output-dir {output directory} +``` + +The images-per-equirect argument is the number of images that will be sampled from each equirectangular image. We have found that 8 images per equirectangular image is sufficient for most use cases so it defaults to that. However, if you find that there isn't enough detail in the nerf or that colmap is having trouble aligning the images, you can try increasing the number of images per equirectangular image to 14. + +### Videos +For videos we reccomend taking a video with the camera held on top of your head. This will result in any unwanted capturer to just be in the bottom of each frame image and therefore can be cropped out. + +For a video, process the data using the following command: + +```bash +ns-process-data video --camera-type equirectangular --images-per-equirect {8, or 14} --num-frames-target {num equirectangular frames to sample from} --crop-factor {top bottom left right} --data {data directory} --output-dir {output directory} +``` +See the equirectangular images section above for a description of the `--images-per-equirect` argument. + +The `num-frames-target` argument is optional but it is recomended to set it to 3*(seconds of video) frames. For example, if you have a 30 second video, you would use `--num-frames-target 90` (3*30=90). This number was chosen from a bit of experimentation and seems to work well for most videos. It is by no means a hard rule and you can experiment with different values. + +The `crop-factor` argument is optional but often very helpful. This is because equirectangular videos taken by 360 cameras tend to have a portion of the bottom of the image that is the person who was holding the camera over their head. + + + +This obscene mesh of human is obviously not helpful in training a nerf so we can remove it by cropping the bottom 20% of the image. This can be done by using the `--crop-factor 0 0.2 0 0` argument. + diff --git a/docs/quickstart/imgs/equirect_crop.jpg b/docs/quickstart/imgs/equirect_crop.jpg new file mode 100644 index 000000000..584dab831 Binary files /dev/null and b/docs/quickstart/imgs/equirect_crop.jpg differ diff --git a/docs/quickstart/installation.md b/docs/quickstart/installation.md index 2a2db4453..950aefb2a 100644 --- a/docs/quickstart/installation.md +++ b/docs/quickstart/installation.md @@ -140,7 +140,7 @@ To restrict to only CUDA architectures that you have available locally, use the build arg and look up [the compute capability for your GPU](https://developer.nvidia.com/cuda-gpus). For example, here's how to build with support for GeForce 30xx series GPUs: ```bash -docker build --build-arg CUDA_ARCHITECTURES=86 -tag nerfstudio-86 . +docker build --build-arg CUDA_ARCHITECTURES=86 --tag nerfstudio-86 -f Dockerfile . ``` ### Using an interactive container diff --git a/nerfstudio/cameras/camera_utils.py b/nerfstudio/cameras/camera_utils.py index b1dd5a267..7b3712cd0 100644 --- a/nerfstudio/cameras/camera_utils.py +++ b/nerfstudio/cameras/camera_utils.py @@ -229,7 +229,7 @@ def get_interpolated_poses_many( traj = np.stack(traj, axis=0) k_interp = np.stack(k_interp, axis=0) - return torch.tensor(traj), torch.tensor(k_interp) + return torch.tensor(traj, dtype=torch.float32), torch.tensor(k_interp, dtype=torch.float32) def normalize(x: torch.Tensor) -> TensorType[...]: @@ -431,7 +431,7 @@ def rotation_matrix(a: TensorType[3], b: TensorType[3]) -> TensorType[3, 3]: def auto_orient_and_center_poses( poses: TensorType["num_poses":..., 4, 4], method: Literal["pca", "up", "none"] = "up", center_poses: bool = True -) -> Tuple[TensorType["num_poses":..., 3, 4], TensorType[4, 4]]: +) -> Tuple[TensorType["num_poses":..., 3, 4], TensorType[3, 4]]: """Orients and centers the poses. We provide two methods for orientation: pca and up. pca: Orient the poses so that the principal component of the points is aligned with the axes. diff --git a/nerfstudio/data/datamanagers/base_datamanager.py b/nerfstudio/data/datamanagers/base_datamanager.py index dd0109c1d..7ff2f9aaf 100644 --- a/nerfstudio/data/datamanagers/base_datamanager.py +++ b/nerfstudio/data/datamanagers/base_datamanager.py @@ -235,24 +235,39 @@ class DataManager(nn.Module): """Sets up the data manager for evaluation""" @abstractmethod - def next_train(self, step: int) -> Tuple: + def next_train(self, step: int) -> Tuple[RayBundle, Dict]: """Returns the next batch of data from the train data manager. - This will be a tuple of all the information that this data manager outputs. + Args: + step: the step number of the eval image to retrieve + Returns: + A tuple of the ray bundle for the image, and a dictionary of additional batch information + such as the groudtruth image. """ raise NotImplementedError @abstractmethod - def next_eval(self, step: int) -> Tuple: + def next_eval(self, step: int) -> Tuple[RayBundle, Dict]: """Returns the next batch of data from the eval data manager. - This will be a tuple of all the information that this data manager outputs. + Args: + step: the step number of the eval image to retrieve + Returns: + A tuple of the ray bundle for the image, and a dictionary of additional batch information + such as the groudtruth image. """ raise NotImplementedError @abstractmethod - def next_eval_image(self, step: int) -> Tuple: - """Returns the next eval image.""" + def next_eval_image(self, step: int) -> Tuple[int, RayBundle, Dict]: + """Retreive the next eval image. + + Args: + step: the step number of the eval image to retrieve + Returns: + A tuple of the step number, the ray bundle for the image, and a dictionary of + additional batch information such as the groudtruth image. + """ raise NotImplementedError @abstractmethod @@ -358,12 +373,12 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method self.sampler = None self.test_mode = test_mode self.test_split = "test" if test_mode in ["test", "inference"] else "val" - self.dataparser = self.config.dataparser + self.dataparser_config = self.config.dataparser if self.config.data is not None: self.config.dataparser.data = Path(self.config.data) else: self.config.data = self.config.dataparser.data - self.dataparser = self.dataparser.setup() + self.dataparser = self.dataparser_config.setup() self.train_dataparser_outputs = self.dataparser.get_dataparser_outputs(split="train") self.train_dataset = self.create_train_dataset() @@ -453,7 +468,6 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method ) self.eval_dataloader = RandIndicesEvalDataloader( input_dataset=self.eval_dataset, - image_indices=self.config.eval_image_indices, device=self.device, num_workers=self.world_size * 4, ) diff --git a/nerfstudio/data/dataparsers/arkitscenes_dataparser.py b/nerfstudio/data/dataparsers/arkitscenes_dataparser.py index 736f9941b..fdf992ebd 100644 --- a/nerfstudio/data/dataparsers/arkitscenes_dataparser.py +++ b/nerfstudio/data/dataparsers/arkitscenes_dataparser.py @@ -208,4 +208,6 @@ class ARKitScenes(DataParser): frame_pose = np.array(poses_from_traj[str(my_key)]) frame_pose[0:3, 1:3] *= -1 + frame_pose = frame_pose[np.array([1, 0, 2, 3]), :] + frame_pose[2, :] *= -1 return frame_pose diff --git a/nerfstudio/data/dataparsers/minimal_dataparser.py b/nerfstudio/data/dataparsers/minimal_dataparser.py index 2fa2218a4..6d8de4ec5 100644 --- a/nerfstudio/data/dataparsers/minimal_dataparser.py +++ b/nerfstudio/data/dataparsers/minimal_dataparser.py @@ -52,15 +52,17 @@ class MinimalDataParser(DataParser): filepath = self.config.data / f"{split}.npz" data = np.load(filepath, allow_pickle=True) - image_filenames = data["image_filenames"].tolist() - mask_filenames = data["mask_filenames"].tolist() if "mask_filenames" in data.keys() else None + image_filenames = [filepath.parent / path for path in data["image_filenames"].tolist()] + mask_filenames = None + if "mask_filenames" in data.keys(): + mask_filenames = [filepath.parent / path for path in data["mask_filenames"].tolist()] metadata = None if "semantics" in data.keys(): semantics = data["semantics"].item() metadata = { "semantics": Semantics( - filenames=semantics["filenames"].tolist(), + filenames=[filepath.parent / path for path in semantics["filenames"].tolist()], classes=semantics["classes"].tolist(), colors=torch.from_numpy(semantics["colors"]), mask_classes=semantics["mask_classes"].tolist(), @@ -87,11 +89,21 @@ class MinimalDataParser(DataParser): times=torch.from_numpy(camera_np["times"]) if "times" in camera_np.keys() else None, ) + applied_scale = 1.0 + applied_transform = np.eye(4, dtype=np.float32)[:3, :] + if "applied_scale" in data.keys(): + applied_scale = float(data["applied_scale"]) + if "applied_transform" in data.keys(): + applied_transform = data["applied_transform"].astype(np.float32) + assert applied_transform.shape == (3, 4) + dataparser_outputs = DataparserOutputs( image_filenames=image_filenames, cameras=cameras, scene_box=scene_box, mask_filenames=mask_filenames, + dataparser_transform=applied_transform, + dataparser_scale=applied_scale, metadata=metadata, ) return dataparser_outputs diff --git a/nerfstudio/data/dataparsers/nerfstudio_dataparser.py b/nerfstudio/data/dataparsers/nerfstudio_dataparser.py index 035e2da2d..6a23a9ed3 100644 --- a/nerfstudio/data/dataparsers/nerfstudio_dataparser.py +++ b/nerfstudio/data/dataparsers/nerfstudio_dataparser.py @@ -182,22 +182,36 @@ class Nerfstudio(DataParser): You should check that depth_file_path is specified for every frame (or zero frames) in transforms.json. """ - # filter image_filenames and poses based on train/eval split percentage - num_images = len(image_filenames) - num_train_images = math.ceil(num_images * self.config.train_split_fraction) - num_eval_images = num_images - num_train_images - i_all = np.arange(num_images) - i_train = np.linspace( - 0, num_images - 1, num_train_images, dtype=int - ) # equally spaced training images starting and ending at 0 and num_images-1 - i_eval = np.setdiff1d(i_all, i_train) # eval images are the remaining images - assert len(i_eval) == num_eval_images - if split == "train": - indices = i_train - elif split in ["val", "test"]: - indices = i_eval + has_split_files_spec = any(f"{split}_filenames" in meta for split in ("train", "val", "test")) + if f"{split}_filenames" in meta: + # Validate split first + split_filenames = set(self._get_fname(PurePath(x), data_dir) for x in meta[f"{split}_filenames"]) + unmatched_filenames = split_filenames.difference(image_filenames) + if unmatched_filenames: + raise RuntimeError(f"Some filenames for split {split} were not found: {unmatched_filenames}.") + + indices = [i for i, path in enumerate(image_filenames) if path in split_filenames] + CONSOLE.log(f"[yellow] Dataset is overriding {split}_indices to {indices}") + indices = np.array(indices, dtype=np.int32) + elif has_split_files_spec: + raise RuntimeError(f"The dataset's list of filenames for split {split} is missing.") else: - raise ValueError(f"Unknown dataparser split {split}") + # filter image_filenames and poses based on train/eval split percentage + num_images = len(image_filenames) + num_train_images = math.ceil(num_images * self.config.train_split_fraction) + num_eval_images = num_images - num_train_images + i_all = np.arange(num_images) + i_train = np.linspace( + 0, num_images - 1, num_train_images, dtype=int + ) # equally spaced training images starting and ending at 0 and num_images-1 + i_eval = np.setdiff1d(i_all, i_train) # eval images are the remaining images + assert len(i_eval) == num_eval_images + if split == "train": + indices = i_train + elif split in ["val", "test"]: + indices = i_eval + else: + raise ValueError(f"Unknown dataparser split {split}") if "orientation_override" in meta: orientation_method = meta["orientation_override"] @@ -274,6 +288,15 @@ class Nerfstudio(DataParser): assert self.downscale_factor is not None cameras.rescale_output_resolution(scaling_factor=1.0 / self.downscale_factor) + if "applied_transform" in meta: + applied_transform = np.array(meta["applied_transform"], dtype=np.float32) + transform_matrix = transform_matrix @ np.concatenate( + (applied_transform, np.array([0, 0, 0, 1], dtype=np.float32)), 0 + ) + if "applied_scale" in meta: + applied_scale = float(meta["applied_scale"]) + scale_factor *= applied_scale + dataparser_outputs = DataparserOutputs( image_filenames=image_filenames, cameras=cameras, diff --git a/nerfstudio/data/dataparsers/scannet_dataparser.py b/nerfstudio/data/dataparsers/scannet_dataparser.py index 80d0dd688..46ff98f8a 100644 --- a/nerfstudio/data/dataparsers/scannet_dataparser.py +++ b/nerfstudio/data/dataparsers/scannet_dataparser.py @@ -77,15 +77,14 @@ class ScanNet(DataParser): first_img = cv2.imread(img_dir_sorted[0].as_posix()) h, w, _ = first_img.shape - image_filenames = [] - depth_filenames = [] - intrinsics = [] - poses = [] + image_filenames, depth_filenames, intrinsics, poses = [], [], [], [] K = np.loadtxt(self.config.data / "intrinsic" / "intrinsic_color.txt") for img, depth, pose in zip(img_dir_sorted, depth_dir_sorted, pose_dir_sorted): pose = np.loadtxt(pose) pose[0:3, 1:3] *= -1 + pose = pose[np.array([1, 0, 2, 3]), :] + pose[2, :] *= -1 # We cannot accept files directly, as some of the poses are invalid if np.isinf(pose).any(): diff --git a/nerfstudio/data/utils/dataloaders.py b/nerfstudio/data/utils/dataloaders.py index c4e155102..23d712d7e 100644 --- a/nerfstudio/data/utils/dataloaders.py +++ b/nerfstudio/data/utils/dataloaders.py @@ -116,7 +116,7 @@ class CacheDataloader(DataLoader): """Returns a collated batch.""" batch_list = self._get_batch_list() collated_batch = self.collate_fn(batch_list) - collated_batch = get_dict_to_torch(collated_batch, device=self.device) + collated_batch = get_dict_to_torch(collated_batch, device=self.device, exclude=["image"]) return collated_batch def __iter__(self): @@ -183,7 +183,7 @@ class EvalDataloader(DataLoader): """ ray_bundle = self.cameras.generate_rays(camera_indices=image_idx, keep_shape=True) batch = self.input_dataset[image_idx] - batch = get_dict_to_torch(batch, device=self.device) + batch = get_dict_to_torch(batch, device=self.device, exclude=["image"]) return ray_bundle, batch diff --git a/nerfstudio/engine/trainer.py b/nerfstudio/engine/trainer.py index a4a09e7b1..97ff9244a 100644 --- a/nerfstudio/engine/trainer.py +++ b/nerfstudio/engine/trainer.py @@ -243,7 +243,9 @@ class Trainer: writer.put_dict(name="Train Loss Dict", scalar_dict=loss_dict, step=step) writer.put_dict(name="Train Metrics Dict", scalar_dict=metrics_dict, step=step) - self.eval_iteration(step) + # Do not perform evaluation if there are no validation images + if self.pipeline.datamanager.eval_dataset: + self.eval_iteration(step) if step_check(step, self.config.steps_per_save): self.save_checkpoint(step) diff --git a/nerfstudio/exporter/exporter_utils.py b/nerfstudio/exporter/exporter_utils.py index 9f4168517..2f45badab 100644 --- a/nerfstudio/exporter/exporter_utils.py +++ b/nerfstudio/exporter/exporter_utils.py @@ -178,8 +178,8 @@ def generate_point_cloud( rgbs = torch.cat(rgbs, dim=0) pcd = o3d.geometry.PointCloud() - pcd.points = o3d.utility.Vector3dVector(points.float().cpu().numpy()) - pcd.colors = o3d.utility.Vector3dVector(rgbs.float().cpu().numpy()) + pcd.points = o3d.utility.Vector3dVector(points.double().cpu().numpy()) + pcd.colors = o3d.utility.Vector3dVector(rgbs.double().cpu().numpy()) ind = None if remove_outliers: @@ -203,7 +203,7 @@ def generate_point_cloud( if ind is not None: # mask out normals for points that were removed with remove_outliers normals = normals[ind] - pcd.normals = o3d.utility.Vector3dVector(normals.float().cpu().numpy()) + pcd.normals = o3d.utility.Vector3dVector(normals.double().cpu().numpy()) return pcd diff --git a/nerfstudio/fields/instant_ngp_field.py b/nerfstudio/fields/instant_ngp_field.py index 102912ba7..b0f4e97e4 100644 --- a/nerfstudio/fields/instant_ngp_field.py +++ b/nerfstudio/fields/instant_ngp_field.py @@ -19,6 +19,7 @@ Instant-NGP field implementations using tiny-cuda-nn, torch, .... from typing import Dict, Optional, Tuple +import numpy as np import torch from nerfacc import ContractionType, contract from torch.nn.parameter import Parameter @@ -63,6 +64,7 @@ class TCNNInstantNGPField(Field): contraction_type: type of contraction num_levels: number of levels of the hashmap for the base mlp log2_hashmap_size: size of the hashmap for the base mlp + max_res: maximum resolution of the hashmap for the base mlp """ def __init__( @@ -79,6 +81,7 @@ class TCNNInstantNGPField(Field): contraction_type: ContractionType = ContractionType.UN_BOUNDED_SPHERE, num_levels: int = 16, log2_hashmap_size: int = 19, + max_res: int = 2048, ) -> None: super().__init__() @@ -93,7 +96,8 @@ class TCNNInstantNGPField(Field): self.appearance_embedding = Embedding(num_images, appearance_embedding_dim) # TODO: set this properly based on the aabb - per_level_scale = 1.4472692012786865 + base_res: int = 16 + per_level_scale = np.exp((np.log(max_res) - np.log(base_res)) / (num_levels - 1)) self.direction_encoding = tcnn.Encoding( n_input_dims=3, @@ -111,7 +115,7 @@ class TCNNInstantNGPField(Field): "n_levels": num_levels, "n_features_per_level": 2, "log2_hashmap_size": log2_hashmap_size, - "base_resolution": 16, + "base_resolution": base_res, "per_level_scale": per_level_scale, }, network_config={ diff --git a/nerfstudio/models/instant_ngp.py b/nerfstudio/models/instant_ngp.py index ace4b77c6..2bdf85b6a 100644 --- a/nerfstudio/models/instant_ngp.py +++ b/nerfstudio/models/instant_ngp.py @@ -65,6 +65,10 @@ class InstantNGPModelConfig(ModelConfig): """Number of samples in field evaluation.""" grid_resolution: int = 128 """Resolution of the grid used for the field.""" + max_res: int = 2048 + """Maximum resolution of the hashmap for the base mlp.""" + log2_hashmap_size: int = 19 + """Size of the hashmap for the base mlp""" contraction_type: ContractionType = ContractionType.UN_BOUNDED_SPHERE """Contraction type used for spatial deformation of the field.""" cone_angle: float = 0.004 @@ -103,6 +107,8 @@ class NGPModel(Model): contraction_type=self.config.contraction_type, use_appearance_embedding=self.config.use_appearance_embedding, num_images=self.num_train_data, + log2_hashmap_size=self.config.log2_hashmap_size, + max_res=self.config.max_res, ) self.scene_aabb = Parameter(self.scene_box.aabb.flatten(), requires_grad=False) diff --git a/nerfstudio/process_data/colmap_utils.py b/nerfstudio/process_data/colmap_utils.py index f8254f0a8..f34d6cd84 100644 --- a/nerfstudio/process_data/colmap_utils.py +++ b/nerfstudio/process_data/colmap_utils.py @@ -455,6 +455,11 @@ 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, 3]), :] + applied_transform[2, :] *= -1 + out["applied_transform"] = applied_transform + with open(output_dir / "transforms.json", "w", encoding="utf-8") as f: json.dump(out, f, indent=4) diff --git a/nerfstudio/process_data/equirect_utils.py b/nerfstudio/process_data/equirect_utils.py index 596af58fb..340dd7884 100644 --- a/nerfstudio/process_data/equirect_utils.py +++ b/nerfstudio/process_data/equirect_utils.py @@ -15,13 +15,15 @@ """Helper utils for processing equirectangular data.""" import os +import sys from pathlib import Path -from typing import Tuple +from typing import List, Tuple import cv2 import numpy as np import torch from equilib import Equi2Pers +from rich.console import Console from rich.progress import ( BarColumn, Progress, @@ -32,11 +34,86 @@ from rich.progress import ( from nerfstudio.utils.rich_utils import ItersPerSecColumn +CONSOLE = Console(width=120) + + +def _crop_bottom(bound_arr: list, fov: int, crop_factor: float) -> List[float]: + """Returns a list of vertical bounds with the bottom cropped. + + Args: + bound_arr (list): List of vertical bounds in ascending order. + fov (int): Field of view of the camera. + crop_factor (float): Portion of the image to crop from the bottom. + + Returns: + list: A new list of bounds with the bottom cropped. + """ + degrees_chopped = 180 * crop_factor + new_bottom_start = 90 - degrees_chopped - fov / 2 + for i, el in reversed(list(enumerate(bound_arr))): + if el > new_bottom_start + fov / 2: + bound_arr[i] = None + elif el > new_bottom_start: + diff = el - new_bottom_start + bound_arr[i] = new_bottom_start + for j in range(i - 1, -1, -1): + bound_arr[j] -= diff / (2 ** (i - j)) + break + + return bound_arr + + +def _crop_top(bound_arr: list, fov: int, crop_factor: float) -> List[float]: + """Returns a list of vertical bounds with the top cropped. + + Args: + bound_arr (list): List of vertical bounds in ascending order. + fov (int): Field of view of the camera. + crop_factor (float): Portion of the image to crop from the top. + + Returns: + list: A new list of bounds with the top cropped. + """ + degrees_chopped = 180 * crop_factor + new_top_start = -90 + degrees_chopped + fov / 2 + for i, el in enumerate(bound_arr): + if el < new_top_start - fov / 2: + bound_arr[i] = None + elif el < new_top_start: + diff = new_top_start - el + bound_arr[i] = new_top_start + for j in range(i + 1, len(bound_arr)): + bound_arr[j] += diff / (2 ** (j - i)) + break + + return bound_arr + + +def _crop_bound_arr_vertical( + bound_arr: list, fov: int, crop_factor: Tuple[float, float, float, float] = (0.0, 0.0, 0.0, 0.0) +) -> list: + """Returns a list of vertical bounds adjusted for cropping. + + Args: + bound_arr (list): Original list of vertical bounds in ascending order. + fov (int): Field of view of the camera. + crop_factor (Tuple[float, float, float, float]): Crop arr (top, bottom, left, right). + + Returns: + list: Cropped bound arr + """ + if crop_factor[1] > 0: + bound_arr = _crop_bottom(bound_arr, fov, crop_factor[1]) + if crop_factor[0] > 0: + bound_arr = _crop_top(bound_arr, fov, crop_factor[0]) + return bound_arr + def generate_planar_projections_from_equirectangular( image_dir: Path, planar_image_size: Tuple[int, int], samples_per_im: int, + crop_factor: Tuple[float, float, float, float] = (0.0, 0.0, 0.0, 0.0), ) -> Path: """Generate planar projections from an equirectangular image. @@ -44,30 +121,53 @@ def generate_planar_projections_from_equirectangular( image_dir: The directory containing the equirectangular image. planar_image_size: The size of the planar projections [width, height]. samples_per_im: The number of samples to take per image. + crop_factor: The portion of the image to crop from the (top, bottom, left, and right). + Values should be in [0, 1]. returns: The path to the planar projections directory. """ + for i in crop_factor: + if i < 0 or i > 1: + CONSOLE.print("[bold red] Invalid crop factor. All values must be in [0,1].") + sys.exit(1) + device = torch.device("cuda") fov = 120 yaw_pitch_pairs = [] + left_bound, right_bound = -180, 180 + if crop_factor[3] > 0: + left_bound = -180 + 360 * crop_factor[3] + if crop_factor[2] > 0: + right_bound = 180 - 360 * crop_factor[2] + if samples_per_im == 8: fov = 120 - for i in np.arange(-180, 180, 90): - yaw_pitch_pairs.append((i, 0)) - for i in np.arange(-180, 180, 180): - yaw_pitch_pairs.append((i, 45)) - for i in np.arange(-180, 180, 180): - yaw_pitch_pairs.append((i, -45)) + bound_arr = [-45, 0, 45] + bound_arr = _crop_bound_arr_vertical(bound_arr, fov, crop_factor) + if bound_arr[1] is not None: + for i in np.arange(left_bound, right_bound, 90): + yaw_pitch_pairs.append((i, bound_arr[1])) + if bound_arr[2] is not None: + for i in np.arange(left_bound, right_bound, 180): + yaw_pitch_pairs.append((i, bound_arr[2])) + if bound_arr[0] is not None: + for i in np.arange(left_bound, right_bound, 180): + yaw_pitch_pairs.append((i, bound_arr[0])) elif samples_per_im == 14: fov = 110 - for i in np.arange(-180, 180, 60): - yaw_pitch_pairs.append((i, 0)) - for i in np.arange(-180, 180, 90): - yaw_pitch_pairs.append((i, 45)) - for i in np.arange(-180, 180, 90): - yaw_pitch_pairs.append((i, -45)) + bound_arr = [-45, 0, 45] + bound_arr = _crop_bound_arr_vertical(bound_arr, fov, crop_factor) + if bound_arr[1] is not None: + for i in np.arange(left_bound, right_bound, 60): + yaw_pitch_pairs.append((i, bound_arr[1])) + if bound_arr[2] is not None: + for i in np.arange(left_bound, right_bound, 90): + yaw_pitch_pairs.append((i, bound_arr[2])) + if bound_arr[0] is not None: + for i in np.arange(left_bound, right_bound, 90): + yaw_pitch_pairs.append((i, bound_arr[0])) equi2pers = Equi2Pers(height=planar_image_size[1], width=planar_image_size[0], fov_x=fov, mode="bilinear") frame_dir = image_dir @@ -82,7 +182,6 @@ def generate_planar_projections_from_equirectangular( TimeRemainingColumn(elapsed_when_finished=True, compact=True), ) - curr_im = 0 with progress: for i in progress.track(os.listdir(frame_dir), description="", total=num_ims): if i.lower().endswith((".jpg", ".png", ".jpeg")): @@ -97,7 +196,6 @@ def generate_planar_projections_from_equirectangular( pers_image = (pers_image.permute(1, 2, 0)).type(torch.uint8).to("cpu").numpy() cv2.imwrite(f"{output_dir}/{i[:-4]}_{count}.jpg", pers_image) count += 1 - curr_im += 1 return output_dir diff --git a/nerfstudio/process_data/process_data_utils.py b/nerfstudio/process_data/process_data_utils.py index 6b94bbaa9..14ae33662 100644 --- a/nerfstudio/process_data/process_data_utils.py +++ b/nerfstudio/process_data/process_data_utils.py @@ -100,7 +100,11 @@ def get_num_frames_in_video(video: Path) -> int: def convert_video_to_images( - video_path: Path, image_dir: Path, num_frames_target: int, verbose: bool = False + video_path: Path, + image_dir: Path, + num_frames_target: int, + crop_factor: Tuple[float, float, float, float] = (0.0, 0.0, 0.0, 0.0), + verbose: bool = False, ) -> Tuple[List[str], int]: """Converts a video into a sequence of images. @@ -108,11 +112,17 @@ def convert_video_to_images( video_path: Path to the video. output_dir: Path to the output directory. num_frames_target: Number of frames to extract. + crop_factor: Portion of the image to crop. Should be in [0,1] (top, bottom, left, right) verbose: If True, logs the output of the command. Returns: A tuple containing summary of the conversion and the number of extracted frames. """ + for i in crop_factor: + if i < 0 or i > 1: + CONSOLE.print("[bold red]Error: Invalid crop factor. All crops must be in [0,1].") + sys.exit(1) + if video_path.is_dir(): CONSOLE.print(f"[bold red]Error: Video path is a directory, not a path: {video_path}") sys.exit(1) @@ -135,16 +145,23 @@ def convert_video_to_images( out_filename = image_dir / "frame_%05d.png" ffmpeg_cmd = f'ffmpeg -i "{video_path}"' - spacing = num_frames // num_frames_target + crop_cmd = "" + if crop_factor != (0.0, 0.0, 0.0, 0.0): + height = 1 - crop_factor[0] - crop_factor[1] + width = 1 - crop_factor[2] - crop_factor[3] + start_x = crop_factor[2] + start_y = crop_factor[0] + crop_cmd = f',"crop=w=iw*{width}:h=ih*{height}:x=iw*{start_x}:y=ih*{start_y}"' + + spacing = num_frames // num_frames_target if spacing > 1: - ffmpeg_cmd += f" -vf thumbnail={spacing},setpts=N/TB -r 1" + ffmpeg_cmd += f" -vf thumbnail={spacing},setpts=N/TB{crop_cmd} -r 1" else: CONSOLE.print("[bold red]Can't satisfy requested number of frames. Extracting all frames.") ffmpeg_cmd += " -pix_fmt bgr8" ffmpeg_cmd += f" {out_filename}" - run_command(ffmpeg_cmd, verbose=verbose) num_final_frames = len(list(image_dir.glob("*.png"))) @@ -399,25 +416,23 @@ def generate_circle_mask(height: int, width: int, percent_radius) -> Optional[np return mask -def generate_crop_mask( - height: int, width: int, percent_crop: Tuple[float, float, float, float] -) -> Optional[np.ndarray]: +def generate_crop_mask(height: int, width: int, crop_factor: Tuple[float, float, float, float]) -> Optional[np.ndarray]: """generate a crop mask of the given size. Args: height: The height of the mask. width: The width of the mask. - percent_crop: The percent of the image to crop in each direction [top, bottom, left, right]. + crop_factor: The percent of the image to crop in each direction [top, bottom, left, right]. Returns: The mask or None if no cropping is performed. """ - if np.all(np.array(percent_crop) == 0.0): + if np.all(np.array(crop_factor) == 0.0): return None - if np.any(np.array(percent_crop) < 0.0) or np.any(np.array(percent_crop) > 1.0): + if np.any(np.array(crop_factor) < 0.0) or np.any(np.array(crop_factor) > 1.0): CONSOLE.log("[bold red]Invalid crop percentage, must be between 0 and 1.") sys.exit(1) - top, bottom, left, right = percent_crop + top, bottom, left, right = crop_factor mask = np.zeros((height, width), dtype=np.uint8) top = int(top * height) bottom = int(bottom * height) @@ -428,20 +443,20 @@ def generate_crop_mask( def generate_mask( - height: int, width: int, percent_crop: Tuple[float, float, float, float], percent_radius: float + height: int, width: int, crop_factor: Tuple[float, float, float, float], percent_radius: float ) -> Optional[np.ndarray]: """generate a mask of the given size. Args: height: The height of the mask. width: The width of the mask. - percent_crop: The percent of the image to crop in each direction [top, bottom, left, right]. + crop_factor: The percent of the image to crop in each direction [top, bottom, left, right]. percent_radius: The radius of the circle as a percentage of the image diagonal size. Returns: The mask or None if no mask is needed. """ - crop_mask = generate_crop_mask(height, width, percent_crop) + crop_mask = generate_crop_mask(height, width, crop_factor) circle_mask = generate_circle_mask(height, width, percent_radius) if crop_mask is None: return circle_mask @@ -453,7 +468,7 @@ def generate_mask( def save_mask( image_dir: Path, num_downscales: int, - percent_crop: Tuple[float, float, float, float] = (0, 0, 0, 0), + crop_factor: Tuple[float, float, float, float] = (0, 0, 0, 0), percent_radius: float = 1.0, ) -> Optional[Path]: """Save a mask for each image in the image directory. @@ -461,7 +476,7 @@ def save_mask( Args: image_dir: The directory containing the images. num_downscales: The number of downscaling levels. - percent_crop: The percent of the image to crop in each direction [top, bottom, left, right]. + crop_factor: The percent of the image to crop in each direction [top, bottom, left, right]. percent_radius: The radius of the circle as a percentage of the image diagonal size. Returns: @@ -470,7 +485,7 @@ def save_mask( image_path = next(image_dir.glob("frame_*")) image = cv2.imread(str(image_path)) height, width = image.shape[:2] - mask = generate_mask(height, width, percent_crop, percent_radius) + mask = generate_mask(height, width, crop_factor, percent_radius) if mask is None: return None mask *= 255 diff --git a/nerfstudio/viewer/app/src/reducer.js b/nerfstudio/viewer/app/src/reducer.js index d8335c80d..76a21fbb8 100644 --- a/nerfstudio/viewer/app/src/reducer.js +++ b/nerfstudio/viewer/app/src/reducer.js @@ -71,9 +71,6 @@ const initialState = { }; function setData(newState, state, path, data) { - if (path === 'colormap_options') { - newState.colormap_choice = 'default'; // eslint-disable-line no-param-reassign - } if (path.length === 1) { newState[path[0]] = data; // eslint-disable-line no-param-reassign } else { diff --git a/nerfstudio/viewer/server/viewer_utils.py b/nerfstudio/viewer/server/viewer_utils.py index 35b63e9f6..e34e53e98 100644 --- a/nerfstudio/viewer/server/viewer_utils.py +++ b/nerfstudio/viewer/server/viewer_utils.py @@ -21,7 +21,6 @@ from __future__ import annotations import asyncio import enum import os -import re import sys import threading import time @@ -351,10 +350,8 @@ class ViewerState: # set the data base dir self.vis["renderingState/data_base_dir"].write(str(self.datapath)) - # get the timestamp of the train run to set default export path name - timestamp_reg = re.compile("[0-9]{4}-[0-9]{2}-[0-9]{2}_[0-9]{6}") - timestamp_match = timestamp_reg.findall(str(self.log_filename.parents[0])) - self.vis["renderingState/export_path"].write(timestamp_match[-1]) + # set default export path name + self.vis["renderingState/export_path"].write(self.log_filename.parent.stem) # clear the current scene self.vis["sceneState/sceneBox"].delete() @@ -768,10 +765,12 @@ class ViewerState: if self.output_type_changed or self.prev_colormap_type is None: self.prev_colormap_type = ColormapTypes.DEFAULT colormap_options = [] + self.vis["renderingState/colormap_options"].write(list(ColormapTypes)) if outputs[reformatted_output].shape[-1] == 3: colormap_options = [ColormapTypes.DEFAULT] if outputs[reformatted_output].shape[-1] == 1 and outputs[reformatted_output].dtype == torch.float: - colormap_options = list(ColormapTypes) + self.prev_colormap_type = ColormapTypes.TURBO + colormap_options = list(ColormapTypes)[1:] self.output_type_changed = False self.vis["renderingState/colormap_choice"].write(self.prev_colormap_type) self.vis["renderingState/colormap_options"].write(colormap_options) diff --git a/scripts/process_data.py b/scripts/process_data.py index 7c39afd1f..570fda59e 100644 --- a/scripts/process_data.py +++ b/scripts/process_data.py @@ -292,8 +292,8 @@ class ProcessVideo: """ percent_radius_crop: float = 1.0 """Create circle crop mask. The radius is the percent of the image diagonal.""" - percent_crop: Tuple[float, float, float, float] = (0.0, 0.0, 0.0, 0.0) - """Percent of the image to crop. (top, bottom, left, right)""" + crop_factor: Tuple[float, float, float, float] = (0.0, 0.0, 0.0, 0.0) + """Portion of the image to crop. All values should be in [0,1]. (top, bottom, left, right)""" use_sfm_depth: bool = False """If True, export and use depth maps induced from SfM points.""" include_depth_debug: bool = False @@ -319,11 +319,19 @@ class ProcessVideo: temp_image_dir = self.output_dir / "temp_images" temp_image_dir.mkdir(parents=True, exist_ok=True) summary_log, num_extracted_frames = process_data_utils.convert_video_to_images( - self.data, image_dir=temp_image_dir, num_frames_target=self.num_frames_target, verbose=self.verbose + self.data, + image_dir=temp_image_dir, + num_frames_target=self.num_frames_target, + crop_factor=(0.0, 0.0, 0.0, 0.0), + verbose=self.verbose, ) else: summary_log, num_extracted_frames = process_data_utils.convert_video_to_images( - self.data, image_dir=image_dir, num_frames_target=self.num_frames_target, verbose=self.verbose + self.data, + image_dir=image_dir, + num_frames_target=self.num_frames_target, + crop_factor=self.crop_factor, + verbose=self.verbose, ) # Generate planar projections if equirectangular @@ -332,7 +340,10 @@ class ProcessVideo: self.output_dir / "temp_images", self.images_per_equirect ) image_dir = equirect_utils.generate_planar_projections_from_equirectangular( - self.output_dir / "temp_images", perspective_image_size, self.images_per_equirect + self.output_dir / "temp_images", + perspective_image_size, + self.images_per_equirect, + crop_factor=self.crop_factor, ) # copy the perspective images to the image directory @@ -346,11 +357,11 @@ class ProcessVideo: # remove the temp_images folder shutil.rmtree(self.output_dir / "temp_images", ignore_errors=True) - # Create mask + # # Create mask mask_path = process_data_utils.save_mask( image_dir=image_dir, num_downscales=self.num_downscales, - percent_crop=self.percent_crop, + crop_factor=(0.0, 0.0, 0.0, 0.0), percent_radius=self.percent_radius_crop, ) if mask_path is not None: diff --git a/tests/dataparsers/test_nerfstudio_dataparser.py b/tests/dataparsers/test_nerfstudio_dataparser.py new file mode 100644 index 000000000..c5dfd9518 --- /dev/null +++ b/tests/dataparsers/test_nerfstudio_dataparser.py @@ -0,0 +1,99 @@ +""" +Nerfstudio dataparser +""" +# pylint: disable=all +import json +from pathlib import Path + +import numpy as np +import pytest +from PIL import Image +from pytest import fixture + + +@fixture +def mocked_dataset(tmp_path: Path): + """Mocked dataset with transforms""" + (tmp_path / "images_4").mkdir() + frames = [] + for i in range(10): + Image.new("RGB", (100, 150)).save(tmp_path / "images_4" / f"img_{i}.png") + frames.append( + { + "file_path": f"img_{i}.png", + "transform_matrix": np.eye(4).tolist(), + } + ) + with (tmp_path / "transforms.json").open("w+", encoding="utf8") as f: + json.dump({"fl_x": 2, "fl_y": 3, "cx": 4, "cy": 5, "h": 150, "w": 100, "frames": frames}, f) + return tmp_path + + +@pytest.mark.parametrize("orientation_method", ["up", "none", "pca"]) +def test_nerfstudio_dataparser_no_filelist(mocked_dataset, orientation_method): + """Tests basic load""" + assert (mocked_dataset / "images_4").exists() + from nerfstudio.data.dataparsers.nerfstudio_dataparser import ( + DataparserOutputs, + Nerfstudio, + NerfstudioDataParserConfig, + ) + + parser: Nerfstudio = NerfstudioDataParserConfig( + data=mocked_dataset, + downscale_factor=4, + orientation_method=orientation_method, + center_poses=False, + auto_scale_poses=False, + ).setup() + + paths = set() + for split in ("train", "test", "val"): + out = parser.get_dataparser_outputs(split) + assert isinstance(out, DataparserOutputs) + assert len(out.image_filenames) > 0 + paths.update(out.image_filenames) + train_files = set(parser.get_dataparser_outputs("train").image_filenames) + assert len(train_files.intersection(parser.get_dataparser_outputs("val").image_filenames)) == 0 + assert len(train_files.intersection(parser.get_dataparser_outputs("test").image_filenames)) == 0 + assert len(paths) == 10 + + +def test_nerfstudio_dataparser_split_filelist(mocked_dataset): + """Tests basic load""" + assert (mocked_dataset / "images_4").exists() + with open(mocked_dataset / "transforms.json", "r+") as f: + data = json.load(f) + data["train_filenames"] = ["img_0.png", "img_1.png"] + data["val_filenames"] = ["img_2.png", "img_3.png"] + data["test_filenames"] = ["img_4.png", "img_5.png"] + f.seek(0) + f.truncate(0) + json.dump(data, f) + + from nerfstudio.data.dataparsers.nerfstudio_dataparser import ( + DataparserOutputs, + Nerfstudio, + NerfstudioDataParserConfig, + ) + + parser: Nerfstudio = NerfstudioDataParserConfig( + data=mocked_dataset, + downscale_factor=4, + orientation_method="none", + center_poses=False, + auto_scale_poses=False, + ).setup() + + assert parser.get_dataparser_outputs("train").image_filenames == [ + mocked_dataset / "images_4/img_0.png", + mocked_dataset / "images_4/img_1.png", + ] + assert parser.get_dataparser_outputs("val").image_filenames == [ + mocked_dataset / "images_4/img_2.png", + mocked_dataset / "images_4/img_3.png", + ] + assert parser.get_dataparser_outputs("test").image_filenames == [ + mocked_dataset / "images_4/img_4.png", + mocked_dataset / "images_4/img_5.png", + ]