From e24f15baadf87fbc7d139283c403172cfe63bd65 Mon Sep 17 00:00:00 2001 From: Nicholas Pfaff Date: Tue, 4 Jul 2023 16:22:42 +0100 Subject: [PATCH] Stop adding all ones alpha channel and clean up blend_background --- nerfstudio/data/datasets/base_dataset.py | 3 --- nerfstudio/model_components/renderers.py | 26 ++++++++++++------------ nerfstudio/models/nerfacto.py | 13 +++++++++--- 3 files changed, 23 insertions(+), 19 deletions(-) diff --git a/nerfstudio/data/datasets/base_dataset.py b/nerfstudio/data/datasets/base_dataset.py index 8b0f820bc..6532d5d50 100644 --- a/nerfstudio/data/datasets/base_dataset.py +++ b/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: diff --git a/nerfstudio/model_components/renderers.py b/nerfstudio/model_components/renderers.py index 501da8c4c..bb813ca7f 100644 --- a/nerfstudio/model_components/renderers.py +++ b/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( diff --git a/nerfstudio/models/nerfacto.py b/nerfstudio/models/nerfacto.py index a4027a6af..d5c5a4d50 100644 --- a/nerfstudio/models/nerfacto.py +++ b/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"]