diff --git a/scripts/run_eval.py b/scripts/run_eval.py index 6deef0e18..c7e31b771 100755 --- a/scripts/run_eval.py +++ b/scripts/run_eval.py @@ -16,21 +16,20 @@ import dcargs import mediapy as media import torch import yaml -from tqdm import tqdm +from rich.console import Console +from rich.progress import track from typing_extensions import assert_never # pylint: disable=unused-import -from nerfactory.cameras.camera_paths import ( - get_interpolated_camera_path, - get_path_from_json, - get_spiral_path, -) +from nerfactory.cameras.camera_paths import get_path_from_json, 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 +console = Console(width=120) + logging.basicConfig(format="[%(filename)s:%(lineno)d] %(message)s", level=logging.INFO) @@ -57,7 +56,7 @@ def _load_checkpoint(config: cfg.TrainerConfig, pipeline: Pipeline) -> Path: """ assert config.load_dir is not None if config.load_step is None: - print("Loading latest checkpoint from load_dir") + console.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: @@ -66,7 +65,7 @@ def _load_checkpoint(config: cfg.TrainerConfig, pipeline: Pipeline) -> Path: 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}") + console.print(f":white_check_mark: Done loading checkpoint from {load_path}") return load_path @@ -82,7 +81,7 @@ def _render_stats_dict(pipeline: Pipeline) -> Dict[str, float]: avg_psnr = 0 avg_rays_per_sec = 0 avg_fps = 0 - for step, (camera_ray_bundle, batch) in tqdm(enumerate(pipeline.datamanager.eval_dataloader)): + for step, (camera_ray_bundle, batch) in track(enumerate(pipeline.datamanager.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]) @@ -115,10 +114,10 @@ def _render_trajectory_video( num_rays_per_chunk: Number of rays to use per chunk. Defaults to 4096. seconds: Number for the output video. Defaults to 5.0. """ - print("Creating trajectory video.") + console.print("[bold green]Creating trajectory video") images = [] cameras.rescale_output_resolution(rendered_resolution_scaling_factor) - for camera_idx in tqdm(range(cameras.size)): + for camera_idx in track(range(cameras.size), description=":movie_camera: Rendering :movie_camera:"): 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(): @@ -127,7 +126,10 @@ def _render_trajectory_video( images.append(image) fps = len(images) / seconds - media.write_video(output_filename, images, fps=fps) + with console.status("[yellow]Saving video", spinner="bouncingBall"): + media.write_video(output_filename, images, fps=fps) + console.rule("[green] :tada: :tada: :tada: Success :tada: :tada: :tada:") + console.print(f"[green]Saved video to {output_filename}", justify="center") def _eval_setup(config_path: Path) -> Tuple[cfg.Config, Pipeline, Path]: @@ -178,9 +180,9 @@ class ComputePSNR: 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}") + console.print(f"Avg. PSNR: {avg_psnr:0.4f}") + console.print(f"Avg. Rays / Sec: {human_format(avg_rays_per_sec)}") + console.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 = { @@ -190,7 +192,7 @@ class ComputePSNR: "results": stats_dict, } self.output_path.write_text(json.dumps(benchmark_info, indent=2), "utf8") - print(f"Saved results to: {self.output_path}") + console.print(f"Saved results to: {self.output_path}") @dataclasses.dataclass