Merge branch 'main' into justin/opensfm

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

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:

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)
 

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],

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]]

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:

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(

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}")
 

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 = [