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

.gitignore

@@ -181,4 +181,7 @@ __MACOSX/
 outputs*/
 node_modules/
 yarn.lock
-package.json
+package.json
+scripts/colmap2nerf.py
+bash/
+pyrad/ngp_pl/

.vscode/launch.json

@@ -57,7 +57,7 @@
             "console": "integratedTerminal",
             "args": [
                 "--config-name=graph_nerfw.yaml",
-                "++data.dataset.use_cache=true"
+                // "++data.dataset.use_cache=true"
             ]
         },
         {

pyproject.toml

@@ -11,7 +11,7 @@ line-length = 120
 max-line-length = 120
 generated-members = ["numpy.*", "torch.*", "cv2.*", "cv.*"]
 good-names-rgxs = "^[_a-zA-Z][_a-z0-9]?$"
-ignore-paths = ["^pyrad/cuda/.*$", "^tests/cuda/.*$", "^pyrad/viewer/.*$", "^tests/viewer/.*$"]
+ignore-paths = ["pyrad/ngp_pl/.*$", "scripts/colmap2nerf.py", "^pyrad/cuda/.*$", "^tests/cuda/.*$", "^pyrad/viewer/.*$", "^tests/viewer/.*$"]
 
 disable = [
   "too-many-arguments",

pyrad/cameras/camera_paths.py

@@ -45,20 +45,24 @@ def get_interpolated_camera_path(camera_a: Camera, camera_b: Camera, steps: int)
     Returns:
         CameraPath: A camera path.
     """
-    Ka = camera_a.get_intrinsics_matrix()
-    pose_a = camera_a.get_camera_to_world_h()
-    Kb = camera_b.get_intrinsics_matrix()
-    pose_b = camera_b.get_camera_to_world_h()
+    device = camera_a.device
+    Ka = camera_a.get_intrinsics_matrix().cpu().numpy()
+    pose_a = camera_a.get_camera_to_world_h().cpu().numpy()
+    Kb = camera_b.get_intrinsics_matrix().cpu().numpy()
+    pose_b = camera_b.get_camera_to_world_h().cpu().numpy()
     poses = [pose_a, pose_b]
     Ks = [Ka, Kb]
     poses, Ks = get_interpolated_poses_many(poses, Ks, steps_per_transition=steps)
-    # create a list of cameras
 
     cameras = []
     for pose, K in zip(poses, Ks):
-        intrinsics = get_intrinsics_from_intrinsics_matrix(K)
-        camera_to_world = torch.from_numpy(pose[:3])
-        camera = get_camera(intrinsics, camera_to_world, None)
+        intrinsics = get_intrinsics_from_intrinsics_matrix(K).to(device).float()
+        # TODO: this makes cx and cy an integer, but this code should be fixed
+        # it was added to avoid floating point errors when rescaling and the image resolution
+        # being different per rendered image in a camera path
+        intrinsics[:2] = torch.round(intrinsics[:2])
+        camera_to_world = torch.from_numpy(pose[:3]).to(device).float()
+        camera = get_camera(intrinsics, camera_to_world, camera_index=camera_a.camera_index)
         cameras.append(camera)
     return CameraPath(cameras=cameras)
 
@@ -94,7 +98,9 @@ def get_spiral_path(
     else:
         raise ValueError("Only one of radius or radiuses must be specified.")
 
+    # TODO: don't hardcode this. pass this in
     up = camera.camera_to_world[:3, 2]  # scene is z up
+    # up = camera.camera_to_world[:3, 1] # this will rotate 90 degrees
     focal = min(camera.fx, camera.fy)
     target = torch.tensor([0, 0, -focal], device=camera.device)  # camera looking in -z direction
 

pyrad/cameras/cameras.py

@@ -91,6 +91,18 @@ class Camera:
         """
         return
 
+    @abstractmethod
+    def rescale_output_resolution(self, scaling_factor: float) -> None:
+        """Rescales the camera intrinsics for output resolution.
+
+        Args:
+            scaling_factor (float): Scaling factor
+
+        Returns:
+            None
+        """
+        return
+
     def get_image_coords(self, pixel_offset: float = 0.5) -> TensorType["image_height", "image_width", 2]:
         """
         Args:
@@ -207,6 +219,12 @@ class PinholeCamera(Camera):
         """
         return 3
 
+    def rescale_output_resolution(self, scaling_factor: float) -> None:
+        self.cx *= scaling_factor
+        self.cy *= scaling_factor
+        self.fx *= scaling_factor
+        self.fy *= scaling_factor
+
     @classmethod
     def generate_rays(
         cls,
@@ -289,6 +307,10 @@ class EquirectangularCamera(Camera):
     def get_intrinsics(self) -> torch.Tensor:
         return torch.tensor([self.height, self.width])
 
+    def rescale_output_resolution(self, scaling_factor: float) -> None:
+        self.height *= scaling_factor
+        self.width *= scaling_factor
+
     @classmethod
     def generate_rays(
         cls,

pyrad/data/format/instant_ngp.py

@@ -57,6 +57,12 @@ def load_instant_ngp_data(
             poses.append(np.array(frame["transform_matrix"]))
     if num_skipped_image_filenames >= 0:
         logging.info("Skipping %s files in dataset split %s.", num_skipped_image_filenames, split)
+    assert (
+        len(image_filenames) != 0
+    ), """
+    No image files found. 
+    You should check the file_paths in the transforms.json file to make sure they are correct.
+    """
     poses = np.array(poses).astype(np.float32)
     poses[:3, 3] *= camera_translation_scalar
 

pyrad/fields/nerfw_field.py

@@ -108,20 +108,30 @@ class VanillaNerfWField(Field):
 
     def get_density(self, ray_samples: RaySamples):
         """Computes and returns the densities."""
-        encoded_xyz = self.position_encoding(ray_samples.frustums.positions())
+        encoded_xyz = self.position_encoding(ray_samples.frustums.get_positions())
+        encoded_xyz = self.position_encoding(ray_samples.frustums.get_positions())
         base_mlp_out = self.mlp_base(encoded_xyz)
         density = self.field_head_density(base_mlp_out)
         return density, base_mlp_out
 
     def get_outputs(
-        self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None
+        self, ray_samples: RaySamples, density_embedding: Optional[TensorType[..., "embedding_size"]] = None
     ) -> Dict[FieldHeadNames, TensorType]:
+        """Returns the outputs of the NeRF-W field.
+
+        Args:
+            ray_samples (RaySamples): Ray samples.
+            density_embedding (TensorType[..., "embedding_size"], optional): Density embedding. Defaults to None.
+
+        Returns:
+            Dict[FieldHeadNames, TensorType]: Outputs of the NeRF-W field.
+        """
         outputs = {}
         encoded_dir = self.direction_encoding(ray_samples.frustums.directions)
-        embedded_appearance = self.embedding_appearance(ray_samples.camera_indices)
+        embedded_appearance = self.embedding_appearance(ray_samples.camera_indices.squeeze())
         mlp_head_out = self.mlp_head(torch.cat([density_embedding, encoded_dir, embedded_appearance], dim=-1))
         outputs[self.field_head_rgb.field_head_name] = self.field_head_rgb(mlp_head_out)  # static rgb
-        embedded_transient = self.embedding_transient(ray_samples.camera_indices)
+        embedded_transient = self.embedding_transient(ray_samples.camera_indices.squeeze())
         transient_mlp_out = self.mlp_transient(torch.cat([density_embedding, embedded_transient], dim=-1))
         outputs[self.field_head_transient_uncertainty.field_head_name] = self.field_head_transient_uncertainty(
             transient_mlp_out

pyrad/graphs/base.py

@@ -59,6 +59,11 @@ class AbstractGraph(nn.Module):
         """Returns the device that the torch parameters are on."""
         return self.device_indicator_param.device
 
+    @property
+    def device(self):
+        """Returns the device that the graph is on."""
+        return self.device_indicator_param.device
+
     @abstractmethod
     def forward(self, ray_indices: TensorType["num_rays", 3], batch: Union[str, Dict[str, torch.tensor]] = None):
         """Process starting with ray indices. Turns them into rays, then performs volume rendering."""
@@ -104,6 +109,7 @@ class Graph(AbstractGraph):
         self.callbacks = None
         # variable for visualizer to fetch TODO(figure out if there is cleaner way to do this)
         self.vis_outputs = None
+        self.default_output_name = None
 
     def register_callbacks(self):  # pylint:disable=no-self-use
         """Option to register callback for training functions"""

pyrad/graphs/nerfw.py

@@ -44,13 +44,24 @@ class NerfWGraph(Graph):
         near_plane=2.0,
         far_plane=6.0,
         num_coarse_samples=64,
-        num_importance_samples=128,
+        num_importance_samples=64,
+        uncertainty_min=0.03,
         **kwargs,
     ) -> None:
+        """A NeRF-W graph.
+
+        Args:
+            ...
+            uncertainty_min (float, optional): This is added to the end of the uncertainty
+                rendering operation. It's called 'beta_min' in other repos.
+                This avoids calling torch.log() on a zero value, which would be undefined.
+                Defaults to 0.03.
+        """
         self.near_plane = near_plane
         self.far_plane = far_plane
         self.num_coarse_samples = num_coarse_samples
         self.num_importance_samples = num_importance_samples
+        self.uncertainty_min = uncertainty_min
         self.field_coarse = None
         self.field_fine = None
         self.num_images = len(intrinsics)
@@ -151,6 +162,7 @@ class NerfWGraph(Graph):
 
         # uncertainty
         uncertainty = self.renderer_uncertainty(field_outputs_fine[FieldHeadNames.UNCERTAINTY], weights_fine_transient)
+        uncertainty += self.uncertainty_min
 
         outputs = {
             "rgb_coarse": rgb_coarse,  # (num_rays, 3)
@@ -206,5 +218,6 @@ class NerfWGraph(Graph):
 
         writer.put_image(name=f"img/image_idx_{image_idx}-nerfw", image=combined_image, step=step)
 
-        mask = batch["mask"].repeat(1, 1, 3)
-        writer.put_image(name=f"mask/image_idx_{image_idx}", image=mask, step=step)
+        if "mask" in batch:
+            mask = batch["mask"].repeat(1, 1, 3)
+            writer.put_image(name=f"mask/image_idx_{image_idx}", image=mask, step=step)

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__":

tests/test_run_train.py

@@ -10,7 +10,7 @@ from omegaconf import DictConfig
 
 from scripts.run_train import main
 
-BLACKLIST = ["graph_nerfw.yaml", "graph_semantic_nerf.yaml", "graph_mipnerf_360.yaml"]
+BLACKLIST = ["graph_semantic_nerf.yaml", "graph_mipnerf_360.yaml"]
 
 
 def set_reduced_config(config: DictConfig):
@@ -18,11 +18,20 @@ def set_reduced_config(config: DictConfig):
     config.machine.num_gpus = 0
     config.trainer.max_num_iterations = 2
     # reduce dataset factors; set dataset to test
+    # switch to using the vanilla ImageDataset class
+    config.data.image_dataset_train._target_ = ( # pylint: disable=protected-access
+        "pyrad.data.image_dataset.ImageDataset"
+    )
     config.data.dataset_inputs_train.data_directory = "tests/data/lego_test"
+    config.data.dataset_inputs_train.dataset_format = "blender"
     config.data.dataset_inputs_train.downscale_factor = 16
     config.data.dataloader_train.image_sampler.num_images_to_sample_from = 1
     config.data.dataloader_train.pixel_sampler.num_rays_per_batch = 4
+    config.data.image_dataset_eval._target_ = ( # pylint: disable=protected-access
+        "pyrad.data.image_dataset.ImageDataset"
+    )
     config.data.dataset_inputs_eval.data_directory = "tests/data/lego_test"
+    config.data.dataset_inputs_eval.dataset_format = "blender"
     config.data.dataset_inputs_eval.downscale_factor = 16
     # reduce graph factors
     config.graph.num_coarse_samples = 4