""" run_eval.py """ import json import logging import os import pickle from pathlib import Path from typing import Dict import mediapy as media import torch from tqdm import tqdm from nerfactory.cameras.camera_paths import ( CameraPath, get_interpolated_camera_path, get_spiral_path, ) from nerfactory.configs import base as cfg from nerfactory.configs.utils import cli_from_base_configs 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, camera_path: CameraPath, 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. camera_path: Index of the image 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 = [] for camera in tqdm(camera_path.cameras): camera.rescale_output_resolution(rendered_resolution_scaling_factor) camera_ray_bundle = camera.get_camera_ray_bundle().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 main(config: cfg.Config): """Evaluate trained model. This evaluation can either render a trajectory or compute the eval psnr. Args: config_name: Name of the config file to use. checkpoint_dir: Directory to load the checkpoint from. rendered_output_name: Name of the renderer output to use (rgb, depth, etc.). method: Method to use for evaluation. PSNR computes metrics, TRAJ renders a trajectory. traj: Trajectory to render. output_filename: Name of the output file. rendered_resolution_scaling_factor: Scaling factor to apply to the camera image resolution. Defaults to 1.0. config_overrides: List of strings to override config values. """ assert config.eval is not None, "Eval config needs to be set" assert config.eval.checkpoint_dir is not None, "Missing eval checkpoint dir: --eval.checkpoint-dir" config.trainer.load_dir = config.eval.checkpoint_dir config.pipeline.dataloader.eval_image_indices = None device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # setup pipeline (which includes the dataloaders) pipeline = config.pipeline.setup(device=device, test_mode=True) pipeline.eval() # load checkpointed information checkpoint_name = _load_checkpoint(config.trainer, pipeline) if config.eval.method == cfg.MethodType.PSNR: assert config.eval.output_filename.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 config.eval.output_filename.parent.mkdir(parents=True, exist_ok=True) benchmark_info = { "experiment_name": config.experiment_name, "method_name": config.method_name, "checkpoint": str(checkpoint_name), "results": stats_dict, } with open(config.eval.output_filename, "w", encoding="utf8") as output_f: json.dump(benchmark_info, output_f, indent=2) print(f"Saved results to: {config.eval.output_filename}") elif config.eval.method == cfg.MethodType.TRAJ: assert config.eval.output_filename.suffix == ".mp4" assert config.eval.rendered_output_name is not None, "Missing eval output type: --eval.rendered-output-name" # TODO(ethan): use camera information from parsing args if config.eval.traj == cfg.TrajectoryType.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 config.eval.traj == cfg.TrajectoryType.INTERP: camera_start = pipeline.dataloader.eval_dataloader.get_camera(image_idx=0) camera_end = pipeline.dataloader.eval_dataloader.get_camera(image_idx=10) camera_path = get_interpolated_camera_path(camera_start, camera_end, steps=30) else: raise NotImplementedError render_trajectory_video( pipeline, camera_path, output_filename=config.eval.output_filename, rendered_output_name=config.eval.rendered_output_name, rendered_resolution_scaling_factor=config.eval.rendered_resolution_scaling_factor, num_rays_per_chunk=pipeline.dataloader.eval_num_rays_per_chunk, ) if __name__ == "__main__": from nerfactory.configs.base_configs import base_configs instantiated_config = cli_from_base_configs(base_configs, eval_mode=True) if instantiated_config.eval.load_config: logging.info(f"Loading pre-set config to: {instantiated_config.eval.load_config}") with open(instantiated_config.eval.load_config, "rb") as config_f: train_instantiated_config = pickle.load(config_f) train_instantiated_config.eval = instantiated_config.eval instantiated_config = train_instantiated_config logging.info("Printing current config setup") print(instantiated_config) main(instantiated_config)