From 5cee32074afd0c7631514fd6b6857b117874c7cf Mon Sep 17 00:00:00 2001 From: Matthew Tancik Date: Tue, 19 Jul 2022 18:23:40 -0700 Subject: [PATCH] Add dcargs (#178) * Add dcargs * Add dcargs to run_actions --- environment/requirements.txt | 1 + scripts/debugging/run_actions.py | 9 +-- scripts/run_eval.py | 100 +++++++++++++++++-------------- 3 files changed, 60 insertions(+), 50 deletions(-) diff --git a/environment/requirements.txt b/environment/requirements.txt index 06cbcefc3..c165d97c2 100644 --- a/environment/requirements.txt +++ b/environment/requirements.txt @@ -2,6 +2,7 @@ aiohttp aiortc av black[jupyter]==22.3.0 +dcargs functorch==0.1.1 h5py>=2.9.0 hydra-colorlog==1.2.0 diff --git a/scripts/debugging/run_actions.py b/scripts/debugging/run_actions.py index acdad7b98..b6d865360 100644 --- a/scripts/debugging/run_actions.py +++ b/scripts/debugging/run_actions.py @@ -1,8 +1,8 @@ """Simple yaml debugger""" import subprocess import sys -from typing import Optional +import dcargs import yaml LOCAL_TESTS = ["Run license checks", "Run Black", "Python Pylint", "Test with pytest"] @@ -13,6 +13,7 @@ def run_command(command: str, continue_on_fail: bool = False) -> None: Args: command: command to run + continue_on_fail: whether to continue running commands if the current one fails. Defaults to False. """ ret_code = subprocess.call(command, shell=True) if not continue_on_fail and ret_code != 0: @@ -20,11 +21,11 @@ def run_command(command: str, continue_on_fail: bool = False) -> None: sys.exit(1) -def main(continue_on_fail: Optional[bool] = False): +def main(continue_on_fail: bool = False): """Run the github actions locally. Args: - continue_on_fail: Whether to continue running actions commands if the current one fails. Defaults to False. + continue_on_fail: Whether to continue running actions commands if the current one fails """ with open(".github/workflows/code_checks.yml", "rb") as f: my_dict = yaml.safe_load(f) @@ -56,4 +57,4 @@ def main(continue_on_fail: Optional[bool] = False): if __name__ == "__main__": - main(continue_on_fail=False) + dcargs.cli(main) diff --git a/scripts/run_eval.py b/scripts/run_eval.py index 8b07841ca..ecf85e91a 100644 --- a/scripts/run_eval.py +++ b/scripts/run_eval.py @@ -1,21 +1,21 @@ """ run_eval.py """ -import argparse +import enum import os -from typing import Dict, Optional +from typing import Dict, List, Optional +import dcargs import mediapy as media import torch from hydra import compose, initialize from omegaconf import DictConfig from tqdm import tqdm -from pyrad.cameras.camera_paths import get_interpolated_camera_path, get_spiral_path, CameraPath -from pyrad.data.dataloader import setup_dataset_eval, setup_dataset_train +from pyrad.cameras.camera_paths import CameraPath, get_interpolated_camera_path, get_spiral_path +from pyrad.data.dataloader import EvalDataloader, setup_dataset_eval, setup_dataset_train from pyrad.graphs.base import Graph, setup_graph from pyrad.utils.writer import TimeWriter -from pyrad.data.dataloader import EvalDataloader def _update_avg(prev_avg: float, new_val: float, step: int) -> float: @@ -138,44 +138,52 @@ def render_trajectory_video( media.write_video(output_filename, images, fps=fps) -def my_func_that_returns_a_parser(): - """Function that returns a parser, which - can be used with sphinx to generate documentation. +class MethodType(enum.Enum): + """Enum for the method type.""" + + PSNR = enum.auto() + TRAJ = enum.auto() + + +class TrajectoryType(enum.Enum): + """Enum for the trajectory type.""" + + SPIRAL = enum.auto() + INTERP = enum.auto() + + +def main( + config_name: str, + checkpoint_dir: str, + rendered_output_name: str, + method: MethodType = MethodType.PSNR, + traj: TrajectoryType = TrajectoryType.SPIRAL, + output_filename: str = "output.mp4", + rendered_resolution_scaling_factor: float = 1.0, + config_overrides: Optional[List[str]] = None, +): + """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. + 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. """ - parser = argparse.ArgumentParser(description="Run the evaluation of a model.") - parser.add_argument( - "--method", - type=str, - default="psnr", - choices=["psnr", "traj"], - help="Specify which type of evaluation method to run.", - ) - parser.add_argument("--traj", type=str, default="spiral", choices=["spiral", "interp"]) - parser.add_argument("--output-filename", type=str, default="output.mp4") - parser.add_argument("--config-name", type=str, default="graph_default.yaml") - parser.add_argument( - "--rendered-output-name", type=str, default=None, help="Name of the rendered output to use from the Graph." - ) - parser.add_argument( - "--rendered-resolution-scaling-factor", - type=float, - default=1.0, - help="Scaling factor to apply to the rendered camera image resolution.", - ) - parser.add_argument("overrides", nargs="*", default=[]) - return parser - - -def main(): - """Main function.""" - parser = my_func_that_returns_a_parser() - args = parser.parse_args() + # parser = my_func_that_returns_a_parser() + # args = parser.parse_args() config_path = "../configs" initialize(version_base="1.2", config_path=config_path) - config = compose(args.config_name, overrides=args.overrides) + config_overrides = config_overrides or [] + config = compose(config_name, overrides=config_overrides) - assert config.trainer.resume_train.load_dir, "Please specify checkpoint load path" + config.trainer.resume_train.load_dir = checkpoint_dir device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # setup graph and dataset @@ -186,7 +194,7 @@ def main(): # load checkpointed information _load_checkpoint(config.trainer.resume_train, graph) - if args.method == "psnr": + if method == MethodType.PSNR: stats_dict = render_stats_dict(graph, dataloader_eval) avg_psnr = stats_dict["avg psnr"] avg_rays_per_sec = stats_dict["avg rays per sec"] @@ -194,25 +202,25 @@ def main(): print(f"Avg. PSNR: {avg_psnr:0.4f}") print(f"Avg. Rays per sec: {avg_rays_per_sec:0.4f}") print(f"Avg. FPS: {avg_fps:0.4f}") - elif args.method == "traj": + elif method == MethodType.TRAJ: # TODO(ethan): pass in camera information into argparse parser - if args.traj == "spiral": + if traj == TrajectoryType.SPIRAL: camera_start = dataloader_eval.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 args.traj == "interp": + elif traj == TrajectoryType.INTERP: camera_start = dataloader_eval.get_camera(image_idx=0) camera_end = dataloader_eval.get_camera(image_idx=10) camera_path = get_interpolated_camera_path(camera_start, camera_end, steps=30) render_trajectory_video( graph, camera_path, - output_filename=args.output_filename, - rendered_output_name=args.rendered_output_name, - rendered_resolution_scaling_factor=args.rendered_resolution_scaling_factor, + output_filename=output_filename, + rendered_output_name=rendered_output_name, + rendered_resolution_scaling_factor=rendered_resolution_scaling_factor, num_rays_per_chunk=config.data.dataloader_eval.num_rays_per_chunk, ) if __name__ == "__main__": - main() + dcargs.cli(main)