diff --git a/.gitignore b/.gitignore index 310a1c41b..ab255fd31 100644 --- a/.gitignore +++ b/.gitignore @@ -181,4 +181,7 @@ __MACOSX/ outputs*/ node_modules/ yarn.lock -package.json \ No newline at end of file +package.json +scripts/colmap2nerf.py +bash/ +pyrad/ngp_pl/ \ No newline at end of file diff --git a/.vscode/launch.json b/.vscode/launch.json index 93adebfe4..5661ddd76 100644 --- a/.vscode/launch.json +++ b/.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" ] }, { diff --git a/pyproject.toml b/pyproject.toml index bd72a5348..2f5915f45 100644 --- a/pyproject.toml +++ b/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", diff --git a/pyrad/cameras/camera_paths.py b/pyrad/cameras/camera_paths.py index 3f66c9044..743277590 100644 --- a/pyrad/cameras/camera_paths.py +++ b/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 diff --git a/pyrad/cameras/cameras.py b/pyrad/cameras/cameras.py index e045d67e0..0a21ff48a 100644 --- a/pyrad/cameras/cameras.py +++ b/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, diff --git a/pyrad/data/format/instant_ngp.py b/pyrad/data/format/instant_ngp.py index 34c44f5a7..4dfb6a7a4 100644 --- a/pyrad/data/format/instant_ngp.py +++ b/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 diff --git a/pyrad/fields/nerfw_field.py b/pyrad/fields/nerfw_field.py index 55728a96a..024ffbc20 100644 --- a/pyrad/fields/nerfw_field.py +++ b/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 diff --git a/pyrad/graphs/base.py b/pyrad/graphs/base.py index e37ff619f..492bdc9fe 100644 --- a/pyrad/graphs/base.py +++ b/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""" diff --git a/pyrad/graphs/nerfw.py b/pyrad/graphs/nerfw.py index 981712ab0..fd4f1d958 100644 --- a/pyrad/graphs/nerfw.py +++ b/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) diff --git a/scripts/benchmarking/run_benchmark.py b/scripts/benchmarking/run_benchmark.py index 797b95b69..b2a8e03bf 100644 --- a/scripts/benchmarking/run_benchmark.py +++ b/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 diff --git a/scripts/debugging/run_actions.py b/scripts/debugging/run_actions.py index 800d558b9..acdad7b98 100644 --- a/scripts/debugging/run_actions.py +++ b/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) diff --git a/scripts/run_eval.py b/scripts/run_eval.py index f51b21742..8b07841ca 100644 --- a/scripts/run_eval.py +++ b/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__": diff --git a/tests/test_run_train.py b/tests/test_run_train.py index 94b4dd3fb..40c87a140 100644 --- a/tests/test_run_train.py +++ b/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