diff --git a/nerfstudio/model_components/renderers.py b/nerfstudio/model_components/renderers.py index 9604a1d5f..21726acc5 100644 --- a/nerfstudio/model_components/renderers.py +++ b/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 diff --git a/nerfstudio/models/instant_ngp.py b/nerfstudio/models/instant_ngp.py index 480c52ee8..f30d2ad9a 100644 --- a/nerfstudio/models/instant_ngp.py +++ b/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() diff --git a/tests/model_components/test_renderers.py b/tests/model_components/test_renderers.py index 26417e4e4..984a45e7c 100644 --- a/tests/model_components/test_renderers.py +++ b/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