Enable SparseNerf depth ranking loss to be used with GT/Sensor depth (#2312)

* add ranking loss to depth-nerfacto with GT depth

* formating

* pair pixel sampler bug fix from pr #2315

---------

Co-authored-by: Sasha Sax <asax@stanford.edu>
Co-authored-by: Ethan Weber <ethanweber@berkeley.edu>
Co-authored-by: AdamRashid96 <abrashid@berkeley.edu>
Co-authored-by: AdamRashid96 <71362382+AdamRashid96@users.noreply.github.com>

nerfstudio/data/datasets/depth_dataset.py

@@ -46,16 +46,20 @@ class DepthDataset(InputDataset):
     def __init__(self, dataparser_outputs: DataparserOutputs, scale_factor: float = 1.0):
         super().__init__(dataparser_outputs, scale_factor)
         # if there are no depth images than we want to generate them all with zoe depth
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
         if len(dataparser_outputs.image_filenames) > 0 and (
             "depth_filenames" not in dataparser_outputs.metadata.keys()
             or dataparser_outputs.metadata["depth_filenames"] is None
         ):
-            losses.DEPTH_METRIC = 0
-            CONSOLE.print("[bold yellow] No depth data found!")
+            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+            CONSOLE.print("[bold yellow] No depth data found! Generating pseudodepth...")
+            losses.FORCE_PSEUDODEPTH_LOSS = True
+            CONSOLE.print("[bold red] Using psueodepth: forcing depth loss to be ranking loss.")
             cache = dataparser_outputs.image_filenames[0].parent / "depths.npy"
+            # Note: this should probably be saved to disk as images, and then loaded with the dataparser.
+            #  That will allow multi-gpu training.
             if cache.exists():
-                CONSOLE.print("[bold yellow] Loading cache!")
+                CONSOLE.print("[bold yellow] Loading pseudodata depth from cache!")
                 # load all the depths
                 self.depths = np.load(cache)
                 self.depths = torch.from_numpy(self.depths).to(device)
@@ -98,9 +102,6 @@ class DepthDataset(InputDataset):
             self.metadata["depth_filenames"] = None
             self.metadata["depth_unit_scale_factor"] = 1.0
 
-        else:
-            losses.DEPTH_METRIC = 0
-
         self.depth_filenames = self.metadata["depth_filenames"]
         self.depth_unit_scale_factor = self.metadata["depth_unit_scale_factor"]
 

nerfstudio/data/pixel_samplers.py

@@ -386,8 +386,8 @@ class PairPixelSampler(PixelSampler):  # pylint: disable=too-few-public-methods
     def __init__(self, config: PairPixelSamplerConfig, **kwargs) -> None:
         self.config = config
         self.radius = self.config.radius
-        self.rays_to_sample = self.config.num_rays_per_batch // 2
         super().__init__(self.config, **kwargs)
+        self.rays_to_sample = self.config.num_rays_per_batch // 2
 
     # overrides base method
     def sample_method(  # pylint: disable=no-self-use

nerfstudio/model_components/losses.py

@@ -35,7 +35,6 @@ EPS = 1.0e-7
 
 # Sigma scale factor from Urban Radiance Fields (Rematas et al., 2022)
 URF_SIGMA_SCALE_FACTOR = 3.0
-DEPTH_METRIC = 1
 
 
 class DepthLossType(Enum):
@@ -43,6 +42,11 @@ class DepthLossType(Enum):
 
     DS_NERF = 1
     URF = 2
+    SPARSENERF_RANKING = 3
+
+
+FORCE_PSEUDODEPTH_LOSS = False
+PSEUDODEPTH_COMPATIBLE_LOSSES = (DepthLossType.SPARSENERF_RANKING,)
 
 
 def outer(

nerfstudio/models/depth_nerfacto.py

@@ -79,36 +79,49 @@ class DepthNerfactoModel(NerfactoModel):
     def get_metrics_dict(self, outputs, batch):
         metrics_dict = super().get_metrics_dict(outputs, batch)
         if self.training:
-            metrics_dict["depth_loss"] = 0.0
-            sigma = self._get_sigma().to(self.device)
-            termination_depth = batch["depth_image"].to(self.device)
-            for i in range(len(outputs["weights_list"])):
-                metrics_dict["depth_loss"] += depth_loss(
-                    weights=outputs["weights_list"][i],
-                    ray_samples=outputs["ray_samples_list"][i],
-                    termination_depth=termination_depth,
-                    predicted_depth=outputs["depth"],
-                    sigma=sigma,
-                    directions_norm=outputs["directions_norm"],
-                    is_euclidean=self.config.is_euclidean_depth,
-                    depth_loss_type=self.config.depth_loss_type,
-                ) / len(outputs["weights_list"])
+            if (
+                losses.FORCE_PSEUDODEPTH_LOSS
+                and self.config.depth_loss_type not in losses.PSEUDODEPTH_COMPATIBLE_LOSSES
+            ):
+                raise ValueError(
+                    f"Forcing pseudodepth loss, but depth loss type ({self.config.depth_loss_type}) must be one of {losses.PSEUDODEPTH_COMPATIBLE_LOSSES}"
+                )
+            if self.config.depth_loss_type in (DepthLossType.DS_NERF, DepthLossType.URF):
+                metrics_dict["depth_loss"] = 0.0
+                sigma = self._get_sigma().to(self.device)
+                termination_depth = batch["depth_image"].to(self.device)
+                for i in range(len(outputs["weights_list"])):
+                    metrics_dict["depth_loss"] += depth_loss(
+                        weights=outputs["weights_list"][i],
+                        ray_samples=outputs["ray_samples_list"][i],
+                        termination_depth=termination_depth,
+                        predicted_depth=outputs["depth"],
+                        sigma=sigma,
+                        directions_norm=outputs["directions_norm"],
+                        is_euclidean=self.config.is_euclidean_depth,
+                        depth_loss_type=self.config.depth_loss_type,
+                    ) / len(outputs["weights_list"])
+            elif self.config.depth_loss_type in (DepthLossType.SPARSENERF_RANKING,):
+                metrics_dict["depth_ranking"] = depth_ranking_loss(
+                    outputs["expected_depth"], batch["depth_image"].to(self.device)
+                )
+            else:
+                raise NotImplementedError(f"Unknown depth loss type {self.config.depth_loss_type}")
 
         return metrics_dict
 
     def get_loss_dict(self, outputs, batch, metrics_dict=None):
         loss_dict = super().get_loss_dict(outputs, batch, metrics_dict)
         if self.training:
-            assert metrics_dict is not None and "depth_loss" in metrics_dict
-            # If real depth used
-            if losses.DEPTH_METRIC:
-                loss_dict["depth_loss"] = self.config.depth_loss_mult * metrics_dict["depth_loss"]
-            else:
-                loss_dict["depth_ranking"] = np.interp(self.step, [0, 2000], [0, 0.2]) * depth_ranking_loss(
-                    outputs["expected_depth"], batch["depth_image"]
+            assert metrics_dict is not None and ("depth_loss" in metrics_dict or "depth_ranking" in metrics_dict)
+            if "depth_ranking" in metrics_dict:
+                loss_dict["depth_ranking"] = (
+                    self.config.depth_loss_mult
+                    * np.interp(self.step, [0, 2000], [0, 0.2])
+                    * metrics_dict["depth_ranking"]
                 )
-            # Else use depth ranking loss
-
+            if "depth_loss" in metrics_dict:
+                loss_dict["depth_loss"] = self.config.depth_loss_mult * metrics_dict["depth_loss"]
         return loss_dict
 
     def get_image_metrics_and_images(