Add option to render evaluated images (#2135)

* Add option to render evaluated images

* add args

* fix formatting

* attempt to fix pyright

* attempt 2

---------

Co-authored-by: Matthew Tancik <matt@lumalabs.ai>

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