diff --git a/nerfstudio/data/dataparsers/colmap_dataparser.py b/nerfstudio/data/dataparsers/colmap_dataparser.py index bb3284f33..dac6aebeb 100644 --- a/nerfstudio/data/dataparsers/colmap_dataparser.py +++ b/nerfstudio/data/dataparsers/colmap_dataparser.py @@ -174,7 +174,6 @@ class ColmapDataParser(DataParser): split_filenames = set(self.config.images_path / x for x in filenames) unmatched_filenames = split_filenames.difference(image_filenames) if unmatched_filenames: - breakpoint() raise RuntimeError( f"Some filenames for split {split} were not found: {set(map(str, unmatched_filenames))}." ) diff --git a/nerfstudio/fields/tensorf_field.py b/nerfstudio/fields/tensorf_field.py index 27d5f3cd4..b3aba4e54 100644 --- a/nerfstudio/fields/tensorf_field.py +++ b/nerfstudio/fields/tensorf_field.py @@ -98,14 +98,14 @@ class TensoRFField(Field): rgb_features = self.color_encoding(positions) rgb_features = self.B(rgb_features) - d_encoded = self.direction_encoding(d) - rgb_features_encoded = self.feature_encoding(rgb_features) - if self.use_sh: sh_mult = self.sh(d)[:, :, None] rgb_sh = rgb_features.view(sh_mult.shape[0], sh_mult.shape[1], 3, sh_mult.shape[-1]) rgb = torch.relu(torch.sum(sh_mult * rgb_sh, dim=-1) + 0.5) else: + d_encoded = self.direction_encoding(d) + rgb_features_encoded = self.feature_encoding(rgb_features) + out = self.mlp_head(torch.cat([rgb_features, d, rgb_features_encoded, d_encoded], dim=-1)) # type: ignore rgb = self.field_output_rgb(out) diff --git a/nerfstudio/models/nerfacto.py b/nerfstudio/models/nerfacto.py index 75e798873..2dbbee939 100644 --- a/nerfstudio/models/nerfacto.py +++ b/nerfstudio/models/nerfacto.py @@ -70,7 +70,7 @@ class NerfactoModelConfig(ModelConfig): num_levels: int = 16 """Number of levels of the hashmap for the base mlp.""" base_res: int = 16 - """Resolution of the base grid for the hasgrid.""" + """Resolution of the base grid for the hashgrid.""" max_res: int = 2048 """Maximum resolution of the hashmap for the base mlp.""" log2_hashmap_size: int = 19 diff --git a/nerfstudio/pipelines/base_pipeline.py b/nerfstudio/pipelines/base_pipeline.py index e9d30a9ab..363c8c9ea 100644 --- a/nerfstudio/pipelines/base_pipeline.py +++ b/nerfstudio/pipelines/base_pipeline.py @@ -20,11 +20,13 @@ from __future__ import annotations import typing from abc import abstractmethod from dataclasses import dataclass, field +from pathlib import Path from time import time from typing import Any, Dict, List, Literal, Mapping, Optional, Tuple, Type, Union, cast import torch import torch.distributed as dist +from PIL import Image from rich.progress import ( BarColumn, MofNCompleteColumn, @@ -173,8 +175,13 @@ class Pipeline(nn.Module): @abstractmethod @profiler.time_function - def get_average_eval_image_metrics(self, step: Optional[int] = None): - """Iterate over all the images in the eval dataset and get the average.""" + def get_average_eval_image_metrics(self, step: Optional[int] = None, output_path: Optional[Path] = None): + """Iterate over all the images in the eval dataset and get the average. + + Args: + step: current training step + output_path: optional path to save rendered images to + """ def load_pipeline(self, loaded_state: Dict[str, Any], step: int) -> None: """Load the checkpoint from the given path @@ -339,9 +346,13 @@ class VanillaPipeline(Pipeline): return metrics_dict, images_dict @profiler.time_function - def get_average_eval_image_metrics(self, step: Optional[int] = None): + def get_average_eval_image_metrics(self, step: Optional[int] = None, output_path: Optional[Path] = None): """Iterate over all the images in the eval dataset and get the average. + Args: + step: current training step + output_path: optional path to save rendered images to + Returns: metrics_dict: dictionary of metrics """ @@ -363,7 +374,15 @@ class VanillaPipeline(Pipeline): height, width = camera_ray_bundle.shape num_rays = height * width outputs = self.model.get_outputs_for_camera_ray_bundle(camera_ray_bundle) - metrics_dict, _ = self.model.get_image_metrics_and_images(outputs, batch) + metrics_dict, images_dict = self.model.get_image_metrics_and_images(outputs, batch) + + if output_path is not None: + camera_indices = camera_ray_bundle.camera_indices + assert camera_indices is not None + for key, val in images_dict.items(): + Image.fromarray((val * 255).byte().cpu().numpy()).save( + output_path / "{0:06d}-{1}.jpg".format(int(camera_indices[0, 0, 0]), key) + ) assert "num_rays_per_sec" not in metrics_dict metrics_dict["num_rays_per_sec"] = num_rays / (time() - inner_start) fps_str = "fps" diff --git a/nerfstudio/scripts/eval.py b/nerfstudio/scripts/eval.py index ef47d7a24..8f5995117 100644 --- a/nerfstudio/scripts/eval.py +++ b/nerfstudio/scripts/eval.py @@ -21,6 +21,7 @@ from __future__ import annotations import json from dataclasses import dataclass from pathlib import Path +from typing import Optional import tyro @@ -36,12 +37,16 @@ class ComputePSNR: load_config: Path # Name of the output file. output_path: Path = Path("output.json") + # Optional path to save rendered outputs to. + render_output_path: Optional[Path] = None def main(self) -> None: """Main function.""" config, pipeline, checkpoint_path, _ = eval_setup(self.load_config) assert self.output_path.suffix == ".json" - metrics_dict = pipeline.get_average_eval_image_metrics() + if self.render_output_path is not None: + self.render_output_path.mkdir(parents=True) + metrics_dict = pipeline.get_average_eval_image_metrics(output_path=self.render_output_path) self.output_path.parent.mkdir(parents=True, exist_ok=True) # Get the output and define the names to save to benchmark_info = {