diff --git a/nerfstudio/model_components/renderers.py b/nerfstudio/model_components/renderers.py index f1e4931e5..5a5dee9d6 100644 --- a/nerfstudio/model_components/renderers.py +++ b/nerfstudio/model_components/renderers.py @@ -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, the predicted color is blended with black (0, 0, 0). Args: rgb: RGB for each sample @@ -101,9 +102,12 @@ 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 blended with black (0, 0, 0). + return comp_rgb + 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 @@ -124,65 +128,68 @@ class RGBRenderer(nn.Module): elif background_color == "last_sample": background_color = rgb[..., -1, :] elif background_color == "random": - background_color = torch.rand_like(rgb[..., -1, :]).to(rgb.device) + # background_color = torch.rand_like(rgb[..., -1, :]).to(rgb.device) + raise RuntimeError("get_background_color should never be called with background_color='random'") 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, - rgb: Tensor, - opacity: Tensor, - background_color: Float[Tensor, "3"], + self, + image: Tensor, + background_color: Optional[Float[Tensor, "3"]] = None, ) -> Float[Tensor, "*bs 3"]: - """Blends the background color into the image. + """Blends the background color into the image if image is RGBA. + Otherwise no blending is performed (we assume opacity of 1). Args: - rgb: RGB per pixel. + image: RGB/RGBA per pixel. opacity: Alpha opacity per pixel. background_color: Background color. Returns: Blended RGB. """ - if len(opacity.shape) < len(rgb.shape): - opacity = opacity.unsqueeze(-1) + if background_color is None: + background_color_name = self.background_color + if background_color_name == "last_sample" or background_color == "random": + background_color_name = "black" + background_color = self.get_background_color( + rgb=image.unsqueeze(-1), background_color=background_color_name + ) + if image.size(-1) < 4: + return image + rgb, opacity = image[..., :3], image[..., 3:] blended_rgb = rgb * opacity + background_color.to(rgb.device) * (1 - opacity) return blended_rgb - def blend_background_for_metric_computation( + def blend_background_for_loss_computation( self, + pred_image: Tensor, + pred_accumulation: Tensor, gt_image: Tensor, - predicted_rgb: Tensor, - predicted_opacity: Tensor, - background_color: Float[Tensor, "3"] = colors.BLACK, ) -> Tuple[Tensor, Tensor]: """Blends a background color into the ground truth and predicted image for - metric computation if the current background color is "random". + loss computation. Args: gt_image: The ground truth image. - predicted_rgb: The predicted RGB values. - predicted_opacity: The predicted opacity/ accumulation. - background_color: The background color to use for blending. This should be - a single color to use for all pixels. - + 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. """ - if gt_image.shape[-1] == 4 and self.background_color == "random": - predicted_rgb = self.blend_background( - rgb=predicted_rgb, opacity=predicted_opacity, background_color=background_color - ) - gt_rgb = self.blend_background( - rgb=gt_image[..., :3], opacity=gt_image[..., 3], background_color=background_color - ) - else: - gt_rgb = gt_image[..., :3] - return predicted_rgb, gt_rgb + if self.background_color == "random": + background_color = torch.rand_like(pred_image) + else: + background_color = self.get_background_color( + rgb=pred_image.unsqueeze(-1), background_color=self.background_color + ).to(pred_image.device) + 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, diff --git a/nerfstudio/models/nerfacto.py b/nerfstudio/models/nerfacto.py index 6674e197a..9bcdccb72 100644 --- a/nerfstudio/models/nerfacto.py +++ b/nerfstudio/models/nerfacto.py @@ -283,16 +283,7 @@ class NerfactoModel(Model): weights_list.append(weights) ray_samples_list.append(ray_samples) - background_color = ( - self.renderer_rgb.get_background_color( - rgb=field_outputs[FieldHeadNames.RGB], background_color=self.config.background_color - ) - if self.training - else None - ) - rgb = self.renderer_rgb( - rgb=field_outputs[FieldHeadNames.RGB], weights=weights, background_color=background_color - ) + rgb = self.renderer_rgb(rgb=field_outputs[FieldHeadNames.RGB], weights=weights) depth = self.renderer_depth(weights=weights, ray_samples=ray_samples) accumulation = self.renderer_accumulation(weights=weights) @@ -301,8 +292,6 @@ class NerfactoModel(Model): "accumulation": accumulation, "depth": depth, } - if self.training and self.config.background_color == "random" and background_color is not None: - outputs["background_color"] = background_color if self.config.predict_normals: normals = self.renderer_normals(normals=field_outputs[FieldHeadNames.NORMALS], weights=weights) @@ -332,10 +321,9 @@ class NerfactoModel(Model): def get_metrics_dict(self, outputs, batch): metrics_dict = {} - image = batch["image"].to(self.device) - predicted_rgb, gt_rgb = self.renderer_rgb.blend_background_for_metric_computation( - gt_image=image, predicted_rgb=outputs["rgb"], predicted_opacity=outputs["accumulation"] - ) + 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: @@ -345,14 +333,13 @@ class NerfactoModel(Model): def get_loss_dict(self, outputs, batch, metrics_dict=None): loss_dict = {} 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 image.shape[-1] == 4 and "background_color" in outputs - else image[..., :3] + 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(blended_rgb, outputs["rgb"]) + + 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"] @@ -374,10 +361,9 @@ 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) - predicted_rgb, gt_rgb = self.renderer_rgb.blend_background_for_metric_computation( - gt_image=image, predicted_rgb=outputs["rgb"], predicted_opacity=outputs["accumulation"] - ) + 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"], diff --git a/tests/field_components/test_mlp.py b/tests/field_components/test_mlp.py index e54c38a54..b52734654 100644 --- a/tests/field_components/test_mlp.py +++ b/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():