Implement nerfacto alpha transparency training (#2165)

* Implement nerfacto alpha transparency training

* Stop adding all ones alpha channel and clean up blend_background

* Make background_color a required argument in blend_background

* Enable alpha transparent training by default

* Improve background_color typing

* Blend uniform background color for metric computation

* Fix typing error

* Clean up typing

* Fix background_color to outputs assignment

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

* Update renderers

* Fix comments

* Add support for transparent BG for other methods

* Minor refac

* Fix get_background_color

* Fix typing error for blend_background

* Fix comments

* Export utils support transparency

* Pylint fix

* Fix pcd export masking

* Filter pcd view directions based on opacity

* Clean up get_rgba_image

---------

Co-authored-by: Jonas Kulhanek <jonas.kulhanek@live.com>

nerfstudio/data/datasets/base_dataset.py

@@ -85,10 +85,7 @@ class InputDataset(Dataset):
         """
         image = torch.from_numpy(self.get_numpy_image(image_idx).astype("float32") / 255.0)
         if self._dataparser_outputs.alpha_color is not None and image.shape[-1] == 4:
-            assert image.shape[-1] == 4
             image = image[:, :, :3] * image[:, :, -1:] + self._dataparser_outputs.alpha_color * (1.0 - image[:, :, -1:])
-        else:
-            image = image[:, :, :3]
         return image
 
     def get_data(self, image_idx: int) -> Dict:

nerfstudio/exporter/exporter_utils.py

@@ -28,13 +28,7 @@ import open3d as o3d
 import pymeshlab
 import torch
 from jaxtyping import Float
-from rich.progress import (
-    BarColumn,
-    Progress,
-    TaskProgressColumn,
-    TextColumn,
-    TimeRemainingColumn,
-)
+from rich.progress import BarColumn, Progress, TaskProgressColumn, TextColumn, TimeRemainingColumn
 from torch import Tensor
 
 from nerfstudio.cameras.cameras import Cameras
@@ -144,7 +138,7 @@ def generate_point_cloud(
                 CONSOLE.print(f"Could not find {depth_output_name} in the model outputs", justify="center")
                 CONSOLE.print(f"Please set --depth_output_name to one of: {outputs.keys()}", justify="center")
                 sys.exit(1)
-            rgb = outputs[rgb_output_name]
+            rgba = pipeline.model.get_rgba_image(outputs, rgb_output_name)
             depth = outputs[depth_output_name]
             if normal_output_name is not None:
                 if normal_output_name not in outputs:
@@ -160,6 +154,14 @@ def generate_point_cloud(
             point = ray_bundle.origins + ray_bundle.directions * depth
             view_direction = ray_bundle.directions
 
+            # Filter points with opacity lower than 0.5
+            mask = rgba[..., -1] > 0.5
+            point = point[mask]
+            view_direction = view_direction[mask]
+            rgb = rgba[mask][..., :3]
+            if normal is not None:
+                normal = normal[mask]
+
             if use_bounding_box:
                 comp_l = torch.tensor(bounding_box_min, device=point.device)
                 comp_m = torch.tensor(bounding_box_max, device=point.device)
@@ -230,6 +232,7 @@ def render_trajectory(
     depth_output_name: str,
     rendered_resolution_scaling_factor: float = 1.0,
     disable_distortion: bool = False,
+    return_rgba_images: bool = False,
 ) -> Tuple[List[np.ndarray], List[np.ndarray]]:
     """Helper function to create a video of a trajectory.
 
@@ -240,6 +243,7 @@ def render_trajectory(
         depth_output_name: Name of the depth output.
         rendered_resolution_scaling_factor: Scaling factor to apply to the camera image resolution.
         disable_distortion: Whether to disable distortion.
+        return_rgba_images: Whether to return RGBA images (default RGB).
 
     Returns:
         List of rgb images, list of depth images.
@@ -272,7 +276,11 @@ def render_trajectory(
                 CONSOLE.print(f"Could not find {depth_output_name} in the model outputs", justify="center")
                 CONSOLE.print(f"Please set --depth_output_name to one of: {outputs.keys()}", justify="center")
                 sys.exit(1)
-            images.append(outputs[rgb_output_name].cpu().numpy())
+            if return_rgba_images:
+                image = pipeline.model.get_rgba_image(outputs, rgb_output_name)
+            else:
+                image = outputs[rgb_output_name]
+            images.append(image.cpu().numpy())
             depths.append(outputs[depth_output_name].cpu().numpy())
     return images, depths
 

nerfstudio/exporter/texture_utils.py

@@ -403,9 +403,11 @@ def export_textured_mesh(
     CONSOLE.print("Creating texture image by rendering with NeRF...")
     with torch.no_grad():
         outputs = pipeline.model.get_outputs_for_camera_ray_bundle(camera_ray_bundle)
+    # TODO: this can be done better by using the alpha channel
+    rgb = pipeline.model.get_rgba_image(outputs, "rgb")[..., :3]
 
     # save the texture image
-    texture_image = outputs["rgb"].cpu().numpy()
+    texture_image = rgb.cpu().numpy()
     media.write_image(str(output_dir / "material_0.png"), texture_image)
 
     CONSOLE.print("Writing relevant OBJ information to files...")

nerfstudio/exporter/tsdf_utils.py

@@ -331,8 +331,12 @@ def export_tsdf_mesh(
         depth_output_name=depth_output_name,
         rendered_resolution_scaling_factor=1.0 / downscale_factor,
         disable_distortion=True,
+        return_rgba_images=True,
     )
 
+    # TODO: this can be done better by removing transparent points from the TSDF
+    color_images = [x[..., :3] for x in color_images]
+
     # camera extrinsics and intrinsics
     c2w: Float[Tensor, "N 3 4"] = cameras.camera_to_worlds.to(device)
     # make c2w homogeneous

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
@@ -39,7 +39,7 @@ from nerfstudio.cameras.rays import RaySamples
 from nerfstudio.utils import colors
 from nerfstudio.utils.math import components_from_spherical_harmonics, safe_normalize
 
-BackgroundColor = Union[Literal["random", "last_sample", "black", "white"], Float[Tensor, "3"]]
+BackgroundColor = Union[Literal["random", "last_sample", "black", "white"], Float[Tensor, "3"], Float[Tensor, "*bs 3"]]
 BACKGROUND_COLOR_OVERRIDE: Optional[Float[Tensor, "3"]] = None
 
 
@@ -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, no BG color is added - as if the background was black!
 
         Args:
             rgb: RGB for each sample
@@ -101,19 +102,99 @@ 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 returned without blending,
+            # as if the background color was black.
+            return comp_rgb
+
+        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=comp_rgb.shape, device=comp_rgb.device)
+        assert isinstance(background_color, torch.Tensor)
+        comp_rgb = comp_rgb + background_color * (1.0 - accumulated_weight)
+        return comp_rgb
+
+    @classmethod
+    def get_background_color(
+        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.
+
+        Note:
+            This function CANNOT be called for background_color being either "last_sample" or "random".
+
+        Args:
+            rgb: RGB for each sample.
+            background_color: The background color specification.
+
+        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
-        if background_color == "last_sample":
-            background_color = rgb[..., -1, :]
-        if background_color == "random":
-            background_color = torch.rand_like(comp_rgb).to(rgb.device)
         if isinstance(background_color, str) and background_color in colors.COLORS_DICT:
-            background_color = colors.COLORS_DICT[background_color].to(rgb.device)
+            background_color = colors.COLORS_DICT[background_color]
+        assert isinstance(background_color, Tensor)
 
+        # Ensure correct shape
+        return background_color.expand(shape).to(device)
+
+    def blend_background(
+        self,
+        image: Tensor,
+        background_color: Optional[BackgroundColor] = None,
+    ) -> Float[Tensor, "*bs 3"]:
+        """Blends the background color into the image if image is RGBA.
+        Otherwise no blending is performed (we assume opacity of 1).
+
+        Args:
+            image: RGB/RGBA per pixel.
+            opacity: Alpha opacity per pixel.
+            background_color: Background color.
+
+        Returns:
+            Blended RGB.
+        """
+        if image.size(-1) < 4:
+            return image
+
+        rgb, opacity = image[..., :3], image[..., 3:]
+        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)
-        comp_rgb = comp_rgb + background_color.to(weights.device) * (1.0 - accumulated_weight)
+        return rgb * opacity + background_color.to(rgb.device) * (1 - opacity)
 
-        return comp_rgb
+    def blend_background_for_loss_computation(
+        self,
+        pred_image: Tensor,
+        pred_accumulation: Tensor,
+        gt_image: Tensor,
+    ) -> Tuple[Tensor, Tensor]:
+        """Blends a background color into the ground truth and predicted image for
+        loss computation.
+
+        Args:
+            gt_image: The ground truth image.
+            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.
+        """
+        background_color = self.background_color
+        if background_color == "last_sample":
+            background_color = "black"  # No background blending for GT
+        elif 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
 
     def forward(
         self,
@@ -121,6 +202,7 @@ class RGBRenderer(nn.Module):
         weights: Float[Tensor, "*bs num_samples 1"],
         ray_indices: Optional[Int[Tensor, "num_samples"]] = None,
         num_rays: Optional[int] = None,
+        background_color: Optional[BackgroundColor] = None,
     ) -> Float[Tensor, "*bs 3"]:
         """Composite samples along ray and render color image
 
@@ -129,15 +211,18 @@ class RGBRenderer(nn.Module):
             weights: Weights for each sample
             ray_indices: Ray index for each sample, used when samples are packed.
             num_rays: Number of rays, used when samples are packed.
+            background_color: The background color to use for rendering.
 
         Returns:
             Outputs of rgb values.
         """
+        if background_color is None:
+            background_color = self.background_color
 
         if not self.training:
             rgb = torch.nan_to_num(rgb)
         rgb = self.combine_rgb(
-            rgb, weights, background_color=self.background_color, ray_indices=ray_indices, num_rays=num_rays
+            rgb, weights, background_color=background_color, ray_indices=ray_indices, num_rays=num_rays
         )
         if not self.training:
             torch.clamp_(rgb, min=0.0, max=1.0)

nerfstudio/models/base_model.py

@@ -187,6 +187,30 @@ class Model(nn.Module):
             outputs[output_name] = torch.cat(outputs_list).view(image_height, image_width, -1)  # type: ignore
         return outputs
 
+    def get_rgba_image(self, outputs: Dict[str, torch.Tensor], output_name: str = "rgb") -> torch.Tensor:
+        """Returns the RGBA image from the outputs of the model.
+
+        Args:
+            outputs: Outputs of the model.
+
+        Returns:
+            RGBA image.
+        """
+        accumulation_name = output_name.replace("rgb", "accumulation")
+        if (
+            not hasattr(self, "renderer_rgb")
+            or not hasattr(self.renderer_rgb, "background_color")
+            or accumulation_name not in outputs
+        ):
+            raise NotImplementedError(f"get_rgba_image is not implemented for model {self.__class__.__name__}")
+        rgb = outputs[output_name]
+        if self.renderer_rgb.background_color == "random":  # type: ignore
+            acc = outputs[accumulation_name]
+            if acc.dim() < rgb.dim():
+                acc = acc.unsqueeze(-1)
+            return torch.cat((rgb / acc.clamp(min=1e-10), acc), dim=-1)
+        return torch.cat((rgb, torch.ones_like(rgb[..., :1])), dim=-1)
+
     @abstractmethod
     def get_image_metrics_and_images(
         self, outputs: Dict[str, torch.Tensor], batch: Dict[str, torch.Tensor]

nerfstudio/models/base_surface_model.py

@@ -293,7 +293,12 @@ class SurfaceModel(Model):
         """
         loss_dict = {}
         image = batch["image"].to(self.device)
-        loss_dict["rgb_loss"] = self.rgb_loss(image, outputs["rgb"])
+        pred_image, image = 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(image, pred_image)
         if self.training:
             # eikonal loss
             grad_theta = outputs["eik_grad"]
@@ -337,6 +342,7 @@ class SurfaceModel(Model):
         """
         metrics_dict = {}
         image = batch["image"].to(self.device)
+        image = self.renderer_rgb.blend_background(image)
         metrics_dict["psnr"] = self.psnr(outputs["rgb"], image)
         return metrics_dict
 
@@ -352,6 +358,7 @@ class SurfaceModel(Model):
             A dictionary of metrics.
         """
         image = batch["image"].to(self.device)
+        image = self.renderer_rgb.blend_background(image)
         rgb = outputs["rgb"]
         acc = colormaps.apply_colormap(outputs["accumulation"])
 

nerfstudio/models/instant_ngp.py

@@ -219,14 +219,20 @@ class NGPModel(Model):
 
     def get_metrics_dict(self, outputs, batch):
         image = batch["image"].to(self.device)
+        image = self.renderer_rgb.blend_background(image)
         metrics_dict = {}
         metrics_dict["psnr"] = self.psnr(outputs["rgb"], image)
         metrics_dict["num_samples_per_batch"] = outputs["num_samples_per_ray"].sum()
         return metrics_dict
 
     def get_loss_dict(self, outputs, batch, metrics_dict=None):
-        image = batch["image"].to(self.device)
-        rgb_loss = self.rgb_loss(image, outputs["rgb"])
+        image = batch["image"][..., :3].to(self.device)
+        pred_rgb, image = self.renderer_rgb.blend_background_for_loss_computation(
+            pred_image=outputs["rgb"],
+            pred_accumulation=outputs["accumulation"],
+            gt_image=image,
+        )
+        rgb_loss = self.rgb_loss(image, pred_rgb)
         loss_dict = {"rgb_loss": rgb_loss}
         return loss_dict
 
@@ -234,6 +240,7 @@ class NGPModel(Model):
         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)
+        image = self.renderer_rgb.blend_background(image)
         rgb = outputs["rgb"]
         acc = colormaps.apply_colormap(outputs["accumulation"])
         depth = colormaps.apply_depth_colormap(

nerfstudio/models/mipnerf.py

@@ -142,8 +142,18 @@ class MipNerfModel(Model):
 
     def get_loss_dict(self, outputs, batch, metrics_dict=None):
         image = batch["image"].to(self.device)
-        rgb_loss_coarse = self.rgb_loss(image, outputs["rgb_coarse"])
-        rgb_loss_fine = self.rgb_loss(image, outputs["rgb_fine"])
+        pred_coarse, image_coarse = self.renderer_rgb.blend_background_for_loss_computation(
+            pred_image=outputs["rgb_coarse"],
+            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_fine"],
+            gt_image=image,
+        )
+        rgb_loss_coarse = self.rgb_loss(image_coarse, pred_coarse)
+        rgb_loss_fine = self.rgb_loss(image_fine, pred_fine)
         loss_dict = {"rgb_loss_coarse": rgb_loss_coarse, "rgb_loss_fine": rgb_loss_fine}
         loss_dict = misc.scale_dict(loss_dict, self.config.loss_coefficients)
         return loss_dict
@@ -153,6 +163,7 @@ class MipNerfModel(Model):
     ) -> Tuple[Dict[str, float], Dict[str, torch.Tensor]]:
         assert self.config.collider_params is not None, "mip-NeRF requires collider parameters to be set."
         image = batch["image"].to(outputs["rgb_coarse"].device)
+        image = self.renderer_rgb.blend_background(image)
         rgb_coarse = outputs["rgb_coarse"]
         rgb_fine = outputs["rgb_fine"]
         acc_coarse = colormaps.apply_colormap(outputs["accumulation_coarse"])

nerfstudio/models/nerfacto.py

@@ -323,8 +323,11 @@ class NerfactoModel(Model):
 
     def get_metrics_dict(self, outputs, batch):
         metrics_dict = {}
-        image = batch["image"].to(self.device)
-        metrics_dict["psnr"] = self.psnr(outputs["rgb"], image)
+        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:
             metrics_dict["distortion"] = distortion_loss(outputs["weights_list"], outputs["ray_samples_list"])
         return metrics_dict
@@ -332,7 +335,13 @@ class NerfactoModel(Model):
     def get_loss_dict(self, outputs, batch, metrics_dict=None):
         loss_dict = {}
         image = batch["image"].to(self.device)
-        loss_dict["rgb_loss"] = self.rgb_loss(image, outputs["rgb"])
+        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(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"]
@@ -354,25 +363,26 @@ 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)
-        rgb = outputs["rgb"]
+        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"],
             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

nerfstudio/models/semantic_nerfw.py

@@ -230,6 +230,7 @@ class SemanticNerfWModel(Model):
     def get_metrics_dict(self, outputs, batch):
         metrics_dict = {}
         image = batch["image"].to(self.device)
+        image = self.renderer_rgb.blend_background(image)
         metrics_dict["psnr"] = self.psnr(outputs["rgb"], image)
         metrics_dict["distortion"] = distortion_loss(outputs["weights_list"], outputs["ray_samples_list"])
         return metrics_dict
@@ -237,6 +238,7 @@ class SemanticNerfWModel(Model):
     def get_loss_dict(self, outputs, batch, metrics_dict=None):
         loss_dict = {}
         image = batch["image"].to(self.device)
+        image = self.renderer_rgb.blend_background(image)
         loss_dict["interlevel_loss"] = self.config.interlevel_loss_mult * interlevel_loss(
             outputs["weights_list"], outputs["ray_samples_list"]
         )
@@ -263,6 +265,12 @@ class SemanticNerfWModel(Model):
     ) -> Tuple[Dict[str, float], Dict[str, torch.Tensor]]:
         image = batch["image"].to(self.device)
         rgb = outputs["rgb"]
+        rgb, image = self.renderer_rgb.blend_background_for_loss_computation(
+            pred_image=rgb,
+            pred_accumulation=outputs["accumulation"],
+            gt_image=image,
+        )
+
         rgb = torch.clamp(rgb, min=0, max=1)
         acc = colormaps.apply_colormap(outputs["accumulation"])
         depth = colormaps.apply_depth_colormap(

nerfstudio/models/tensorf.py

@@ -303,9 +303,14 @@ class TensoRFModel(Model):
     def get_loss_dict(self, outputs, batch, metrics_dict=None) -> Dict[str, torch.Tensor]:
         # Scaling metrics by coefficients to create the losses.
         device = outputs["rgb"].device
-        image = batch["image"].to(device)
+        image = batch["image"][..., :3].to(device)
+        pred_image, image = self.renderer_rgb.blend_background_for_loss_computation(
+            pred_image=outputs["rgb"],
+            pred_accumulation=outputs["accumulation"],
+            gt_image=image,
+        )
 
-        rgb_loss = self.rgb_loss(image, outputs["rgb"])
+        rgb_loss = self.rgb_loss(image, pred_image)
 
         loss_dict = {"rgb_loss": rgb_loss}
 
@@ -334,6 +339,7 @@ class TensoRFModel(Model):
         self, outputs: Dict[str, torch.Tensor], batch: Dict[str, torch.Tensor]
     ) -> Tuple[Dict[str, float], Dict[str, torch.Tensor]]:
         image = batch["image"].to(outputs["rgb"].device)
+        image = self.renderer_rgb.blend_background(image)
         rgb = outputs["rgb"]
         acc = colormaps.apply_colormap(outputs["accumulation"])
         assert self.config.collider_params is not None

nerfstudio/models/vanilla_nerf.py

@@ -192,9 +192,19 @@ class NeRFModel(Model):
         # Scaling metrics by coefficients to create the losses.
         device = outputs["rgb_coarse"].device
         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_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_fine"],
+            gt_image=image,
+        )
 
-        rgb_loss_coarse = self.rgb_loss(image, outputs["rgb_coarse"])
-        rgb_loss_fine = self.rgb_loss(image, outputs["rgb_fine"])
+        rgb_loss_coarse = self.rgb_loss(coarse_image, coarse_pred)
+        rgb_loss_fine = self.rgb_loss(fine_image, fine_pred)
 
         loss_dict = {"rgb_loss_coarse": rgb_loss_coarse, "rgb_loss_fine": rgb_loss_fine}
         loss_dict = misc.scale_dict(loss_dict, self.config.loss_coefficients)
@@ -204,6 +214,7 @@ class NeRFModel(Model):
         self, outputs: Dict[str, torch.Tensor], batch: Dict[str, torch.Tensor]
     ) -> Tuple[Dict[str, float], Dict[str, torch.Tensor]]:
         image = batch["image"].to(outputs["rgb_coarse"].device)
+        image = self.renderer_rgb.blend_background(image)
         rgb_coarse = outputs["rgb_coarse"]
         rgb_fine = outputs["rgb_fine"]
         acc_coarse = colormaps.apply_colormap(outputs["accumulation_coarse"])

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():