diff --git a/nerfstudio/model_components/renderers.py b/nerfstudio/model_components/renderers.py index ca8acadc1..5f0589442 100644 --- a/nerfstudio/model_components/renderers.py +++ b/nerfstudio/model_components/renderers.py @@ -55,6 +55,24 @@ def background_color_override_context(mode: Float[Tensor, "3"]) -> Generator[Non BACKGROUND_COLOR_OVERRIDE = old_background_color +def compute_median(quantity: torch.Tensor, weights: torch.Tensor) -> torch.Tensor: + """ + Compute the median of a quantity along rays. + + Args: + quantity: Quantity to compute median of. + weights: Weights along the rays + """ + + cumulative_weights = torch.cumsum(weights[..., 0], dim=-1) # [..., num_samples] + split = torch.ones((*weights.shape[:-2], 1), device=weights.device) * 0.5 # [..., 1] + median_index = torch.searchsorted(cumulative_weights, split, side="left") # [..., 1] + median_index = torch.clamp(median_index, 0, quantity.shape[-2] - 1) # [..., 1] + median_quantity = torch.gather(quantity[..., 0], dim=-1, index=median_index) # [..., 1] + + return median_quantity + + class RGBRenderer(nn.Module): """Standard volumetric rendering. @@ -335,7 +353,6 @@ class DepthRenderer(nn.Module): ray_samples: RaySamples, ray_indices: Optional[Int[Tensor, "num_samples"]] = None, num_rays: Optional[int] = None, - quantity: Optional[Float[Tensor, "*batch num_samples 1"]] = None, ) -> Float[Tensor, "*bs 1"]: """Composite samples along ray and calculate depths. @@ -350,37 +367,74 @@ class DepthRenderer(nn.Module): Outputs of depth values. """ - if quantity is None: - # set the quantity to steps - quantity = (ray_samples.frustums.starts + ray_samples.frustums.ends) / 2 + steps = (ray_samples.frustums.starts + ray_samples.frustums.ends) / 2 if self.method == "median": if ray_indices is not None and num_rays is not None: raise NotImplementedError("Median depth calculation is not implemented for packed samples.") - cumulative_weights = torch.cumsum(weights[..., 0], dim=-1) # [..., num_samples] - split = torch.ones((*weights.shape[:-2], 1), device=weights.device) * 0.5 # [..., 1] - median_index = torch.searchsorted(cumulative_weights, split, side="left") # [..., 1] - median_index = torch.clamp(median_index, 0, quantity.shape[-2] - 1) # [..., 1] - median_depth = torch.gather(quantity[..., 0], dim=-1, index=median_index) # [..., 1] + median_depth = compute_median(steps, weights) return median_depth if self.method == "expected": eps = 1e-10 if ray_indices is not None and num_rays is not None: # Necessary for packed samples from volumetric ray sampler - depth = nerfacc.accumulate_along_rays(weights, ray_indices, quantity, num_rays) + depth = nerfacc.accumulate_along_rays(weights, ray_indices, steps, num_rays) accumulation = nerfacc.accumulate_along_rays(weights, ray_indices, None, num_rays) depth = depth / (accumulation + eps) else: - depth = torch.sum(weights * quantity, dim=-2) / (torch.sum(weights, -2) + eps) + depth = torch.sum(weights * steps, dim=-2) / (torch.sum(weights, -2) + eps) - depth = torch.clip(depth, quantity.min(), quantity.max()) + depth = torch.clip(depth, steps.min(), steps.max()) return depth raise NotImplementedError(f"Method {self.method} not implemented") +class VisibilityRenderer(nn.Module): + """Calculate visibility along ray. + + Visibity Method: + - median: Visibility is set to the num_views where the accumulated weight reaches 0.5. + + Args: + method: Visibility calculation method. + """ + + def __init__(self, method: Literal["median"] = "median") -> None: + super().__init__() + self.method = method + + def forward( + self, + visibility_samples: Float[Tensor, "*batch num_samples 1"], + weights: Float[Tensor, "*batch num_samples 1"], + ray_indices: Optional[Int[Tensor, "num_samples"]] = None, + num_rays: Optional[int] = None, + ) -> Float[Tensor, "*bs 1"]: + """Composite samples along ray and calculate depths. + + Args: + weights: Weights for each sample. + ray_samples: Set of ray samples. + ray_indices: Ray index for each sample, used when samples are packed. + num_rays: Number of rays, used when samples are packed. + num_views: Number of views each point is visible in. + + Returns: + Outputs of depth values. + """ + + if self.method == "median": + if ray_indices is not None and num_rays is not None: + raise NotImplementedError("Median visibility calculation is not implemented for packed samples.") + median_visibility = compute_median(visibility_samples, weights) + return median_visibility + else: + raise NotImplementedError(f"Method {self.method} not implemented") + + class UncertaintyRenderer(nn.Module): """Calculate uncertainty along the ray.""" diff --git a/nerfstudio/models/nerfacto.py b/nerfstudio/models/nerfacto.py index e90c443e9..d0151dff2 100644 --- a/nerfstudio/models/nerfacto.py +++ b/nerfstudio/models/nerfacto.py @@ -44,7 +44,13 @@ from nerfstudio.model_components.losses import ( scale_gradients_by_distance_squared, ) from nerfstudio.model_components.ray_samplers import ProposalNetworkSampler, UniformSampler -from nerfstudio.model_components.renderers import AccumulationRenderer, DepthRenderer, NormalsRenderer, RGBRenderer +from nerfstudio.model_components.renderers import ( + AccumulationRenderer, + DepthRenderer, + NormalsRenderer, + RGBRenderer, + VisibilityRenderer, +) from nerfstudio.model_components.scene_colliders import NearFarCollider from nerfstudio.model_components.shaders import NormalsShader from nerfstudio.models.base_model import Model, ModelConfig @@ -226,8 +232,7 @@ class NerfactoModel(Model): self.renderer_rgb = RGBRenderer(background_color=self.config.background_color) self.renderer_accumulation = AccumulationRenderer() self.renderer_depth = DepthRenderer() - self.renderer_depth_expected = DepthRenderer(method="expected") - self.renderer_depth_median = DepthRenderer(method="median") + self.renderer_visibility = VisibilityRenderer(method="median") self.renderer_normals = NormalsRenderer() # shaders @@ -308,20 +313,11 @@ class NerfactoModel(Model): if self.visibility_field is not None: assert isinstance(self.visibility_field, VisibilityField), "self.visibility_field must be a VisibilityField" visibility_samples = self.visibility_field.forward(ray_samples_list[-1]) - visibility = self.renderer_depth_expected( - weights=weights, ray_samples=ray_samples, quantity=visibility_samples - ) + visibility = self.renderer_visibility(weights=weights, visibility_samples=visibility_samples) visibility = visibility.float() / len( self.visibility_field.c2ws ) # range [0, 1] where 1 is seen by all cameras - outputs["visibility_expected"] = visibility - visibility_median = self.renderer_depth_median( - weights=weights, ray_samples=ray_samples, quantity=visibility_samples - ) - visibility_median = visibility_median.float() / len( - self.visibility_field.c2ws - ) # range [0, 1] where 1 is seen by all cameras - outputs["visibility_median"] = visibility_median + outputs["visibility"] = visibility if self.config.predict_normals: normals = self.renderer_normals(normals=field_outputs[FieldHeadNames.NORMALS], weights=weights) diff --git a/nerfstudio/pipelines/base_pipeline.py b/nerfstudio/pipelines/base_pipeline.py index be5b10485..4f82df5fe 100644 --- a/nerfstudio/pipelines/base_pipeline.py +++ b/nerfstudio/pipelines/base_pipeline.py @@ -226,15 +226,13 @@ class VanillaPipelineConfig(cfg.InstantiateConfig): """whether to apply the visibility loss""" visibility_steps_per_loss: int = 10 """how often to apply the visibility loss""" - visibility_loss_quantity: Literal["weights", "densities"] = "densities" - """whether to apply the visibility loss to weights or densities""" visibility_loss_mult: float = 1.0 """multiplier for the visibility loss""" visibility_min_views: int = 1 - """minimum number of views for the visibility loss""" + """apply visibility loss where fewer than this min_views are seen.""" visibility_num_rays: int = 10 """number of rays per batch to use for the visibility loss""" - visibility_radius: float = 1.0 + visibility_radius: float = 2.0 """radius for the visibility loss""" @@ -293,7 +291,6 @@ class VanillaPipeline(Pipeline): if self.config.use_visibility_loss: # initialize the visibility field cameras = typing.cast(Cameras, self.datamanager.train_dataparser_outputs.cameras.to(self.device)) # type: ignore - # TODO(ethan): make sure only the training cameras are being used self.model.visibility_field = VisibilityField(cameras).to(self.device) @property @@ -349,22 +346,17 @@ class VanillaPipeline(Pipeline): # We only use the first camera index, but it doesn't matter since we don't care about appearance embeddings # when using the visibility loss. camera_indices = torch.tensor([0] * self.config.visibility_num_rays).unsqueeze(-1) - ray_bundle_rays = cameras.generate_rays(camera_indices) - # ray_bundle is (1, 1, self.config.visibility_num_rays) - ray_bundle_rays = ray_bundle_rays.flatten() + ray_bundle_rays = cameras.generate_rays(camera_indices).flatten() model_outputs_rays = self.model(ray_bundle_rays) - quantity_list = model_outputs_rays[ - f"{self.config.visibility_loss_quantity}_list" - ] # weights_list or densities_list visibility_loss = torch.zeros(1, device=self.device) - for i in range(len(quantity_list)): - quantity_samples = quantity_list[i] + for i in range(len(model_outputs_rays["densities_list"])): + density_samples = model_outputs_rays["densities_list"][i] ray_samples = model_outputs_rays["ray_samples_list"][i] visibility_samples = self.model.visibility_field(ray_samples) # type: ignore - quantity_samples_masked = quantity_samples[visibility_samples < self.config.visibility_min_views] - if quantity_samples_masked.numel() != 0: - visibility_loss_i = self.config.visibility_loss_mult * quantity_samples_masked.mean() - visibility_loss += visibility_loss_i + # penalize densities seen in less than min_views + density_samples_masked = density_samples[visibility_samples < self.config.visibility_min_views] + if density_samples_masked.numel() != 0: + visibility_loss += self.config.visibility_loss_mult * density_samples_masked.mean() loss_dict.update({"visibility_loss": visibility_loss}) return model_outputs, loss_dict, metrics_dict