Minor refac

nerfstudio/model_components/renderers.py

@@ -107,13 +107,18 @@ class RGBRenderer(nn.Module):
             # as if the background color was black.
             return comp_rgb
 
-        background_color = cls.get_background_color(rgb=rgb, background_color=background_color)
+        elif background_color == "last_sample":
+            # Note, this is only supported for non-packed samples.
+            background_color = rgb[..., -1, :]
+        else:
+            background_color = cls.get_background_color(background_color, shape=rgb.shape, device=rgb.device)
+        assert isinstance(background_color, torch.Tensor)
         comp_rgb = comp_rgb + background_color.to(weights.device) * (1.0 - accumulated_weight)
         return comp_rgb
 
     @classmethod
     def get_background_color(
-        cls, rgb: Float[Tensor, "*bs num_samples 3"], background_color: BackgroundColor
+        cls, background_color: BackgroundColor, shape: Tuple[int, ...], device: torch.device
     ) -> Union[Float[Tensor, "3"], Float[Tensor, "*bs 3"]]:
         """Returns the RGB background color for a specified background color.
 
@@ -124,16 +129,16 @@ class RGBRenderer(nn.Module):
         Returns:
             Background color as RGB.
         """
+        assert background_color not in {"last_sample", "random"}
+        assert shape[-1] == 3, "Background color must be RGB."
         if BACKGROUND_COLOR_OVERRIDE is not None:
             background_color = BACKGROUND_COLOR_OVERRIDE
-        elif background_color == "last_sample":
-            background_color = rgb[..., -1, :]
-        elif background_color == "random":
-            background_color = torch.rand_like(rgb[..., -1, :]).to(rgb.device)
-        elif isinstance(background_color, str) and background_color in colors.COLORS_DICT:
-            background_color = colors.COLORS_DICT[background_color].to(rgb.device)
+        if isinstance(background_color, str) and background_color in colors.COLORS_DICT:
+            background_color = colors.COLORS_DICT[background_color]
         assert isinstance(background_color, Tensor)
-        return background_color
+
+        # Ensure correct shape
+        return background_color.expand(shape).to(device)
 
     def blend_background(
         self,
@@ -154,16 +159,14 @@ class RGBRenderer(nn.Module):
         if image.size(-1) < 4:
             return image
 
-        if background_color is None:
-            background_color_name = self.background_color
-            if background_color_name == "last_sample" or background_color_name == "random":
-                background_color_name = "black"
-            background_color = self.get_background_color(
-                rgb=image.unsqueeze(-1), background_color=background_color_name
-            )
         rgb, opacity = image[..., :3], image[..., 3:]
-        blended_rgb = rgb * opacity + background_color.to(rgb.device) * (1 - opacity)
-        return blended_rgb
+        if background_color is None:
+            background_color = self.background_color
+            if background_color in {"last_sample", "random"}:
+                background_color = "black"
+        background_color = self.get_background_color(background_color, shape=rgb.shape, device=rgb.device)
+        assert isinstance(background_color, torch.Tensor)
+        return rgb * opacity + background_color.to(rgb.device) * (1 - opacity)
 
     def blend_background_for_loss_computation(
         self,
@@ -181,11 +184,11 @@ class RGBRenderer(nn.Module):
         Returns:
             A tuple of the predicted and ground truth RGB values.
         """
-
-        background_color = self.get_background_color(
-            rgb=pred_image.unsqueeze(-1), background_color=self.background_color
-        ).to(pred_image.device)
-        if self.background_color == "random":
+        background_color = self.background_color
+        if self.background_color == "last_sample":
+            background_color = "black"  # No background blending for GT
+        elif self.background_color == "random":
+            background_color = torch.rand_like(pred_image)
             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

nerfstudio/models/mipnerf.py

@@ -144,12 +144,12 @@ class MipNerfModel(Model):
         image = batch["image"].to(self.device)
         pred_coarse, image_coarse = self.renderer_rgb.blend_background_for_loss_computation(
             pred_image=outputs["rgb_coarse"],
-            pred_accumulation=outputs["accumulation"],
+            pred_accumulation=outputs["accumulation_coarse"],
             gt_image=image,
         )
         pred_fine, image_fine = self.renderer_rgb.blend_background_for_loss_computation(
             pred_image=outputs["rgb_fine"],
-            pred_accumulation=outputs["accumulation"],
+            pred_accumulation=outputs["accumulation_fine"],
             gt_image=image,
         )
         rgb_loss_coarse = self.rgb_loss(image_coarse, pred_coarse)

nerfstudio/models/vanilla_nerf.py

@@ -194,12 +194,12 @@ class NeRFModel(Model):
         image = batch["image"].to(device)
         coarse_pred, coarse_image = self.renderer_rgb.blend_background_for_loss_computation(
             pred_image=outputs["rgb_coarse"],
-            pred_accumulation=outputs["accumulation"],
+            pred_accumulation=outputs["accumulation_coarse"],
             gt_image=image,
         )
         fine_pred, fine_image = self.renderer_rgb.blend_background_for_loss_computation(
             pred_image=outputs["rgb_fine"],
-            pred_accumulation=outputs["accumulation"],
+            pred_accumulation=outputs["accumulation_fine"],
             gt_image=image,
         )