From 324f1d49491c629ce3e0d96d4869cd5cc0b99924 Mon Sep 17 00:00:00 2001 From: Nicholas Pfaff Date: Wed, 5 Jul 2023 08:35:42 +0100 Subject: [PATCH] Blend uniform background color for metric computation --- nerfstudio/model_components/renderers.py | 40 +++++++++++++++++++++--- nerfstudio/models/nerfacto.py | 28 ++++++++++------- 2 files changed, 52 insertions(+), 16 deletions(-) diff --git a/nerfstudio/model_components/renderers.py b/nerfstudio/model_components/renderers.py index 431e2d4f4..ac7a19e5f 100644 --- a/nerfstudio/model_components/renderers.py +++ b/nerfstudio/model_components/renderers.py @@ -28,7 +28,7 @@ Example: """ import contextlib import math -from typing import Generator, Literal, Optional, Union +from typing import Generator, Literal, Optional, Tuple, Union import nerfacc import torch @@ -107,9 +107,7 @@ class RGBRenderer(nn.Module): return comp_rgb @classmethod - def get_background_color( - cls, rgb: Float[Tensor, "*bs num_samples 3"], background_color: BackgroundColor - ) -> Float[Tensor, "3"]: + def get_background_color(cls, rgb: Float[Tensor, "*bs num_samples 3"], background_color: BackgroundColor) -> Tensor: """Returns the RGB background color for a specified background color. Args: @@ -147,11 +145,43 @@ class RGBRenderer(nn.Module): Returns: Blended RGB. """ - if len(opacity.shape) == 1: + if len(opacity.shape) < len(rgb.shape): opacity = opacity.unsqueeze(-1) blended_rgb = rgb * opacity + background_color.to(rgb.device) * (1 - opacity) return blended_rgb + def blend_background_for_metric_computation( + self, + gt_image: torch.Tensor, + predicted_rgb: torch.Tensor, + predicted_opacity: torch.Tensor, + background_color: torch.Tensor = colors.BLACK, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Blends a background color into the ground truth and predicted image for + metric computation if the current background color is "random". + + Args: + gt_image: The ground truth image. + predicted_rgb: The predicted RGB values. + predicted_opacity: The predicted opacity/ accumulation. + background_color: The background color to use for blending. This should be + a single color to use for all pixels. + + Returns: + A tuple of the predicted and ground truth RGB values. + """ + if gt_image.shape[-1] == 4 and self.background_color == "random": + predicted_rgb = self.blend_background( + rgb=predicted_rgb, opacity=predicted_opacity, background_color=background_color + ) + gt_rgb = self.blend_background( + rgb=gt_image[..., :3], opacity=gt_image[..., 3], background_color=background_color + ) + else: + gt_rgb = gt_image[..., :3] + + return predicted_rgb, gt_rgb + def forward( self, rgb: Float[Tensor, "*bs num_samples 3"], diff --git a/nerfstudio/models/nerfacto.py b/nerfstudio/models/nerfacto.py index 3e7f8051e..289f992fc 100644 --- a/nerfstudio/models/nerfacto.py +++ b/nerfstudio/models/nerfacto.py @@ -332,8 +332,12 @@ class NerfactoModel(Model): def get_metrics_dict(self, outputs, batch): metrics_dict = {} - image = batch["image"][..., :3].to(self.device) - metrics_dict["psnr"] = self.psnr(outputs["rgb"], image) + image = batch["image"].to(self.device) + predicted_rgb, gt_rgb = self.renderer_rgb.blend_background_for_metric_computation( + gt_image=image, predicted_rgb=outputs["rgb"], predicted_opacity=outputs["accumulation"] + ) + metrics_dict["psnr"] = self.psnr(predicted_rgb, gt_rgb) + if self.training: metrics_dict["distortion"] = distortion_loss(outputs["weights_list"], outputs["ray_samples_list"]) return metrics_dict @@ -345,7 +349,7 @@ class NerfactoModel(Model): self.renderer_rgb.blend_background( rgb=image[..., :3], opacity=image[..., 3], background_color=outputs["background_color"] ) - if "background_color" in outputs and image.shape[-1] == 4 + if image.shape[-1] == 4 and "background_color" in outputs else image[..., :3] ) loss_dict["rgb_loss"] = self.rgb_loss(blended_rgb, outputs["rgb"]) @@ -370,25 +374,27 @@ class NerfactoModel(Model): def get_image_metrics_and_images( self, outputs: Dict[str, torch.Tensor], batch: Dict[str, torch.Tensor] ) -> Tuple[Dict[str, float], Dict[str, torch.Tensor]]: - image = batch["image"][..., :3].to(self.device) - rgb = outputs["rgb"] + image = batch["image"].to(self.device) + predicted_rgb, gt_rgb = self.renderer_rgb.blend_background_for_metric_computation( + gt_image=image, predicted_rgb=outputs["rgb"], predicted_opacity=outputs["accumulation"] + ) acc = colormaps.apply_colormap(outputs["accumulation"]) depth = colormaps.apply_depth_colormap( outputs["depth"], accumulation=outputs["accumulation"], ) - combined_rgb = torch.cat([image, rgb], dim=1) + combined_rgb = torch.cat([gt_rgb, predicted_rgb], dim=1) combined_acc = torch.cat([acc], dim=1) combined_depth = torch.cat([depth], dim=1) # Switch images from [H, W, C] to [1, C, H, W] for metrics computations - image = torch.moveaxis(image, -1, 0)[None, ...] - rgb = torch.moveaxis(rgb, -1, 0)[None, ...] + gt_rgb = torch.moveaxis(gt_rgb, -1, 0)[None, ...] + predicted_rgb = torch.moveaxis(predicted_rgb, -1, 0)[None, ...] - psnr = self.psnr(image, rgb) - ssim = self.ssim(image, rgb) - lpips = self.lpips(image, rgb) + psnr = self.psnr(gt_rgb, predicted_rgb) + ssim = self.ssim(gt_rgb, predicted_rgb) + lpips = self.lpips(gt_rgb, predicted_rgb) # all of these metrics will be logged as scalars metrics_dict = {"psnr": float(psnr.item()), "ssim": float(ssim)} # type: ignore