diff --git a/nerfstudio/pipelines/base_pipeline.py b/nerfstudio/pipelines/base_pipeline.py index e9d30a9ab..363c8c9ea 100644 --- a/nerfstudio/pipelines/base_pipeline.py +++ b/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" diff --git a/nerfstudio/scripts/eval.py b/nerfstudio/scripts/eval.py index ef47d7a24..8f5995117 100644 --- a/nerfstudio/scripts/eval.py +++ b/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 = {