diff --git a/nerfstudio/data/datasets/base_dataset.py b/nerfstudio/data/datasets/base_dataset.py index 95269a162..6532d5d50 100644 --- a/nerfstudio/data/datasets/base_dataset.py +++ b/nerfstudio/data/datasets/base_dataset.py @@ -85,10 +85,7 @@ 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] return image def get_data(self, image_idx: int) -> Dict: diff --git a/nerfstudio/exporter/exporter_utils.py b/nerfstudio/exporter/exporter_utils.py index 21ea545b9..1baeeb2b5 100644 --- a/nerfstudio/exporter/exporter_utils.py +++ b/nerfstudio/exporter/exporter_utils.py @@ -28,13 +28,7 @@ import open3d as o3d import pymeshlab import torch from jaxtyping import Float -from rich.progress import ( - BarColumn, - Progress, - TaskProgressColumn, - TextColumn, - TimeRemainingColumn, -) +from rich.progress import BarColumn, Progress, TaskProgressColumn, TextColumn, TimeRemainingColumn from torch import Tensor from nerfstudio.cameras.cameras import Cameras @@ -144,7 +138,7 @@ def generate_point_cloud( CONSOLE.print(f"Could not find {depth_output_name} in the model outputs", justify="center") CONSOLE.print(f"Please set --depth_output_name to one of: {outputs.keys()}", justify="center") sys.exit(1) - rgb = outputs[rgb_output_name] + rgba = pipeline.model.get_rgba_image(outputs, rgb_output_name) depth = outputs[depth_output_name] if normal_output_name is not None: if normal_output_name not in outputs: @@ -160,6 +154,14 @@ def generate_point_cloud( point = ray_bundle.origins + ray_bundle.directions * depth view_direction = ray_bundle.directions + # Filter points with opacity lower than 0.5 + mask = rgba[..., -1] > 0.5 + point = point[mask] + view_direction = view_direction[mask] + rgb = rgba[mask][..., :3] + if normal is not None: + normal = normal[mask] + if use_bounding_box: comp_l = torch.tensor(bounding_box_min, device=point.device) comp_m = torch.tensor(bounding_box_max, device=point.device) @@ -230,6 +232,7 @@ def render_trajectory( depth_output_name: str, rendered_resolution_scaling_factor: float = 1.0, disable_distortion: bool = False, + return_rgba_images: bool = False, ) -> Tuple[List[np.ndarray], List[np.ndarray]]: """Helper function to create a video of a trajectory. @@ -240,6 +243,7 @@ def render_trajectory( depth_output_name: Name of the depth output. rendered_resolution_scaling_factor: Scaling factor to apply to the camera image resolution. disable_distortion: Whether to disable distortion. + return_rgba_images: Whether to return RGBA images (default RGB). Returns: List of rgb images, list of depth images. @@ -272,7 +276,11 @@ def render_trajectory( CONSOLE.print(f"Could not find {depth_output_name} in the model outputs", justify="center") CONSOLE.print(f"Please set --depth_output_name to one of: {outputs.keys()}", justify="center") sys.exit(1) - images.append(outputs[rgb_output_name].cpu().numpy()) + if return_rgba_images: + image = pipeline.model.get_rgba_image(outputs, rgb_output_name) + else: + image = outputs[rgb_output_name] + images.append(image.cpu().numpy()) depths.append(outputs[depth_output_name].cpu().numpy()) return images, depths diff --git a/nerfstudio/exporter/texture_utils.py b/nerfstudio/exporter/texture_utils.py index 8a304a7b8..120696752 100644 --- a/nerfstudio/exporter/texture_utils.py +++ b/nerfstudio/exporter/texture_utils.py @@ -403,9 +403,11 @@ def export_textured_mesh( CONSOLE.print("Creating texture image by rendering with NeRF...") with torch.no_grad(): outputs = pipeline.model.get_outputs_for_camera_ray_bundle(camera_ray_bundle) + # TODO: this can be done better by using the alpha channel + rgb = pipeline.model.get_rgba_image(outputs, "rgb")[..., :3] # save the texture image - texture_image = outputs["rgb"].cpu().numpy() + texture_image = rgb.cpu().numpy() media.write_image(str(output_dir / "material_0.png"), texture_image) CONSOLE.print("Writing relevant OBJ information to files...") diff --git a/nerfstudio/exporter/tsdf_utils.py b/nerfstudio/exporter/tsdf_utils.py index d144686c9..301591316 100644 --- a/nerfstudio/exporter/tsdf_utils.py +++ b/nerfstudio/exporter/tsdf_utils.py @@ -331,8 +331,12 @@ def export_tsdf_mesh( depth_output_name=depth_output_name, rendered_resolution_scaling_factor=1.0 / downscale_factor, disable_distortion=True, + return_rgba_images=True, ) + # TODO: this can be done better by removing transparent points from the TSDF + color_images = [x[..., :3] for x in color_images] + # camera extrinsics and intrinsics c2w: Float[Tensor, "N 3 4"] = cameras.camera_to_worlds.to(device) # make c2w homogeneous diff --git a/nerfstudio/model_components/renderers.py b/nerfstudio/model_components/renderers.py index da3277908..ae50b0e12 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, Tuple, Union import nerfacc import torch @@ -39,7 +39,7 @@ from nerfstudio.cameras.rays import RaySamples from nerfstudio.utils import colors from nerfstudio.utils.math import components_from_spherical_harmonics, safe_normalize -BackgroundColor = Union[Literal["random", "last_sample", "black", "white"], Float[Tensor, "3"]] +BackgroundColor = Union[Literal["random", "last_sample", "black", "white"], Float[Tensor, "3"], Float[Tensor, "*bs 3"]] BACKGROUND_COLOR_OVERRIDE: Optional[Float[Tensor, "3"]] = None @@ -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, no BG color is added - as if the background was black! Args: rgb: RGB for each sample @@ -101,19 +102,99 @@ 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 returned without blending, + # as if the background color was black. + return comp_rgb + + elif background_color == "last_sample": + # Note, this is only supported for non-packed samples. + background_color = rgb[..., -1, :] + else: + background_color = cls.get_background_color(background_color, shape=comp_rgb.shape, device=comp_rgb.device) + assert isinstance(background_color, torch.Tensor) + comp_rgb = comp_rgb + background_color * (1.0 - accumulated_weight) + return comp_rgb + + @classmethod + def get_background_color( + cls, background_color: BackgroundColor, shape: Tuple[int, ...], device: torch.device + ) -> Union[Float[Tensor, "3"], Float[Tensor, "*bs 3"]]: + """Returns the RGB background color for a specified background color. + + Note: + This function CANNOT be called for background_color being either "last_sample" or "random". + + Args: + rgb: RGB for each sample. + background_color: The background color specification. + + Returns: + Background color as RGB. + """ + assert background_color not in {"last_sample", "random"} + assert shape[-1] == 3, "Background color must be RGB." 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) + background_color = colors.COLORS_DICT[background_color] + assert isinstance(background_color, Tensor) + # Ensure correct shape + return background_color.expand(shape).to(device) + + def blend_background( + self, + image: Tensor, + background_color: Optional[BackgroundColor] = None, + ) -> Float[Tensor, "*bs 3"]: + """Blends the background color into the image if image is RGBA. + Otherwise no blending is performed (we assume opacity of 1). + + Args: + image: RGB/RGBA per pixel. + opacity: Alpha opacity per pixel. + background_color: Background color. + + Returns: + Blended RGB. + """ + if image.size(-1) < 4: + return image + + rgb, opacity = image[..., :3], image[..., 3:] + if background_color is None: + background_color = self.background_color + if background_color in {"last_sample", "random"}: + background_color = "black" + background_color = self.get_background_color(background_color, shape=rgb.shape, device=rgb.device) assert isinstance(background_color, torch.Tensor) - comp_rgb = comp_rgb + background_color.to(weights.device) * (1.0 - accumulated_weight) + return rgb * opacity + background_color.to(rgb.device) * (1 - opacity) - return comp_rgb + def blend_background_for_loss_computation( + self, + pred_image: Tensor, + pred_accumulation: Tensor, + gt_image: Tensor, + ) -> Tuple[Tensor, Tensor]: + """Blends a background color into the ground truth and predicted image for + loss computation. + + Args: + gt_image: The ground truth image. + 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. + """ + background_color = self.background_color + if background_color == "last_sample": + background_color = "black" # No background blending for GT + elif background_color == "random": + background_color = torch.rand_like(pred_image) + 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, @@ -121,6 +202,7 @@ class RGBRenderer(nn.Module): weights: Float[Tensor, "*bs num_samples 1"], ray_indices: Optional[Int[Tensor, "num_samples"]] = None, num_rays: Optional[int] = None, + background_color: Optional[BackgroundColor] = None, ) -> Float[Tensor, "*bs 3"]: """Composite samples along ray and render color image @@ -129,15 +211,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_model.py b/nerfstudio/models/base_model.py index 6f31c3b30..ebda710b3 100644 --- a/nerfstudio/models/base_model.py +++ b/nerfstudio/models/base_model.py @@ -187,6 +187,30 @@ class Model(nn.Module): outputs[output_name] = torch.cat(outputs_list).view(image_height, image_width, -1) # type: ignore return outputs + def get_rgba_image(self, outputs: Dict[str, torch.Tensor], output_name: str = "rgb") -> torch.Tensor: + """Returns the RGBA image from the outputs of the model. + + Args: + outputs: Outputs of the model. + + Returns: + RGBA image. + """ + accumulation_name = output_name.replace("rgb", "accumulation") + if ( + not hasattr(self, "renderer_rgb") + or not hasattr(self.renderer_rgb, "background_color") + or accumulation_name not in outputs + ): + raise NotImplementedError(f"get_rgba_image is not implemented for model {self.__class__.__name__}") + rgb = outputs[output_name] + if self.renderer_rgb.background_color == "random": # type: ignore + acc = outputs[accumulation_name] + if acc.dim() < rgb.dim(): + acc = acc.unsqueeze(-1) + return torch.cat((rgb / acc.clamp(min=1e-10), acc), dim=-1) + return torch.cat((rgb, torch.ones_like(rgb[..., :1])), dim=-1) + @abstractmethod def get_image_metrics_and_images( self, outputs: Dict[str, torch.Tensor], batch: Dict[str, torch.Tensor] diff --git a/nerfstudio/models/base_surface_model.py b/nerfstudio/models/base_surface_model.py index f8d3bca10..638aa2713 100644 --- a/nerfstudio/models/base_surface_model.py +++ b/nerfstudio/models/base_surface_model.py @@ -293,7 +293,12 @@ class SurfaceModel(Model): """ loss_dict = {} image = batch["image"].to(self.device) - loss_dict["rgb_loss"] = self.rgb_loss(image, outputs["rgb"]) + pred_image, image = 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(image, pred_image) if self.training: # eikonal loss grad_theta = outputs["eik_grad"] @@ -337,6 +342,7 @@ class SurfaceModel(Model): """ metrics_dict = {} image = batch["image"].to(self.device) + image = self.renderer_rgb.blend_background(image) metrics_dict["psnr"] = self.psnr(outputs["rgb"], image) return metrics_dict @@ -352,6 +358,7 @@ class SurfaceModel(Model): A dictionary of metrics. """ image = batch["image"].to(self.device) + image = self.renderer_rgb.blend_background(image) rgb = outputs["rgb"] acc = colormaps.apply_colormap(outputs["accumulation"]) diff --git a/nerfstudio/models/instant_ngp.py b/nerfstudio/models/instant_ngp.py index abdc892dd..c67d45db5 100644 --- a/nerfstudio/models/instant_ngp.py +++ b/nerfstudio/models/instant_ngp.py @@ -219,14 +219,20 @@ class NGPModel(Model): def get_metrics_dict(self, outputs, batch): image = batch["image"].to(self.device) + image = self.renderer_rgb.blend_background(image) 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) - rgb_loss = self.rgb_loss(image, outputs["rgb"]) + image = batch["image"][..., :3].to(self.device) + pred_rgb, image = self.renderer_rgb.blend_background_for_loss_computation( + pred_image=outputs["rgb"], + pred_accumulation=outputs["accumulation"], + gt_image=image, + ) + rgb_loss = self.rgb_loss(image, pred_rgb) loss_dict = {"rgb_loss": rgb_loss} return loss_dict @@ -234,6 +240,7 @@ class NGPModel(Model): 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 = self.renderer_rgb.blend_background(image) 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 9a3fdabe2..a3a783461 100644 --- a/nerfstudio/models/mipnerf.py +++ b/nerfstudio/models/mipnerf.py @@ -142,8 +142,18 @@ class MipNerfModel(Model): def get_loss_dict(self, outputs, batch, metrics_dict=None): image = batch["image"].to(self.device) - rgb_loss_coarse = self.rgb_loss(image, outputs["rgb_coarse"]) - rgb_loss_fine = self.rgb_loss(image, outputs["rgb_fine"]) + pred_coarse, image_coarse = self.renderer_rgb.blend_background_for_loss_computation( + pred_image=outputs["rgb_coarse"], + pred_accumulation=outputs["accumulation_coarse"], + gt_image=image, + ) + pred_fine, image_fine = self.renderer_rgb.blend_background_for_loss_computation( + pred_image=outputs["rgb_fine"], + pred_accumulation=outputs["accumulation_fine"], + gt_image=image, + ) + rgb_loss_coarse = self.rgb_loss(image_coarse, pred_coarse) + rgb_loss_fine = self.rgb_loss(image_fine, pred_fine) loss_dict = {"rgb_loss_coarse": rgb_loss_coarse, "rgb_loss_fine": rgb_loss_fine} loss_dict = misc.scale_dict(loss_dict, self.config.loss_coefficients) return loss_dict @@ -153,6 +163,7 @@ class MipNerfModel(Model): ) -> 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 = self.renderer_rgb.blend_background(image) 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 3024c77cd..065fb571e 100644 --- a/nerfstudio/models/nerfacto.py +++ b/nerfstudio/models/nerfacto.py @@ -323,8 +323,11 @@ class NerfactoModel(Model): def get_metrics_dict(self, outputs, batch): metrics_dict = {} - image = batch["image"].to(self.device) - metrics_dict["psnr"] = self.psnr(outputs["rgb"], image) + 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: metrics_dict["distortion"] = distortion_loss(outputs["weights_list"], outputs["ray_samples_list"]) return metrics_dict @@ -332,7 +335,13 @@ class NerfactoModel(Model): def get_loss_dict(self, outputs, batch, metrics_dict=None): loss_dict = {} image = batch["image"].to(self.device) - loss_dict["rgb_loss"] = self.rgb_loss(image, outputs["rgb"]) + 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(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"] @@ -354,25 +363,26 @@ 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) - rgb = outputs["rgb"] + 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"], accumulation=outputs["accumulation"], ) - combined_rgb = torch.cat([image, rgb], dim=1) + combined_rgb = torch.cat([gt_rgb, predicted_rgb], dim=1) combined_acc = torch.cat([acc], dim=1) combined_depth = torch.cat([depth], dim=1) # Switch images from [H, W, C] to [1, C, H, W] for metrics computations - image = torch.moveaxis(image, -1, 0)[None, ...] - rgb = torch.moveaxis(rgb, -1, 0)[None, ...] + gt_rgb = torch.moveaxis(gt_rgb, -1, 0)[None, ...] + predicted_rgb = torch.moveaxis(predicted_rgb, -1, 0)[None, ...] - psnr = self.psnr(image, rgb) - ssim = self.ssim(image, rgb) - lpips = self.lpips(image, rgb) + psnr = self.psnr(gt_rgb, predicted_rgb) + ssim = self.ssim(gt_rgb, predicted_rgb) + lpips = self.lpips(gt_rgb, predicted_rgb) # all of these metrics will be logged as scalars metrics_dict = {"psnr": float(psnr.item()), "ssim": float(ssim)} # type: ignore diff --git a/nerfstudio/models/semantic_nerfw.py b/nerfstudio/models/semantic_nerfw.py index 106f1794a..87388d0e9 100644 --- a/nerfstudio/models/semantic_nerfw.py +++ b/nerfstudio/models/semantic_nerfw.py @@ -230,6 +230,7 @@ class SemanticNerfWModel(Model): def get_metrics_dict(self, outputs, batch): metrics_dict = {} image = batch["image"].to(self.device) + image = self.renderer_rgb.blend_background(image) metrics_dict["psnr"] = self.psnr(outputs["rgb"], image) metrics_dict["distortion"] = distortion_loss(outputs["weights_list"], outputs["ray_samples_list"]) return metrics_dict @@ -237,6 +238,7 @@ class SemanticNerfWModel(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(image) loss_dict["interlevel_loss"] = self.config.interlevel_loss_mult * interlevel_loss( outputs["weights_list"], outputs["ray_samples_list"] ) @@ -263,6 +265,12 @@ class SemanticNerfWModel(Model): ) -> Tuple[Dict[str, float], Dict[str, torch.Tensor]]: image = batch["image"].to(self.device) rgb = outputs["rgb"] + rgb, image = self.renderer_rgb.blend_background_for_loss_computation( + pred_image=rgb, + pred_accumulation=outputs["accumulation"], + gt_image=image, + ) + rgb = torch.clamp(rgb, min=0, max=1) acc = colormaps.apply_colormap(outputs["accumulation"]) depth = colormaps.apply_depth_colormap( diff --git a/nerfstudio/models/tensorf.py b/nerfstudio/models/tensorf.py index 160ef1b86..a2cd9901d 100644 --- a/nerfstudio/models/tensorf.py +++ b/nerfstudio/models/tensorf.py @@ -303,9 +303,14 @@ 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) + pred_image, image = self.renderer_rgb.blend_background_for_loss_computation( + pred_image=outputs["rgb"], + pred_accumulation=outputs["accumulation"], + gt_image=image, + ) - rgb_loss = self.rgb_loss(image, outputs["rgb"]) + rgb_loss = self.rgb_loss(image, pred_image) loss_dict = {"rgb_loss": rgb_loss} @@ -334,6 +339,7 @@ class TensoRFModel(Model): 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 = self.renderer_rgb.blend_background(image) 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 62d6b96b7..ea5344fb8 100644 --- a/nerfstudio/models/vanilla_nerf.py +++ b/nerfstudio/models/vanilla_nerf.py @@ -192,9 +192,19 @@ class NeRFModel(Model): # Scaling metrics by coefficients to create the losses. device = outputs["rgb_coarse"].device image = batch["image"].to(device) + coarse_pred, coarse_image = self.renderer_rgb.blend_background_for_loss_computation( + pred_image=outputs["rgb_coarse"], + pred_accumulation=outputs["accumulation_coarse"], + gt_image=image, + ) + fine_pred, fine_image = self.renderer_rgb.blend_background_for_loss_computation( + pred_image=outputs["rgb_fine"], + pred_accumulation=outputs["accumulation_fine"], + gt_image=image, + ) - rgb_loss_coarse = self.rgb_loss(image, outputs["rgb_coarse"]) - rgb_loss_fine = self.rgb_loss(image, outputs["rgb_fine"]) + rgb_loss_coarse = self.rgb_loss(coarse_image, coarse_pred) + rgb_loss_fine = self.rgb_loss(fine_image, fine_pred) loss_dict = {"rgb_loss_coarse": rgb_loss_coarse, "rgb_loss_fine": rgb_loss_fine} loss_dict = misc.scale_dict(loss_dict, self.config.loss_coefficients) @@ -204,6 +214,7 @@ class NeRFModel(Model): 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 = self.renderer_rgb.blend_background(image) rgb_coarse = outputs["rgb_coarse"] rgb_fine = outputs["rgb_fine"] acc_coarse = colormaps.apply_colormap(outputs["accumulation_coarse"]) 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():