Stop adding all ones alpha channel and clean up blend_background

nerfstudio/data/datasets/base_dataset.py

@@ -86,9 +86,6 @@ 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:
             image = image[:, :, :3] * image[:, :, -1:] + self._dataparser_outputs.alpha_color * (1.0 - image[:, :, -1:])
-        elif image.shape[-1] == 3:
-            alpha_channel = torch.ones(size=(*image.shape[:2], 1), dtype=image.dtype)
-            image = torch.concatenate([image, alpha_channel], dim=2)
         return image
 
     def get_data(self, image_idx: int) -> Dict:

nerfstudio/model_components/renderers.py

@@ -28,7 +28,7 @@ Example:
 """
 import contextlib
 import math
-from typing import Generator, Literal, Optional, Union, Dict
+from typing import Generator, Literal, Optional, Union
 
 import nerfacc
 import torch
@@ -131,25 +131,25 @@ class RGBRenderer(nn.Module):
     @classmethod
     def blend_background(
         cls,
-        image: Tensor,
-        outputs: Dict[str, Union[Tensor, list]],
+        rgb: Tensor,
+        opacity: Tensor,
+        background_color: Optional[Tensor] = None,
     ) -> Float[Tensor, "*bs 3"]:
-        """Blends the background color into the image using the alpha channel. Removes
-        the alpha channel if there is no `background_color` in `outputs`.
+        """Blends the background color into the image.
 
         Args:
-            image: RGBA.
-            outputs: Dict containing the model outputs.
+            rgb: RGB per pixel.
+            opacity: Alpha opacity per pixel.
+            background_color: Background color.
 
         Returns:
             Blended RGB.
         """
-        if outputs.get("background_color", None) is not None:
-            background_color = outputs["background_color"]
-            assert isinstance(background_color, Tensor)
-            blended_rgb = image[:, :3] * image[:, 3:] + background_color.to(image.device) * (1 - image[:, 3:])
-        else:
-            blended_rgb = image[:, :3]
+        if len(opacity.shape) == 1:
+            opacity = opacity.unsqueeze(-1)
+        blended_rgb = (
+            rgb * opacity + background_color.to(rgb.device) * (1 - opacity) if background_color is not None else rgb
+        )
         return blended_rgb
 
     def forward(

nerfstudio/models/nerfacto.py

@@ -301,7 +301,7 @@ class NerfactoModel(Model):
             "accumulation": accumulation,
             "depth": depth,
         }
-        if self.config.use_alpha_transparency_training:
+        if background_color is not None:
             outputs["background_color"] = background_color
 
         if self.config.predict_normals:
@@ -340,8 +340,15 @@ class NerfactoModel(Model):
 
     def get_loss_dict(self, outputs, batch, metrics_dict=None):
         loss_dict = {}
-        image = self.renderer_rgb.blend_background(batch["image"].to(self.device), outputs)
-        loss_dict["rgb_loss"] = self.rgb_loss(image, outputs["rgb"])
+        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 "background_color" in outputs
+            else image[..., :3]
+        )
+        loss_dict["rgb_loss"] = self.rgb_loss(blended_rgb, outputs["rgb"])
         if self.training:
             loss_dict["interlevel_loss"] = self.config.interlevel_loss_mult * interlevel_loss(
                 outputs["weights_list"], outputs["ray_samples_list"]