""" run_eval.py """ from __future__ import annotations import dataclasses import json import logging import os from dataclasses import dataclass from pathlib import Path from typing import Dict, Literal, Tuple, Union import dcargs import mediapy as media import torch import yaml from tqdm import tqdm from typing_extensions import assert_never from nerfactory.cameras.camera_paths import ( get_interpolated_camera_path, get_spiral_path, ) from nerfactory.cameras.cameras import Cameras from nerfactory.configs import base as cfg from nerfactory.pipelines.base import Pipeline from nerfactory.utils.misc import human_format from nerfactory.utils.writer import TimeWriter logging.basicConfig(format="[%(filename)s:%(lineno)d] %(message)s", level=logging.INFO) def _update_avg(prev_avg: float, new_val: float, step: int) -> float: """helper to calculate the running average Args: prev_avg (float): previous average value new_val (float): new value to update the average with step (int): current step number Returns: float: new updated average """ return (step * prev_avg + new_val) / (step + 1) def _load_checkpoint(config: cfg.TrainerConfig, pipeline: Pipeline) -> Path: """Helper function to load checkpointed pipeline Args: config (DictConfig): Configuration of pipeline to load pipeline (Pipeline): Pipeline instance of which to load weights """ assert config.load_dir is not None if config.load_step is None: print("Loading latest checkpoint from load_dir") # NOTE: this is specific to the checkpoint name format load_step = sorted(int(x[x.find("-") + 1 : x.find(".")]) for x in os.listdir(config.load_dir))[-1] else: load_step = config.load_step load_path = config.load_dir / f"step-{load_step:09d}.ckpt" assert load_path.exists(), f"Checkpoint {load_path} does not exist" loaded_state = torch.load(load_path, map_location="cpu") pipeline.load_pipeline(loaded_state["pipeline"]) print(f"done loading checkpoint from {load_path}") return load_path def _render_stats_dict(pipeline: Pipeline) -> Dict[str, float]: """Helper function to evaluate the pipeline on a dataloader. Args: pipeline (Pipeline): Pipeline to evaluate Returns: dict: returns the average psnr and average rays per second """ avg_psnr = 0 avg_rays_per_sec = 0 avg_fps = 0 for step, (camera_ray_bundle, batch) in tqdm(enumerate(pipeline.dataloader.eval_dataloader)): with TimeWriter(writer=None, name=None, write=False) as t: with torch.no_grad(): image_idx = int(camera_ray_bundle.camera_indices[0, 0]) outputs = pipeline.model.get_outputs_for_camera_ray_bundle(camera_ray_bundle) psnr = pipeline.model.log_test_image_outputs(image_idx, step, batch, outputs) avg_rays_per_sec = _update_avg(avg_rays_per_sec, camera_ray_bundle.origins.shape[0] / t.duration, step) avg_psnr = _update_avg(avg_psnr, psnr, step) avg_fps = _update_avg(avg_fps, 1 / t.duration, step) return {"avg psnr": avg_psnr, "avg rays per sec": avg_rays_per_sec, "avg fps": avg_fps} def _render_trajectory_video( pipeline: Pipeline, cameras: Cameras, output_filename: Path, rendered_output_name: str, rendered_resolution_scaling_factor: float = 1.0, num_rays_per_chunk: int = 4096, ) -> None: """Helper function to create a video of the spiral trajectory. Args: pipeline: Pipeline to evaluate with. cameras: Cameras to render. output_filename: Name of the output file. rendered_output_name: Name of the renderer output to use. rendered_resolution_scaling_factor: Scaling factor to apply to the camera image resolution. Defaults to 1.0. num_rays_per_chunk: Number of rays to use per chunk. Defaults to 4096. """ print("Creating trajectory video.") images = [] cameras.rescale_output_resolution(rendered_resolution_scaling_factor) for camera_idx in tqdm(range(cameras.size)): camera_ray_bundle = cameras.generate_rays(camera_indices=camera_idx).to(pipeline.device) camera_ray_bundle.num_rays_per_chunk = num_rays_per_chunk with torch.no_grad(): outputs = pipeline.model.get_outputs_for_camera_ray_bundle(camera_ray_bundle) image = outputs[rendered_output_name].cpu().numpy() images.append(image) seconds = 5.0 fps = len(images) / seconds media.write_video(output_filename, images, fps=fps) def _eval_setup(config_path: Path) -> Tuple[cfg.Config, Pipeline, Path]: """Shared setup for loading a saved pipeline for evaluation. Args: config_path: Path to config YAML file. Returns: Loaded config, pipeline module, and corresponding checkpoint. """ # load save config config = yaml.load(config_path.read_text(), Loader=yaml.Loader) assert isinstance(config, cfg.Config) # load checkpoints from wherever they were saved # TODO: expose the ability to choose an arbitrary checkpoint config.trainer.load_dir = config.trainer.model_dir config.pipeline.dataloader.eval_image_indices = None # setup pipeline (which includes the dataloaders) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") pipeline = config.pipeline.setup(device=device, test_mode=True) assert isinstance(pipeline, Pipeline) pipeline.eval() # load checkpointed information checkpoint_path = _load_checkpoint(config.trainer, pipeline) return config, pipeline, checkpoint_path @dataclass class ComputePSNR: """Load a checkpoint, compute some PSNR metrics, and save to a JSON.""" # Path to config YAML file. load_config: Path # Name of the output file. output_path: Path = Path("output.json") def main(self) -> None: """Main function.""" config, pipeline, checkpoint_path = _eval_setup(self.load_config) assert self.output_path.suffix == ".json" stats_dict = _render_stats_dict(pipeline) avg_psnr = stats_dict["avg psnr"] avg_rays_per_sec = stats_dict["avg rays per sec"] avg_fps = stats_dict["avg fps"] print(f"Avg. PSNR: {avg_psnr:0.4f}") print(f"Avg. Rays / Sec: {human_format(avg_rays_per_sec)}") print(f"Avg. FPS: {avg_fps:0.4f}") # save output to some file self.output_path.parent.mkdir(parents=True, exist_ok=True) benchmark_info = { "experiment_name": config.experiment_name, "method_name": config.method_name, "checkpoint": str(checkpoint_path), "results": stats_dict, } self.output_path.write_text(json.dumps(benchmark_info, indent=2), "utf8") print(f"Saved results to: {self.output_path}") @dataclasses.dataclass class RenderTrajectory: """Load a checkpoint, render a trajectory, and save to a video file.""" # Path to config YAML file. load_config: Path # Name of the renderer output to use. rgb, depth, etc. rendered_output_name: str = "rgb" # Trajectory to render. traj: Literal["spiral", "interp"] = "spiral" # Scaling factor to apply to the camera image resolution. rendered_resolution_scaling_factor: float = 1.0 # Name of the output file. output_path: Path = Path("output.mp4") def main(self) -> None: """Main function.""" config, pipeline, _ = _eval_setup(self.load_config) # TODO(ethan): use camera information from parsing args if self.traj == "spiral": camera_start = pipeline.dataloader.eval_dataloader.get_camera(image_idx=0) # TODO(ethan): pass in the up direction of the camera camera_path = get_spiral_path(camera_start, steps=30, radius=0.1) elif self.traj == "interp": cameras = pipeline.dataloader.eval_dataloader.get_camera(image_idx=[0, 10]) camera_path = get_interpolated_camera_path(cameras, steps=30) else: assert_never(self.traj) _render_trajectory_video( pipeline, camera_path, output_filename=self.output_path, rendered_output_name=self.rendered_output_name, rendered_resolution_scaling_factor=self.rendered_resolution_scaling_factor, num_rays_per_chunk=config.pipeline.dataloader.eval_num_rays_per_chunk, ) if __name__ == "__main__": # A Union over dataclass types will create a subcommand for each type. # # TODO: it would make sense to split this script up into separate scripts. # - To reduce duplicate code, some "upstream" refactor could also simplify the # shared `_load_checkpoint()` and `_eval_setup()` helpers. The high-level # operations implemented by each seem fairly universal; ideally the checkpoint # loading logic, for example, would be the same as what's used for loading a # checkpoint when resuming a training run. dcargs.cli(Union[ComputePSNR, RenderTrajectory]).main()