Merge branch 'main' into jkulhanek/unify-plugin-registry

nerfstudio/data/dataparsers/colmap_dataparser.py

@@ -174,7 +174,6 @@ class ColmapDataParser(DataParser):
             split_filenames = set(self.config.images_path / x for x in filenames)
             unmatched_filenames = split_filenames.difference(image_filenames)
             if unmatched_filenames:
-                breakpoint()
                 raise RuntimeError(
                     f"Some filenames for split {split} were not found: {set(map(str, unmatched_filenames))}."
                 )

nerfstudio/fields/tensorf_field.py

@@ -98,14 +98,14 @@ class TensoRFField(Field):
         rgb_features = self.color_encoding(positions)
         rgb_features = self.B(rgb_features)
 
-        d_encoded = self.direction_encoding(d)
-        rgb_features_encoded = self.feature_encoding(rgb_features)
-
         if self.use_sh:
             sh_mult = self.sh(d)[:, :, None]
             rgb_sh = rgb_features.view(sh_mult.shape[0], sh_mult.shape[1], 3, sh_mult.shape[-1])
             rgb = torch.relu(torch.sum(sh_mult * rgb_sh, dim=-1) + 0.5)
         else:
+            d_encoded = self.direction_encoding(d)
+            rgb_features_encoded = self.feature_encoding(rgb_features)
+
             out = self.mlp_head(torch.cat([rgb_features, d, rgb_features_encoded, d_encoded], dim=-1))  # type: ignore
             rgb = self.field_output_rgb(out)
 

nerfstudio/models/nerfacto.py

@@ -70,7 +70,7 @@ class NerfactoModelConfig(ModelConfig):
     num_levels: int = 16
     """Number of levels of the hashmap for the base mlp."""
     base_res: int = 16
-    """Resolution of the base grid for the hasgrid."""
+    """Resolution of the base grid for the hashgrid."""
     max_res: int = 2048
     """Maximum resolution of the hashmap for the base mlp."""
     log2_hashmap_size: int = 19

nerfstudio/pipelines/base_pipeline.py

@@ -20,11 +20,13 @@ from __future__ import annotations
 import typing
 from abc import abstractmethod
 from dataclasses import dataclass, field
+from pathlib import Path
 from time import time
 from typing import Any, Dict, List, Literal, Mapping, Optional, Tuple, Type, Union, cast
 
 import torch
 import torch.distributed as dist
+from PIL import Image
 from rich.progress import (
     BarColumn,
     MofNCompleteColumn,
@@ -173,8 +175,13 @@ class Pipeline(nn.Module):
 
     @abstractmethod
     @profiler.time_function
-    def get_average_eval_image_metrics(self, step: Optional[int] = None):
-        """Iterate over all the images in the eval dataset and get the average."""
+    def get_average_eval_image_metrics(self, step: Optional[int] = None, output_path: Optional[Path] = None):
+        """Iterate over all the images in the eval dataset and get the average.
+
+        Args:
+            step: current training step
+            output_path: optional path to save rendered images to
+        """
 
     def load_pipeline(self, loaded_state: Dict[str, Any], step: int) -> None:
         """Load the checkpoint from the given path
@@ -339,9 +346,13 @@ class VanillaPipeline(Pipeline):
         return metrics_dict, images_dict
 
     @profiler.time_function
-    def get_average_eval_image_metrics(self, step: Optional[int] = None):
+    def get_average_eval_image_metrics(self, step: Optional[int] = None, output_path: Optional[Path] = None):
         """Iterate over all the images in the eval dataset and get the average.
 
+        Args:
+            step: current training step
+            output_path: optional path to save rendered images to
+
         Returns:
             metrics_dict: dictionary of metrics
         """
@@ -363,7 +374,15 @@ class VanillaPipeline(Pipeline):
                 height, width = camera_ray_bundle.shape
                 num_rays = height * width
                 outputs = self.model.get_outputs_for_camera_ray_bundle(camera_ray_bundle)
-                metrics_dict, _ = self.model.get_image_metrics_and_images(outputs, batch)
+                metrics_dict, images_dict = self.model.get_image_metrics_and_images(outputs, batch)
+
+                if output_path is not None:
+                    camera_indices = camera_ray_bundle.camera_indices
+                    assert camera_indices is not None
+                    for key, val in images_dict.items():
+                        Image.fromarray((val * 255).byte().cpu().numpy()).save(
+                            output_path / "{0:06d}-{1}.jpg".format(int(camera_indices[0, 0, 0]), key)
+                        )
                 assert "num_rays_per_sec" not in metrics_dict
                 metrics_dict["num_rays_per_sec"] = num_rays / (time() - inner_start)
                 fps_str = "fps"

nerfstudio/scripts/eval.py

@@ -21,6 +21,7 @@ from __future__ import annotations
 import json
 from dataclasses import dataclass
 from pathlib import Path
+from typing import Optional
 
 import tyro
 
@@ -36,12 +37,16 @@ class ComputePSNR:
     load_config: Path
     # Name of the output file.
     output_path: Path = Path("output.json")
+    # Optional path to save rendered outputs to.
+    render_output_path: Optional[Path] = None
 
     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()
+        if self.render_output_path is not None:
+            self.render_output_path.mkdir(parents=True)
+        metrics_dict = pipeline.get_average_eval_image_metrics(output_path=self.render_output_path)
         self.output_path.parent.mkdir(parents=True, exist_ok=True)
         # Get the output and define the names to save to
         benchmark_info = {