diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index db903c4a2..f0815b51f 100644 --- a/nerfstudio/configs/method_configs.py +++ b/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, diff --git a/nerfstudio/configs/sampling_config.py b/nerfstudio/configs/sampling_config.py new file mode 100644 index 000000000..bc04ee804 --- /dev/null +++ b/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""" diff --git a/nerfstudio/data/datasets/depth_dataset.py b/nerfstudio/data/datasets/depth_dataset.py index dff82314a..948dc01a5 100644 --- a/nerfstudio/data/datasets/depth_dataset.py +++ b/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"] diff --git a/nerfstudio/model_components/losses.py b/nerfstudio/model_components/losses.py index e948cbfdb..9823e76e3 100644 --- a/nerfstudio/model_components/losses.py +++ b/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): diff --git a/nerfstudio/models/depth_nerfacto.py b/nerfstudio/models/depth_nerfacto.py index e3328fd6d..6bfcc3ba2 100644 --- a/nerfstudio/models/depth_nerfacto.py +++ b/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 diff --git a/nerfstudio/models/nerfacto.py b/nerfstudio/models/nerfacto.py index 175d671cc..a4cfb342b 100644 --- a/nerfstudio/models/nerfacto.py +++ b/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