create a renderer specific to visibility

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."""
 

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)

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