Add median depth calculation (#877)

* Add median depth calculation

* lint

* Fix test

nerfstudio/model_components/renderers.py

@@ -200,11 +200,15 @@ class AccumulationRenderer(nn.Module):
 class DepthRenderer(nn.Module):
     """Calculate depth along ray.
 
+    Depth Method:
+        - median: Depth is set to the distance where the accumulated weight reaches 0.5.
+        - expected: Expected depth along ray. Same procedure as rendering rgb, but with depth.
+
     Args:
-        method (str, optional): Depth calculation method.
+        method: Depth calculation method.
     """
 
-    def __init__(self, method: Literal["expected"] = "expected") -> None:
+    def __init__(self, method: Literal["median", "expected"] = "median") -> None:
         super().__init__()
         self.method = method
 
@@ -215,7 +219,7 @@ class DepthRenderer(nn.Module):
         ray_indices: Optional[TensorType["num_samples"]] = None,
         num_rays: Optional[int] = None,
     ) -> TensorType[..., 1]:
-        """Composite samples along ray and calculate disparities.
+        """Composite samples along ray and calculate depths.
 
         Args:
             weights: Weights for each sample.
@@ -227,6 +231,17 @@ class DepthRenderer(nn.Module):
             Outputs of depth values.
         """
 
+        if self.method == "median":
+            steps = (ray_samples.frustums.starts + ray_samples.frustums.ends) / 2
+
+            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, steps.shape[-2] - 1)  # [..., 1]
+            median_depth = torch.gather(steps[..., 0], dim=-1, index=median_index)  # [..., 1]
+            return median_depth
         if self.method == "expected":
             eps = 1e-10
             steps = (ray_samples.frustums.starts + ray_samples.frustums.ends) / 2

nerfstudio/models/instant_ngp.py

@@ -125,7 +125,7 @@ class NGPModel(Model):
         background_color = "random" if self.config.randomize_background else colors.WHITE
         self.renderer_rgb = RGBRenderer(background_color=background_color)
         self.renderer_accumulation = AccumulationRenderer()
-        self.renderer_depth = DepthRenderer()
+        self.renderer_depth = DepthRenderer(method="expected")
 
         # losses
         self.rgb_loss = MSELoss()

tests/model_components/test_renderers.py

@@ -60,7 +60,7 @@ def test_depth_renderer():
     """Test depth rendering"""
 
     num_samples = 10
-    weights = torch.ones((3, num_samples, 1))
+    weights = torch.ones((num_samples, 1))
     weights /= torch.sum(weights, dim=-2, keepdim=True)
 
     frustums = Frustums.get_mock_frustum()
@@ -73,8 +73,11 @@ def test_depth_renderer():
         deltas=torch.ones((num_samples, 1)),
     )
 
-    depth_renderer = renderers.DepthRenderer()
+    depth_renderer = renderers.DepthRenderer(method="median")
+    depth = depth_renderer(weights=weights, ray_samples=ray_samples)
+    assert torch.min(depth) > 0
 
+    depth_renderer = renderers.DepthRenderer(method="expected")
     depth = depth_renderer(weights=weights, ray_samples=ray_samples)
     assert torch.min(depth) > 0