From 8c9c2ba2ab052317ff47f2f93d0ff5e0c7655471 Mon Sep 17 00:00:00 2001
From: Yiming Zhou <107962411+yimingzhou1@users.noreply.github.com>
Date: Tue, 7 Feb 2023 16:26:05 -0800
Subject: [PATCH] Add support for sparse depth maps (#1373)

Add support for sparse depth maps (#2)

* sparse depth maps

* pr comments

* unified loss

* minor

* lint
---
 nerfstudio/model_components/losses.py | 18 ++++++++++++------
 nerfstudio/models/depth_nerfacto.py   |  9 ++++++---
 2 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/nerfstudio/model_components/losses.py b/nerfstudio/model_components/losses.py
index 6a981d664..353dd7db9 100644
--- a/nerfstudio/model_components/losses.py
+++ b/nerfstudio/model_components/losses.py
@@ -34,7 +34,7 @@ EPS = 1.0e-7
 URF_SIGMA_SCALE_FACTOR = 3.0
 
 
-class DephtLossType(Enum):
+class DepthLossType(Enum):
     """Types of depth losses for depth supervision."""
 
     DS_NERF = 1
@@ -225,8 +225,11 @@ def ds_nerf_depth_loss(
     Returns:
         Depth loss scalar.
     """
+    depth_mask = termination_depth > 0
+
     loss = -torch.log(weights + EPS) * torch.exp(-((steps - termination_depth[:, None]) ** 2) / (2 * sigma)) * lengths
-    return torch.mean(loss.sum(-2))
+    loss = loss.sum(-2) * depth_mask
+    return torch.mean(loss)
 
 
 def urban_radiance_field_depth_loss(
@@ -247,6 +250,8 @@ def urban_radiance_field_depth_loss(
     Returns:
         Depth loss scalar.
     """
+    depth_mask = termination_depth > 0
+
     # Expected depth loss
     expected_depth_loss = (termination_depth - predicted_depth) ** 2
 
@@ -262,7 +267,8 @@ def urban_radiance_field_depth_loss(
     line_of_sight_loss_empty = (line_of_sight_loss_empty_mask * weights**2).sum(-2)
     line_of_sight_loss = line_of_sight_loss_near + line_of_sight_loss_empty
 
-    return torch.mean(expected_depth_loss + line_of_sight_loss)
+    loss = (expected_depth_loss + line_of_sight_loss) * depth_mask
+    return torch.mean(loss)
 
 
 def depth_loss(
@@ -273,7 +279,7 @@ def depth_loss(
     sigma: TensorType[0],
     directions_norm: TensorType[..., 1],
     is_euclidean: bool,
-    depth_loss_type: DephtLossType,
+    depth_loss_type: DepthLossType,
 ) -> TensorType[0]:
     """Implementation of depth losses.
 
@@ -294,11 +300,11 @@ def depth_loss(
         termination_depth = termination_depth * directions_norm
     steps = (ray_samples.frustums.starts + ray_samples.frustums.ends) / 2
 
-    if depth_loss_type == DephtLossType.DS_NERF:
+    if depth_loss_type == DepthLossType.DS_NERF:
         lengths = ray_samples.frustums.ends - ray_samples.frustums.starts
         return ds_nerf_depth_loss(weights, termination_depth, steps, lengths, sigma)
 
-    if depth_loss_type == DephtLossType.URF:
+    if depth_loss_type == DepthLossType.URF:
         return urban_radiance_field_depth_loss(weights, termination_depth, predicted_depth, steps, sigma)
 
     raise NotImplementedError("Provided depth loss type not implemented.")
diff --git a/nerfstudio/models/depth_nerfacto.py b/nerfstudio/models/depth_nerfacto.py
index ac0768fdf..2bea84c7b 100644
--- a/nerfstudio/models/depth_nerfacto.py
+++ b/nerfstudio/models/depth_nerfacto.py
@@ -24,7 +24,7 @@ from typing import Dict, Tuple, Type
 import torch
 
 from nerfstudio.cameras.rays import RayBundle
-from nerfstudio.model_components.losses import DephtLossType, depth_loss
+from nerfstudio.model_components.losses import DepthLossType, depth_loss
 from nerfstudio.models.nerfacto import NerfactoModel, NerfactoModelConfig
 from nerfstudio.utils import colormaps
 
@@ -46,7 +46,7 @@ class DepthNerfactoModelConfig(NerfactoModelConfig):
     """Starting uncertainty around depth values in meters (defaults to 0.2m)."""
     sigma_decay_rate: float = 0.99985
     """Rate of exponential decay."""
-    depth_loss_type: DephtLossType = DephtLossType.DS_NERF
+    depth_loss_type: DepthLossType = DepthLossType.DS_NERF
     """Depth loss type."""
 
 
@@ -118,7 +118,10 @@ class DepthNerfactoModel(NerfactoModel):
             far_plane=torch.max(ground_truth_depth),
         )
         images["depth"] = torch.cat([ground_truth_depth_colormap, predicted_depth_colormap], dim=1)
-        metrics["depth_mse"] = torch.nn.functional.mse_loss(outputs["depth"], ground_truth_depth)
+        depth_mask = ground_truth_depth > 0
+        metrics["depth_mse"] = torch.nn.functional.mse_loss(
+            outputs["depth"][depth_mask], ground_truth_depth[depth_mask]
+        )
         return metrics, images
 
     def _get_sigma(self):