Model outputs returns unblended RGB (black bg) for random bg

nerfstudio/model_components/renderers.py

@@ -75,7 +75,8 @@ class RGBRenderer(nn.Module):
         ray_indices: Optional[Int[Tensor, "num_samples"]] = None,
         num_rays: Optional[int] = None,
     ) -> Float[Tensor, "*bs 3"]:
-        """Composite samples along ray and render color image
+        """Composite samples along ray and render color image.
+        If background color is random, the predicted color is blended with black (0, 0, 0).
 
         Args:
             rgb: RGB for each sample
@@ -101,9 +102,12 @@ class RGBRenderer(nn.Module):
             comp_rgb = torch.sum(weights * rgb, dim=-2)
             accumulated_weight = torch.sum(weights, dim=-2)
 
+        if background_color == "random":
+            # If background color is random, the predicted color is blended with black (0, 0, 0).
+            return comp_rgb
+
         background_color = cls.get_background_color(rgb=rgb, background_color=background_color)
         comp_rgb = comp_rgb + background_color.to(weights.device) * (1.0 - accumulated_weight)
-
         return comp_rgb
 
     @classmethod
@@ -124,65 +128,68 @@ class RGBRenderer(nn.Module):
         elif background_color == "last_sample":
             background_color = rgb[..., -1, :]
         elif background_color == "random":
-            background_color = torch.rand_like(rgb[..., -1, :]).to(rgb.device)
+            # background_color = torch.rand_like(rgb[..., -1, :]).to(rgb.device)
+            raise RuntimeError("get_background_color should never be called with background_color='random'")
         elif isinstance(background_color, str) and background_color in colors.COLORS_DICT:
             background_color = colors.COLORS_DICT[background_color].to(rgb.device)
         assert isinstance(background_color, Tensor)
         return background_color
 
-    @classmethod
     def blend_background(
-        cls,
-        rgb: Tensor,
-        opacity: Tensor,
-        background_color: Float[Tensor, "3"],
+        self,
+        image: Tensor,
+        background_color: Optional[Float[Tensor, "3"]] = None,
     ) -> Float[Tensor, "*bs 3"]:
-        """Blends the background color into the image.
+        """Blends the background color into the image if image is RGBA.
+        Otherwise no blending is performed (we assume opacity of 1).
 
         Args:
-            rgb: RGB per pixel.
+            image: RGB/RGBA per pixel.
             opacity: Alpha opacity per pixel.
             background_color: Background color.
 
         Returns:
             Blended RGB.
         """
-        if len(opacity.shape) < len(rgb.shape):
-            opacity = opacity.unsqueeze(-1)
+        if background_color is None:
+            background_color_name = self.background_color
+            if background_color_name == "last_sample" or background_color == "random":
+                background_color_name = "black"
+            background_color = self.get_background_color(
+                rgb=image.unsqueeze(-1), background_color=background_color_name
+            )
+        if image.size(-1) < 4:
+            return image
+        rgb, opacity = image[..., :3], image[..., 3:]
         blended_rgb = rgb * opacity + background_color.to(rgb.device) * (1 - opacity)
         return blended_rgb
 
-    def blend_background_for_metric_computation(
+    def blend_background_for_loss_computation(
         self,
+        pred_image: Tensor,
+        pred_accumulation: Tensor,
         gt_image: Tensor,
-        predicted_rgb: Tensor,
-        predicted_opacity: Tensor,
-        background_color: Float[Tensor, "3"] = colors.BLACK,
     ) -> Tuple[Tensor, Tensor]:
         """Blends a background color into the ground truth and predicted image for
-        metric computation if the current background color is "random".
+        loss computation.
 
         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.
-
+            pred_image: The predicted RGB values (without background blending).
+            pred_accumulation: The predicted opacity/ accumulation.
         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
+        if self.background_color == "random":
+            background_color = torch.rand_like(pred_image)
+        else:
+            background_color = self.get_background_color(
+                rgb=pred_image.unsqueeze(-1), background_color=self.background_color
+            ).to(pred_image.device)
+        pred_image = pred_image + background_color * (1.0 - pred_accumulation)
+        gt_image = self.blend_background(gt_image, background_color=background_color)
+        return pred_image, gt_image
 
     def forward(
         self,

nerfstudio/models/nerfacto.py

@@ -283,16 +283,7 @@ class NerfactoModel(Model):
         weights_list.append(weights)
         ray_samples_list.append(ray_samples)
 
-        background_color = (
-            self.renderer_rgb.get_background_color(
-                rgb=field_outputs[FieldHeadNames.RGB], background_color=self.config.background_color
-            )
-            if self.training
-            else None
-        )
-        rgb = self.renderer_rgb(
-            rgb=field_outputs[FieldHeadNames.RGB], weights=weights, background_color=background_color
-        )
+        rgb = self.renderer_rgb(rgb=field_outputs[FieldHeadNames.RGB], weights=weights)
         depth = self.renderer_depth(weights=weights, ray_samples=ray_samples)
         accumulation = self.renderer_accumulation(weights=weights)
 
@@ -301,8 +292,6 @@ class NerfactoModel(Model):
             "accumulation": accumulation,
             "depth": depth,
         }
-        if self.training and self.config.background_color == "random" and background_color is not None:
-            outputs["background_color"] = background_color
 
         if self.config.predict_normals:
             normals = self.renderer_normals(normals=field_outputs[FieldHeadNames.NORMALS], weights=weights)
@@ -332,10 +321,9 @@ class NerfactoModel(Model):
 
     def get_metrics_dict(self, outputs, batch):
         metrics_dict = {}
-        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"]
-        )
+        gt_rgb = batch["image"].to(self.device)  # RGB or RGBA image
+        gt_rgb = self.renderer_rgb.blend_background(gt_rgb)  # Blend if RGBA
+        predicted_rgb = outputs["rgb"]
         metrics_dict["psnr"] = self.psnr(predicted_rgb, gt_rgb)
 
         if self.training:
@@ -345,14 +333,13 @@ class NerfactoModel(Model):
     def get_loss_dict(self, outputs, batch, metrics_dict=None):
         loss_dict = {}
         image = batch["image"].to(self.device)
-        blended_rgb = (
-            self.renderer_rgb.blend_background(
-                rgb=image[..., :3], opacity=image[..., 3], background_color=outputs["background_color"]
-            )
-            if image.shape[-1] == 4 and "background_color" in outputs
-            else image[..., :3]
+        pred_rgb, gt_rgb = self.renderer_rgb.blend_background_for_loss_computation(
+            pred_image=outputs["rgb"],
+            pred_accumulation=outputs["accumulation"],
+            gt_image=image,
         )
-        loss_dict["rgb_loss"] = self.rgb_loss(blended_rgb, outputs["rgb"])
+
+        loss_dict["rgb_loss"] = self.rgb_loss(gt_rgb, pred_rgb)
         if self.training:
             loss_dict["interlevel_loss"] = self.config.interlevel_loss_mult * interlevel_loss(
                 outputs["weights_list"], outputs["ray_samples_list"]
@@ -374,10 +361,9 @@ 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"].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"]
-        )
+        gt_rgb = batch["image"].to(self.device)
+        predicted_rgb = outputs["rgb"]  # Blended with background (black if random background)
+        gt_rgb = self.renderer_rgb.blend_background(gt_rgb)
         acc = colormaps.apply_colormap(outputs["accumulation"])
         depth = colormaps.apply_depth_colormap(
             outputs["depth"],

tests/field_components/test_mlp.py

@@ -4,7 +4,7 @@ MLP Test
 import torch
 from torch import nn
 
-from nerfstudio.field_components import MLP
+from nerfstudio.field_components.mlp import MLP
 
 
 def test_mlp():