tink/nerfstudio
4
0
Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 18:44:21 +08:00
Code

visibility depth nerfacto

Browse Source
This commit is contained in:
Ethan Weber 2024-03-05 01:57:05 +00:00
parent 1717a2fc2f
commit 27aa29988a
3 changed files with 30 additions and 26 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

53
nerfstudio/models/depth_nerfacto.py
View File

@ -71,10 +71,6 @@ class DepthNerfactoModel(NerfactoModel):
else:
self.depth_sigma = torch.tensor([self.config.depth_sigma])
# self.collider = NearFarCollider(near_plane=self.config.near_plane, far_plane=self.config.far_plane)
# import pdb; pdb.set_trace();
self.collider = AABBBoxCollider(scene_box=self.scene_box)
def get_outputs(self, ray_bundle: RayBundle):
outputs = super().get_outputs(ray_bundle)
if ray_bundle.metadata is not None and "directions_norm" in ray_bundle.metadata:
@ -91,27 +87,34 @@ class DepthNerfactoModel(NerfactoModel):
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["expected_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}")
# 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["expected_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}")
termination_depth = batch["depth_image"].to(self.device)
predicted_depth = outputs["expected_depth"]
# import pdb; pdb.set_trace();
if not self.config.is_euclidean_depth:
termination_depth = termination_depth * outputs["directions_norm"]
metrics_dict["depth_loss"] = torch.mean((termination_depth - predicted_depth) ** 2)
return metrics_dict

2
nerfstudio/models/nerfacto.py
View File

@ -179,7 +179,7 @@ class NerfactoModel(Model):
implementation=self.config.implementation,
)
self.camera_optimizer: CameraOptimizer = self.config.camera_optimizer.setup(
num_cameras=self.num_train_data, device="cpu"
num_cameras=self.num_train_data, device="cpu")
# this can be set by the pipeline
# if set, then we create a visibility outputs
self.visibility_field: Optional[VisibilityField] = None

1
nerfstudio/pipelines/base_pipeline.py
View File

@ -44,6 +44,7 @@ from nerfstudio.cameras.cameras import Cameras
from nerfstudio.fields.visibility_field import VisibilityField
from nerfstudio.utils.misc import step_check
from nerfstudio.data.datamanagers.random_cameras_datamanager import random_train_pose
import math
def module_wrapper(ddp_or_model: Union[DDP, Model]) -> Model: