Implement nerfacto alpha transparency training

nerfstudio/data/datasets/base_dataset.py

@@ -85,10 +85,10 @@ 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]
+        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
+from typing import Generator, Literal, Optional, Union, Dict
 
 import nerfacc
 import torch
@@ -101,26 +101,64 @@ class RGBRenderer(nn.Module):
             comp_rgb = torch.sum(weights * rgb, dim=-2)
             accumulated_weight = torch.sum(weights, dim=-2)
 
-        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)
-
-        assert isinstance(background_color, torch.Tensor)
+        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
+    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:
+            rgb: RGB for each sample.
+            background_color: The background color specification.
+
+        Returns:
+            Background color as 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)
+        assert isinstance(background_color, Tensor)
+        return background_color
+
+    @classmethod
+    def blend_background(
+        cls,
+        image: Tensor,
+        outputs: Dict[str, Union[Tensor, list]],
+    ) -> 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`.
+
+        Args:
+            image: RGBA.
+            outputs: Dict containing the model outputs.
+
+        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]
+        return blended_rgb
+
     def forward(
         self,
         rgb: Float[Tensor, "*bs num_samples 3"],
         weights: Float[Tensor, "*bs num_samples 1"],
         ray_indices: Optional[Int[Tensor, "num_samples"]] = None,
         num_rays: Optional[int] = None,
+        background_color: Union[Float[Tensor, "3"], BackgroundColor, None] = None,
     ) -> Float[Tensor, "*bs 3"]:
         """Composite samples along ray and render color image
 
@@ -129,15 +167,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_surface_model.py

@@ -292,7 +292,7 @@ class SurfaceModel(Model):
             metrics_dict: dictionary of metrics, some of which we can use for loss
         """
         loss_dict = {}
-        image = batch["image"].to(self.device)
+        image = batch["image"][..., :3].to(self.device)
         loss_dict["rgb_loss"] = self.rgb_loss(image, outputs["rgb"])
         if self.training:
             # eikonal loss
@@ -336,7 +336,7 @@ class SurfaceModel(Model):
             batch: ground truth batch corresponding to outputs
         """
         metrics_dict = {}
-        image = batch["image"].to(self.device)
+        image = batch["image"][..., :3].to(self.device)
         metrics_dict["psnr"] = self.psnr(outputs["rgb"], image)
         return metrics_dict
 
@@ -351,7 +351,7 @@ class SurfaceModel(Model):
         Returns:
             A dictionary of metrics.
         """
-        image = batch["image"].to(self.device)
+        image = batch["image"][..., :3].to(self.device)
         rgb = outputs["rgb"]
         acc = colormaps.apply_colormap(outputs["accumulation"])
 

nerfstudio/models/instant_ngp.py

@@ -218,14 +218,14 @@ class NGPModel(Model):
         return outputs
 
     def get_metrics_dict(self, outputs, batch):
-        image = batch["image"].to(self.device)
+        image = batch["image"][..., :3].to(self.device)
         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)
+        image = batch["image"][..., :3].to(self.device)
         rgb_loss = self.rgb_loss(image, outputs["rgb"])
         loss_dict = {"rgb_loss": rgb_loss}
         return loss_dict
@@ -233,7 +233,7 @@ class NGPModel(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)
+        image = batch["image"][..., :3].to(self.device)
         rgb = outputs["rgb"]
         acc = colormaps.apply_colormap(outputs["accumulation"])
         depth = colormaps.apply_depth_colormap(

nerfstudio/models/mipnerf.py

@@ -141,7 +141,7 @@ class MipNerfModel(Model):
         return outputs
 
     def get_loss_dict(self, outputs, batch, metrics_dict=None):
-        image = batch["image"].to(self.device)
+        image = batch["image"][..., :3].to(self.device)
         rgb_loss_coarse = self.rgb_loss(image, outputs["rgb_coarse"])
         rgb_loss_fine = self.rgb_loss(image, outputs["rgb_fine"])
         loss_dict = {"rgb_loss_coarse": rgb_loss_coarse, "rgb_loss_fine": rgb_loss_fine}
@@ -152,7 +152,7 @@ class MipNerfModel(Model):
         self, outputs: Dict[str, torch.Tensor], batch: Dict[str, torch.Tensor]
     ) -> 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 = batch["image"][..., :3].to(outputs["rgb_coarse"].device)
         rgb_coarse = outputs["rgb_coarse"]
         rgb_fine = outputs["rgb_fine"]
         acc_coarse = colormaps.apply_colormap(outputs["accumulation_coarse"])

nerfstudio/models/nerfacto.py

@@ -106,6 +106,8 @@ class NerfactoModelConfig(ModelConfig):
     """Orientation loss multiplier on computed normals."""
     pred_normal_loss_mult: float = 0.001
     """Predicted normal loss multiplier."""
+    use_alpha_transparency_training: bool = False
+    """Whether to use alpha channel transparency training."""
     use_proposal_weight_anneal: bool = True
     """Whether to use proposal weight annealing."""
     use_average_appearance_embedding: bool = True
@@ -283,7 +285,14 @@ class NerfactoModel(Model):
         weights_list.append(weights)
         ray_samples_list.append(ray_samples)
 
-        rgb = self.renderer_rgb(rgb=field_outputs[FieldHeadNames.RGB], weights=weights)
+        background_color = (
+            self.renderer_rgb.get_background_color(rgb=field_outputs[FieldHeadNames.RGB], background_color="random")
+            if self.config.use_alpha_transparency_training and self.training
+            else None
+        )
+        rgb = self.renderer_rgb(
+            rgb=field_outputs[FieldHeadNames.RGB], weights=weights, background_color=background_color
+        )
         depth = self.renderer_depth(weights=weights, ray_samples=ray_samples)
         accumulation = self.renderer_accumulation(weights=weights)
 
@@ -292,6 +301,8 @@ class NerfactoModel(Model):
             "accumulation": accumulation,
             "depth": depth,
         }
+        if self.config.use_alpha_transparency_training:
+            outputs["background_color"] = background_color
 
         if self.config.predict_normals:
             normals = self.renderer_normals(normals=field_outputs[FieldHeadNames.NORMALS], weights=weights)
@@ -321,7 +332,7 @@ class NerfactoModel(Model):
 
     def get_metrics_dict(self, outputs, batch):
         metrics_dict = {}
-        image = batch["image"].to(self.device)
+        image = batch["image"][..., :3].to(self.device)
         metrics_dict["psnr"] = self.psnr(outputs["rgb"], image)
         if self.training:
             metrics_dict["distortion"] = distortion_loss(outputs["weights_list"], outputs["ray_samples_list"])
@@ -329,7 +340,7 @@ class NerfactoModel(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(batch["image"].to(self.device), outputs)
         loss_dict["rgb_loss"] = self.rgb_loss(image, outputs["rgb"])
         if self.training:
             loss_dict["interlevel_loss"] = self.config.interlevel_loss_mult * interlevel_loss(
@@ -352,7 +363,7 @@ 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)
+        image = batch["image"][..., :3].to(self.device)
         rgb = outputs["rgb"]
         acc = colormaps.apply_colormap(outputs["accumulation"])
         depth = colormaps.apply_depth_colormap(

nerfstudio/models/semantic_nerfw.py

@@ -229,14 +229,14 @@ class SemanticNerfWModel(Model):
 
     def get_metrics_dict(self, outputs, batch):
         metrics_dict = {}
-        image = batch["image"].to(self.device)
+        image = batch["image"][..., :3].to(self.device)
         metrics_dict["psnr"] = self.psnr(outputs["rgb"], image)
         metrics_dict["distortion"] = distortion_loss(outputs["weights_list"], outputs["ray_samples_list"])
         return metrics_dict
 
     def get_loss_dict(self, outputs, batch, metrics_dict=None):
         loss_dict = {}
-        image = batch["image"].to(self.device)
+        image = batch["image"][..., :3].to(self.device)
         loss_dict["interlevel_loss"] = self.config.interlevel_loss_mult * interlevel_loss(
             outputs["weights_list"], outputs["ray_samples_list"]
         )
@@ -261,7 +261,7 @@ class SemanticNerfWModel(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)
+        image = batch["image"][..., :3].to(self.device)
         rgb = outputs["rgb"]
         rgb = torch.clamp(rgb, min=0, max=1)
         acc = colormaps.apply_colormap(outputs["accumulation"])

nerfstudio/models/tensorf.py

@@ -303,7 +303,7 @@ 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)
 
         rgb_loss = self.rgb_loss(image, outputs["rgb"])
 
@@ -333,7 +333,7 @@ class TensoRFModel(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(outputs["rgb"].device)
+        image = batch["image"][..., :3].to(outputs["rgb"].device)
         rgb = outputs["rgb"]
         acc = colormaps.apply_colormap(outputs["accumulation"])
         assert self.config.collider_params is not None

nerfstudio/models/vanilla_nerf.py

@@ -191,7 +191,7 @@ class NeRFModel(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_coarse"].device
-        image = batch["image"].to(device)
+        image = batch["image"][..., :3].to(device)
 
         rgb_loss_coarse = self.rgb_loss(image, outputs["rgb_coarse"])
         rgb_loss_fine = self.rgb_loss(image, outputs["rgb_fine"])
@@ -203,7 +203,7 @@ class NeRFModel(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(outputs["rgb_coarse"].device)
+        image = batch["image"][..., :3].to(outputs["rgb_coarse"].device)
         rgb_coarse = outputs["rgb_coarse"]
         rgb_fine = outputs["rgb_fine"]
         acc_coarse = colormaps.apply_colormap(outputs["accumulation_coarse"])