From dbd4092dbb68a7ff9cb9667cdbcc56dec53bd2e8 Mon Sep 17 00:00:00 2001 From: Khoa Nguyen-Tuan <12138158+Khoa-NT@users.noreply.github.com> Date: Tue, 18 Jul 2023 07:02:32 +0900 Subject: [PATCH] Add std for eval (#2239) * Add std for eval * change default value of get_std --------- Co-authored-by: khoa --- nerfstudio/pipelines/base_pipeline.py | 23 ++++++++++++++++++----- nerfstudio/scripts/eval.py | 2 +- 2 files changed, 19 insertions(+), 6 deletions(-) diff --git a/nerfstudio/pipelines/base_pipeline.py b/nerfstudio/pipelines/base_pipeline.py index 11ba841e0..5e202b06c 100644 --- a/nerfstudio/pipelines/base_pipeline.py +++ b/nerfstudio/pipelines/base_pipeline.py @@ -183,12 +183,15 @@ class Pipeline(nn.Module): @abstractmethod @profiler.time_function - def get_average_eval_image_metrics(self, step: Optional[int] = None, output_path: Optional[Path] = None): + def get_average_eval_image_metrics( + self, step: Optional[int] = None, output_path: Optional[Path] = None, get_std: bool = False + ): """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 + get_std: Set True if you want to return std with the mean metric. """ def load_pipeline(self, loaded_state: Dict[str, Any], step: int) -> None: @@ -354,12 +357,15 @@ class VanillaPipeline(Pipeline): return metrics_dict, images_dict @profiler.time_function - def get_average_eval_image_metrics(self, step: Optional[int] = None, output_path: Optional[Path] = None): + def get_average_eval_image_metrics( + self, step: Optional[int] = None, output_path: Optional[Path] = None, get_std: bool = False + ): """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 + get_std: Set True if you want to return std with the mean metric. Returns: metrics_dict: dictionary of metrics @@ -401,9 +407,16 @@ class VanillaPipeline(Pipeline): # average the metrics list metrics_dict = {} for key in metrics_dict_list[0].keys(): - metrics_dict[key] = float( - torch.mean(torch.tensor([metrics_dict[key] for metrics_dict in metrics_dict_list])) - ) + if get_std: + key_std, key_mean = torch.std_mean( + torch.tensor([metrics_dict[key] for metrics_dict in metrics_dict_list]) + ) + metrics_dict[key] = float(key_mean) + metrics_dict[f"{key}_std"] = float(key_std) + else: + metrics_dict[key] = float( + torch.mean(torch.tensor([metrics_dict[key] for metrics_dict in metrics_dict_list])) + ) self.train() return metrics_dict diff --git a/nerfstudio/scripts/eval.py b/nerfstudio/scripts/eval.py index 8f5995117..751b318a2 100644 --- a/nerfstudio/scripts/eval.py +++ b/nerfstudio/scripts/eval.py @@ -46,7 +46,7 @@ class ComputePSNR: assert self.output_path.suffix == ".json" 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) + metrics_dict = pipeline.get_average_eval_image_metrics(output_path=self.render_output_path, get_std=True) self.output_path.parent.mkdir(parents=True, exist_ok=True) # Get the output and define the names to save to benchmark_info = {