Fixing nerfw and adding to GH Actions (#167)

* added beta_min to nerfw, which I think is needed

* added camera interpolation path and restructured run_eval code

* adding to gitignore

* addressing github comments. mostly documentation

* fixing nerfw dataset for test cases

scripts/benchmarking/run_benchmark.py

@@ -13,7 +13,7 @@ import numpy as np
 from omegaconf import DictConfig
 from tqdm import tqdm
 
-from scripts.run_eval import run_inference
+from scripts.run_eval import run_inference_from_config
 
 
 OBJECT_LIST = ["mic", "ficus", "chair", "hotdog", "materials", "drums", "ship", "lego"]
@@ -78,7 +78,7 @@ def main(args):
         ckpt = _load_best_ckpt(hydra_dir, config.trainer)
 
         # run evaluation
-        stats_dict = run_inference(config)
+        stats_dict = run_inference_from_config(config)
         stats_dict["checkpoint"] = ckpt
         benchmarks[dataset] = stats_dict
 

scripts/debugging/run_actions.py

@@ -1,25 +1,31 @@
 """Simple yaml debugger"""
 import subprocess
 import sys
+from typing import Optional
 
 import yaml
 
 LOCAL_TESTS = ["Run license checks", "Run Black", "Python Pylint", "Test with pytest"]
 
 
-def run_command(command) -> None:
+def run_command(command: str, continue_on_fail: bool = False) -> None:
     """Run a command kill actions if it fails
 
     Args:
-        command (str): command to run
+        command: command to run
     """
     ret_code = subprocess.call(command, shell=True)
-    if ret_code != 0:
+    if not continue_on_fail and ret_code != 0:
         print(f"\033[31mError: `{command}` failed. Exiting...\033[0m")
         sys.exit(1)
 
 
-if __name__ == "__main__":
+def main(continue_on_fail: Optional[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.
+    """
     with open(".github/workflows/code_checks.yml", "rb") as f:
         my_dict = yaml.safe_load(f)
     steps = my_dict["jobs"]["build"]["steps"]
@@ -34,7 +40,7 @@ if __name__ == "__main__":
 
             print("*" * 100)
             print_green(f"Running: {curr_command}")
-            run_command(curr_command)
+            run_command(curr_command, continue_on_fail=continue_on_fail)
         else:
             print("*" * 100)
             print(f"Skipping {step}")
@@ -47,3 +53,7 @@ if __name__ == "__main__":
     print_green("=" * 100)
     print_green("ALL CHECKS PASSED")
     print_green("=" * 100)
+
+
+if __name__ == "__main__":
+    main(continue_on_fail=False)

scripts/run_eval.py

@@ -3,7 +3,7 @@ run_eval.py
 """
 import argparse
 import os
-from typing import Dict
+from typing import Dict, Optional
 
 import mediapy as media
 import torch
@@ -11,10 +11,11 @@ from hydra import compose, initialize
 from omegaconf import DictConfig
 from tqdm import tqdm
 
-from pyrad.cameras.camera_paths import get_spiral_path
+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.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:
@@ -52,7 +53,7 @@ def _load_checkpoint(config: DictConfig, graph: Graph) -> None:
     print(f"done loading checkpoint from {load_path}")
 
 
-def run_inference(config: DictConfig, local_rank: int = 0, world_size: int = 1) -> Dict[str, float]:
+def run_inference_from_config(config: DictConfig) -> Dict[str, float]:
     """helper function to run inference given config specifications (also used in benchmarking)
 
     Args:
@@ -61,7 +62,8 @@ def run_inference(config: DictConfig, local_rank: int = 0, world_size: int = 1)
     Returns:
         Tuple[float, float]: returns both the avg psnr and avg rays per second
     """
-    device = "cpu" if world_size == 0 else f"cuda:{local_rank}"
+    print("Running inference.")
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     # setup graph and dataset
     dataset_inputs_train, _ = setup_dataset_train(config.data, device=device)
     _, dataloader_eval = setup_dataset_eval(config.data, test_mode=True, device=device)
@@ -72,7 +74,19 @@ def run_inference(config: DictConfig, local_rank: int = 0, world_size: int = 1)
     _load_checkpoint(config.trainer.resume_train, graph)
 
     # calculate average psnr across test dataset
-    # TODO(ethan): trajector specification
+    return render_stats_dict(graph, dataloader_eval)
+
+
+def render_stats_dict(graph: Graph, dataloader_eval: EvalDataloader) -> Dict[str, float]:
+    """Helper function to evaluate the graph on a dataloader.
+
+    Args:
+        graph (Graph): Graph to evaluate
+        dataloader_eval (EvalDataloader): Dataloader to evaluate on
+
+    Returns:
+        dict: returns the average psnr and average rays per second
+    """
     avg_psnr = 0
     avg_rays_per_sec = 0
     avg_fps = 0
@@ -88,44 +102,35 @@ def run_inference(config: DictConfig, local_rank: int = 0, world_size: int = 1)
     return {"avg psnr": avg_psnr, "avg rays per sec": avg_rays_per_sec, "avg fps": avg_fps}
 
 
-def create_spiral_video(
-    config: DictConfig, local_rank: int = 0, world_size: int = 1, output_filename: str = None
+def render_trajectory_video(
+    graph: Graph,
+    camera_path: CameraPath,
+    output_filename: Optional[str] = None,
+    rendered_output_name: Optional[str] = None,
+    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:
-        config (DictConfig): Configuration for loading the evaluation.
-        local_rank (int): Local rank of the process.
-        world_size (int): Total number of GPUs.
-        output_filename (str): Name of the output file.
+        graph: Graph 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 spiral video")
-    device = "cpu" if world_size == 0 else f"cuda:{local_rank}"
-    # setup graph and dataset
-    dataset_inputs_train, _ = setup_dataset_train(config.data, device=device)
-    _, dataloader_eval = setup_dataset_eval(config.data, test_mode=True, device=device)
-    graph = setup_graph(config.graph, dataset_inputs_train, device=device)
-
-    # load checkpointed information
-    _load_checkpoint(config.trainer.resume_train, graph)
-
-    # get a trajecory
-    start_camera = dataloader_eval.get_camera(image_idx=0)
-    # TODO(ethan): replace with radius with radiuses, based on camera pose percentiles
-    # see original nerf paper code for details
-    camera_path = get_spiral_path(start_camera, steps=60, radius=0.5)
-
+    print("Creating trajectory video.")
     images = []
     for camera in tqdm(camera_path.cameras):
-        camera.cx /= 4
-        camera.cy /= 4
-        camera.fx /= 4
-        camera.fy /= 4
-        camera_ray_bundle = camera.get_camera_ray_bundle().to(device)
-        camera_ray_bundle.num_rays_per_chunk = 4096
-        outputs = graph.get_outputs_for_camera_ray_bundle(camera_ray_bundle)
+        camera.rescale_output_resolution(rendered_resolution_scaling_factor)
+        camera_ray_bundle = camera.get_camera_ray_bundle().to(graph.device)
+        camera_ray_bundle.num_rays_per_chunk = num_rays_per_chunk
+        with torch.no_grad():
+            outputs = graph.get_outputs_for_camera_ray_bundle(camera_ray_bundle)
         # TODO: don't hardcode the key! this will break for some nerf Graphs
-        image = outputs["rgb"].cpu().numpy()
+        image = outputs[rendered_output_name].cpu().numpy()
         images.append(image)
 
     seconds = 5.0
@@ -133,9 +138,10 @@ def create_spiral_video(
     media.write_video(output_filename, images, fps=fps)
 
 
-def main():
-    """Main function."""
-
+def my_func_that_returns_a_parser():
+    """Function that returns a parser, which
+    can be used with sphinx to generate documentation.
+    """
     parser = argparse.ArgumentParser(description="Run the evaluation of a model.")
     parser.add_argument(
         "--method",
@@ -147,7 +153,22 @@ def main():
     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()
 
     config_path = "../configs"
@@ -155,10 +176,18 @@ def main():
     config = compose(args.config_name, overrides=args.overrides)
 
     assert config.trainer.resume_train.load_dir, "Please specify checkpoint load path"
-    assert args.traj != "interp", "Camera pose interpolation trajectory isn't yet implemented."
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    # setup graph and dataset
+    dataset_inputs_train, _ = setup_dataset_train(config.data, device=device)
+    _, dataloader_eval = setup_dataset_eval(config.data, test_mode=True, device=device)
+    graph = setup_graph(config.graph, dataset_inputs_train, device=device)
+    graph.eval()
+    # load checkpointed information
+    _load_checkpoint(config.trainer.resume_train, graph)
 
     if args.method == "psnr":
-        stats_dict = run_inference(config)
+        stats_dict = render_stats_dict(graph, dataloader_eval)
         avg_psnr = stats_dict["avg psnr"]
         avg_rays_per_sec = stats_dict["avg rays per sec"]
         avg_fps = stats_dict["avg fps"]
@@ -166,7 +195,23 @@ def main():
         print(f"Avg. Rays per sec: {avg_rays_per_sec:0.4f}")
         print(f"Avg. FPS: {avg_fps:0.4f}")
     elif args.method == "traj":
-        create_spiral_video(config, output_filename=args.output_filename)
+        # TODO(ethan): pass in camera information into argparse parser
+        if args.traj == "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":
+            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,
+            num_rays_per_chunk=config.data.dataloader_eval.num_rays_per_chunk,
+        )
 
 
 if __name__ == "__main__":