diff --git a/nerfstudio/configs/base_config.py b/nerfstudio/configs/base_config.py index 3f1aff36a..340806577 100644 --- a/nerfstudio/configs/base_config.py +++ b/nerfstudio/configs/base_config.py @@ -175,6 +175,9 @@ class ViewerConfig(PrintableConfig): """the ip address where the bridge server is running""" num_rays_per_chunk: int = 32768 """number of rays per chunk to render with viewer""" + max_num_display_images: int = 512 + """Maximum number of training images to display in the viewer, to avoid lag. This does not change which images are + actually used in training/evaluation. If -1, display all.""" from nerfstudio.engine.optimizers import OptimizerConfig diff --git a/nerfstudio/data/dataparsers/nerfstudio_dataparser.py b/nerfstudio/data/dataparsers/nerfstudio_dataparser.py index 66cfbd967..afb66c657 100644 --- a/nerfstudio/data/dataparsers/nerfstudio_dataparser.py +++ b/nerfstudio/data/dataparsers/nerfstudio_dataparser.py @@ -17,7 +17,7 @@ from __future__ import annotations import math from dataclasses import dataclass, field -from pathlib import Path, PureWindowsPath +from pathlib import Path, PurePath from typing import Optional, Type import numpy as np @@ -80,12 +80,9 @@ class Nerfstudio(DataParser): num_skipped_image_filenames = 0 for frame in meta["frames"]: - if "\\" in frame["file_path"]: - filepath = PureWindowsPath(frame["file_path"]) - else: - filepath = Path(frame["file_path"]) + filepath = PurePath(frame["file_path"]) fname = self._get_fname(filepath) - if not fname: + if not fname.exists(): num_skipped_image_filenames += 1 else: image_filenames.append(fname) @@ -177,7 +174,7 @@ class Nerfstudio(DataParser): ) return dataparser_outputs - def _get_fname(self, filepath): + def _get_fname(self, filepath: PurePath) -> Path: """Get the filename of the image file.""" if self.downscale_factor is None: diff --git a/nerfstudio/model_components/renderers.py b/nerfstudio/model_components/renderers.py index 025194ee4..9604a1d5f 100644 --- a/nerfstudio/model_components/renderers.py +++ b/nerfstudio/model_components/renderers.py @@ -234,6 +234,8 @@ class DepthRenderer(nn.Module): if ray_indices is not None and num_rays is not None: # Necessary for packed samples from volumetric ray sampler depth = nerfacc.accumulate_along_rays(weights, ray_indices, steps, num_rays) + accumulation = nerfacc.accumulate_along_rays(weights, ray_indices, None, num_rays) + depth = depth / (accumulation + eps) else: depth = torch.sum(weights * steps, dim=-2) / (torch.sum(weights, -2) + eps) diff --git a/nerfstudio/models/semantic_nerfw.py b/nerfstudio/models/semantic_nerfw.py index 8dd866935..c30804439 100644 --- a/nerfstudio/models/semantic_nerfw.py +++ b/nerfstudio/models/semantic_nerfw.py @@ -102,9 +102,9 @@ class SemanticNerfWModel(Model): if self.config.use_same_proposal_network: network = HashMLPDensityField(self.scene_box.aabb, spatial_distortion=scene_contraction) self.proposal_networks.append(network) - self.density_fns = [network.density_fn for _ in range(self.config.num_proposal_network_iterations)] + self.density_fns = [network.density_fn for _ in range(self.config.num_proposal_iterations)] else: - for _ in range(self.config.num_proposal_network_iterations): + for _ in range(self.config.num_proposal_iterations): network = HashMLPDensityField(self.scene_box.aabb, spatial_distortion=scene_contraction) self.proposal_networks.append(network) self.density_fns = [network.density_fn for network in self.proposal_networks] @@ -116,7 +116,7 @@ class SemanticNerfWModel(Model): self.proposal_sampler = ProposalNetworkSampler( num_nerf_samples_per_ray=self.config.num_nerf_samples_per_ray, num_proposal_samples_per_ray=self.config.num_proposal_samples_per_ray, - num_proposal_network_iterations=self.config.num_proposal_network_iterations, + num_proposal_network_iterations=self.config.num_proposal_iterations, single_jitter=self.config.use_single_jitter, ) @@ -193,7 +193,7 @@ class SemanticNerfWModel(Model): outputs["weights_list"] = weights_list outputs["ray_samples_list"] = ray_samples_list - for i in range(self.config.num_proposal_network_iterations): + for i in range(self.config.num_proposal_iterations): outputs[f"prop_depth_{i}"] = self.renderer_depth(weights=weights_list[i], ray_samples=ray_samples_list[i]) # transients @@ -286,7 +286,7 @@ class SemanticNerfWModel(Model): images_dict = {"img": combined_rgb, "accumulation": combined_acc, "depth": combined_depth} - for i in range(self.config.num_proposal_network_iterations): + for i in range(self.config.num_proposal_iterations): key = f"prop_depth_{i}" prop_depth_i = colormaps.apply_depth_colormap( outputs[key], diff --git a/nerfstudio/viewer/server/viewer_utils.py b/nerfstudio/viewer/server/viewer_utils.py index 70848d08d..4824262b7 100644 --- a/nerfstudio/viewer/server/viewer_utils.py +++ b/nerfstudio/viewer/server/viewer_utils.py @@ -274,6 +274,22 @@ class ViewerState: self.webrtc_thread = None self.kill_webrtc_signal = False + def _pick_drawn_image_idxs(self, total_num: int) -> list[int]: + """Determine indicies of images to display in viewer. + + Args: + total_num: total number of training images. + + Returns: + List of indices from [0, total_num-1]. + """ + if self.config.max_num_display_images < 0: + num_display_images = total_num + else: + num_display_images = min(self.config.max_num_display_images, total_num) + # draw indices, roughly evenly spaced + return np.linspace(0, total_num - 1, num_display_images, dtype=np.int32).tolist() + def init_scene(self, dataset: InputDataset, start_train=True) -> None: """Draw some images and the scene aabb in the viewer. @@ -290,7 +306,7 @@ class ViewerState: self.vis["sceneState/cameras"].delete() # draw the training cameras and images - image_indices = range(len(dataset)) + image_indices = self._pick_drawn_image_idxs(len(dataset)) for idx in image_indices: image = dataset[idx]["image"] bgr = image[..., [2, 1, 0]] diff --git a/scripts/completions/install.py b/scripts/completions/install.py index 37cfd9fbc..adfb76e7c 100755 --- a/scripts/completions/install.py +++ b/scripts/completions/install.py @@ -133,46 +133,52 @@ def _update_rc( """ # Install or uninstall `source_line`. - source_lines = "\n".join( - [ - "", - "# Source nerfstudio autocompletions.", - "# From https://unix.stackexchange.com/questions/339954/zsh-command-not-found-compinstall-compinit-compdef", - "autoload -Uz compinit", - "compinit", - f"source {completions_dir / 'setup'}.{shell}", - ] - ) + header_line = "# Source nerfstudio autocompletions." + if shell == "zsh": + source_lines = "\n".join( + [ + "", + header_line, + "if ! command -v compdef &> /dev/null; then", + " autoload -Uz compinit", + " compinit", + "fi", + f"source {completions_dir / 'setup.zsh'}", + ] + ) + elif shell == "bash": + source_lines = "\n".join( + [ + "", + header_line, + f"source {completions_dir / 'setup.bash'}", + ] + ) + else: + assert_never(shell) + rc_path = pathlib.Path(os.environ["HOME"]) / f".{shell}rc" + + # Always try to uninstall previous completions. + rc_source = rc_path.read_text() + while header_line in rc_source: + before_install, _, after_install = rc_source.partition(header_line) + source_file, _, after_install = after_install.partition("\nsource ")[2].partition("\n") + assert source_file.endswith(f"/completions/setup.{shell}") + rc_source = before_install + after_install + rc_path.write_text(rc_source) + CONSOLE.log(f":broom: Existing completions uninstalled from {rc_path}.") + + # Install completions. if mode == "install": - if source_lines in rc_path.read_text(): - CONSOLE.log(f":call_me_hand: Completions are already installed in {rc_path}. {_exclamation()}") - return - - if not Confirm.ask(f"[bold yellow]Install to {rc_path}?", default=True): - CONSOLE.log(f"[bold red]Skipping install for {rc_path.name}.") - return - - rc_path.write_text(rc_path.read_text() + source_lines) + assert source_lines not in rc_source + rc_path.write_text(rc_source.rstrip() + "\n" + source_lines) CONSOLE.log( f":person_gesturing_ok: Completions installed to {rc_path}. {_exclamation()} Open a new shell to try them" " out." ) - - elif mode == "uninstall": - if source_lines not in rc_path.read_text(): - CONSOLE.log(f":heavy_check_mark: No completions to uninstall from {rc_path.name}.") - return - - if not Confirm.ask(f"[bold yellow]Uninstall from {rc_path}?", default=True): - CONSOLE.log(f"[dim red]Skipping uninstall for {rc_path.name}.") - return - - rc_path.write_text(rc_path.read_text().replace(source_lines, "")) - CONSOLE.log(f":broom: Completions uninstalled from {rc_path}.") - else: - assert_never(mode) + assert mode == "uninstall" def main(mode: ConfigureMode = "install") -> None: diff --git a/scripts/downloads/download_data.py b/scripts/downloads/download_data.py index aa96b25c1..5a8c80727 100755 --- a/scripts/downloads/download_data.py +++ b/scripts/downloads/download_data.py @@ -32,7 +32,8 @@ def download_blender(save_dir: Path): unzip_path = save_dir / Path("nerf_synthetic") final_path = save_dir / Path("blender") unzip_path.rename(final_path) - download_path.unlink(missing_ok=True) + if download_path.exists(): + download_path.unlink() def download_friends(save_dir: Path): @@ -123,7 +124,8 @@ def download_dnerf(save_dir: Path): unzip_path = save_dir / Path("data") final_path = save_dir / Path("dnerf") unzip_path.rename(final_path) - download_path.unlink(missing_ok=True) + if download_path.exists(): + download_path.unlink() def main( diff --git a/scripts/eval.py b/scripts/eval.py index b34a04373..e9ba12f78 100755 --- a/scripts/eval.py +++ b/scripts/eval.py @@ -18,7 +18,7 @@ CONSOLE = Console(width=120) @dataclass class ComputePSNR: - """Load a checkpoint, compute some PSNR metrics, and save to a JSON.""" + """Load a checkpoint, compute some PSNR metrics, and save it to a JSON file.""" # Path to config YAML file. load_config: Path @@ -28,17 +28,17 @@ class ComputePSNR: def main(self) -> None: """Main function.""" config, pipeline, checkpoint_path = eval_setup(self.load_config) - assert self.output_path.suffix == ".json" metrics_dict = pipeline.get_average_eval_image_metrics() - # save output to some file self.output_path.parent.mkdir(parents=True, exist_ok=True) + # Get the output and define the names to save to benchmark_info = { "experiment_name": config.experiment_name, "method_name": config.method_name, "checkpoint": str(checkpoint_path), "results": metrics_dict, } + # Save output to output file self.output_path.write_text(json.dumps(benchmark_info, indent=2), "utf8") CONSOLE.print(f"Saved results to: {self.output_path}") diff --git a/scripts/process_data.py b/scripts/process_data.py index 222d3635c..c0f457760 100755 --- a/scripts/process_data.py +++ b/scripts/process_data.py @@ -230,11 +230,9 @@ def convert_insta360_to_images( front_vf_cmds = vf_cmds + ["transpose=2"] back_vf_cmds = vf_cmds + ["transpose=1"] - front_ffmpeg_cmd = ( - f"ffmpeg -i {video_front} -vf '{','.join(front_vf_cmds)}' -r 1 {image_dir / 'frame_%05d.png'}" - ) + front_ffmpeg_cmd = f"ffmpeg -i {video_front} -vf {','.join(front_vf_cmds)} -r 1 {image_dir / 'frame_%05d.png'}" back_ffmpeg_cmd = ( - f"ffmpeg -i {video_back} -vf '{','.join(back_vf_cmds)}' -r 1 {image_dir / 'back_frame_%05d.png'}" + f"ffmpeg -i {video_back} -vf {','.join(back_vf_cmds)} -r 1 {image_dir / 'back_frame_%05d.png'}" ) run_command(front_ffmpeg_cmd, verbose=verbose) @@ -347,7 +345,10 @@ def run_colmap( colmap_version = get_colmap_version() - (colmap_dir / "database.db").unlink(missing_ok=True) + colmap_database_path = colmap_dir / "database.db" + if colmap_database_path.exists(): + # Can't use missing_ok argument because of Python 3.7 compatibility. + colmap_database_path.unlink() # Feature extraction feature_extractor_cmd = [