Blend uniform background color for metric computation

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"],

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