diff --git a/nerfstudio/data/datasets/depth_dataset.py b/nerfstudio/data/datasets/depth_dataset.py index 0bb33c776..e17ac6315 100644 --- a/nerfstudio/data/datasets/depth_dataset.py +++ b/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"] diff --git a/nerfstudio/data/pixel_samplers.py b/nerfstudio/data/pixel_samplers.py index 30b33bf23..02f1dbc6e 100644 --- a/nerfstudio/data/pixel_samplers.py +++ b/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 diff --git a/nerfstudio/model_components/losses.py b/nerfstudio/model_components/losses.py index 8c53c24f1..e90ec8e61 100644 --- a/nerfstudio/model_components/losses.py +++ b/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( diff --git a/nerfstudio/models/depth_nerfacto.py b/nerfstudio/models/depth_nerfacto.py index f3f791e80..d925c0516 100644 --- a/nerfstudio/models/depth_nerfacto.py +++ b/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(