From 0bbaba6f1aa494499711dd1e027a71b0d540e60a Mon Sep 17 00:00:00 2001 From: Nicholas Pfaff Date: Mon, 3 Jul 2023 14:36:03 +0100 Subject: [PATCH] Implement nerfacto alpha transparency training --- nerfstudio/data/datasets/base_dataset.py | 6 +-- nerfstudio/model_components/renderers.py | 65 +++++++++++++++++++----- nerfstudio/models/base_surface_model.py | 6 +-- nerfstudio/models/instant_ngp.py | 6 +-- nerfstudio/models/mipnerf.py | 4 +- nerfstudio/models/nerfacto.py | 19 +++++-- nerfstudio/models/semantic_nerfw.py | 6 +-- nerfstudio/models/tensorf.py | 4 +- nerfstudio/models/vanilla_nerf.py | 4 +- 9 files changed, 86 insertions(+), 34 deletions(-) diff --git a/nerfstudio/data/datasets/base_dataset.py b/nerfstudio/data/datasets/base_dataset.py index 95269a162..8b0f820bc 100644 --- a/nerfstudio/data/datasets/base_dataset.py +++ b/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: diff --git a/nerfstudio/model_components/renderers.py b/nerfstudio/model_components/renderers.py index da3277908..501da8c4c 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 +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) diff --git a/nerfstudio/models/base_surface_model.py b/nerfstudio/models/base_surface_model.py index f7af6e2c1..7f37faf2e 100644 --- a/nerfstudio/models/base_surface_model.py +++ b/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"]) diff --git a/nerfstudio/models/instant_ngp.py b/nerfstudio/models/instant_ngp.py index 95f059b60..91aa49767 100644 --- a/nerfstudio/models/instant_ngp.py +++ b/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( diff --git a/nerfstudio/models/mipnerf.py b/nerfstudio/models/mipnerf.py index f0deb9bc4..083f7ecd2 100644 --- a/nerfstudio/models/mipnerf.py +++ b/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"]) diff --git a/nerfstudio/models/nerfacto.py b/nerfstudio/models/nerfacto.py index 2dbbee939..a4027a6af 100644 --- a/nerfstudio/models/nerfacto.py +++ b/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( diff --git a/nerfstudio/models/semantic_nerfw.py b/nerfstudio/models/semantic_nerfw.py index bf7914b1c..4c3de1e41 100644 --- a/nerfstudio/models/semantic_nerfw.py +++ b/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"]) diff --git a/nerfstudio/models/tensorf.py b/nerfstudio/models/tensorf.py index 742818b64..89a5bca52 100644 --- a/nerfstudio/models/tensorf.py +++ b/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 diff --git a/nerfstudio/models/vanilla_nerf.py b/nerfstudio/models/vanilla_nerf.py index 4a23bf0e5..097c8f551 100644 --- a/nerfstudio/models/vanilla_nerf.py +++ b/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"])