moving to depth nerfacto

nerfstudio/configs/method_configs.py

@@ -173,6 +173,7 @@ method_configs["depth-nerfacto"] = TrainerConfig(
     pipeline=VanillaPipelineConfig(
         datamanager=VanillaDataManagerConfig(
             _target=VanillaDataManager[DepthDataset],
+            patch_size=2,
             dataparser=NerfstudioDataParserConfig(),
             train_num_rays_per_batch=4096,
             eval_num_rays_per_batch=4096,

nerfstudio/configs/sampling_config.py

@@ -0,0 +1,32 @@
+# Copyright 2022 the Regents of the University of California, Nerfstudio Team and contributors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Config used for sampling rays."""
+
+from dataclasses import dataclass
+from typing import Optional
+
+from nerfstudio.configs.base_config import (
+    InstantiateConfig,
+)
+
+
+@dataclass
+class SamplingConfig(InstantiateConfig):
+    """Full config contents for defining a pixel sampler."""
+
+    patch_size: int = 1
+    """Patch size for sampling rays."""
+    method_name: Optional[str] = None
+    """Method name. Required to set in python or via cli"""

nerfstudio/data/datasets/depth_dataset.py

@@ -21,8 +21,11 @@ from typing import Dict
 import numpy as np
 
 from nerfstudio.data.dataparsers.base_dataparser import DataparserOutputs
+from nerfstudio.model_components import losses
 from nerfstudio.data.datasets.base_dataset import InputDataset
 from nerfstudio.data.utils.data_utils import get_depth_image_from_path
+from nerfstudio.utils.rich_utils import CONSOLE
+
 from typing import Union
 from PIL import Image
 import torch
@@ -32,7 +35,7 @@ import json
 
 
 class DepthDataset(InputDataset):
-    """Dataset that returns images and depths.
+    """Dataset that returns images and depths. If no depths are found, then we generate them with Zoe Depth.
 
     Args:
         dataparser_outputs: description of where and how to read input images.
@@ -47,8 +50,11 @@ class DepthDataset(InputDataset):
             "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!")
             cache = dataparser_outputs.image_filenames[0].parent / "depths.npy"
             if cache.exists():
+                CONSOLE.print("[bold yellow] Loading cache!")
                 # load all the depths
                 self.depths = np.load(cache)
                 self.depths = torch.from_numpy(self.depths).to(device)
@@ -86,12 +92,14 @@ class DepthDataset(InputDataset):
 
                 self.depths = torch.stack(depth_tensors)
                 np.save(cache, self.depths.cpu().numpy())
-
             dataparser_outputs.metadata["depth_filenames"] = None
             dataparser_outputs.metadata["depth_unit_scale_factor"] = 1.0
             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/model_components/losses.py

@@ -35,6 +35,7 @@ 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):

nerfstudio/models/depth_nerfacto.py

@@ -22,9 +22,11 @@ from dataclasses import dataclass, field
 from typing import Dict, Tuple, Type
 
 import torch
+import numpy as np
 
 from nerfstudio.cameras.rays import RayBundle
-from nerfstudio.model_components.losses import DepthLossType, depth_loss
+from nerfstudio.model_components import losses
+from nerfstudio.model_components.losses import DepthLossType, depth_loss, depth_ranking_loss
 from nerfstudio.models.nerfacto import NerfactoModel, NerfactoModelConfig
 from nerfstudio.utils import colormaps
 
@@ -98,7 +100,14 @@ class DepthNerfactoModel(NerfactoModel):
         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
-            loss_dict["depth_loss"] = self.config.depth_loss_mult * metrics_dict["depth_loss"]
+            # 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["depth"], batch["depth_image"]
+                )
+            # Else use depth ranking loss
 
         return loss_dict
 

nerfstudio/models/nerfacto.py

@@ -301,8 +301,8 @@ class NerfactoModel(Model):
         outputs = {
             "rgb": rgb,
             "accumulation": accumulation,
-            "depth": depth,
-            "depth_med": depth_med,
+            "expected_depth": depth,
+            "depth": depth_med,
         }
 
         if self.config.predict_normals:
@@ -359,9 +359,8 @@ class NerfactoModel(Model):
                 loss_dict["pred_normal_loss"] = self.config.pred_normal_loss_mult * torch.mean(
                     outputs["rendered_pred_normal_loss"]
                 )
-
-            if "depth_image" in batch and self.step >= 500:
-                loss_dict["depth_ranking"] = np.interp(self.step, [500, 1000], [0, 0.2]) * depth_ranking_loss(
+            if "depth_image" in batch:
+                loss_dict["depth_ranking"] = np.interp(self.step, [0, 2000], [0, 0.2]) * depth_ranking_loss(
                     outputs["depth"], batch["depth_image"]
                 )
         return loss_dict