From 1f63072852cba3cc583e6f52eb7fbf20d1ba9197 Mon Sep 17 00:00:00 2001 From: Justin Kerr Date: Mon, 27 May 2024 11:03:39 -0700 Subject: [PATCH] alsdjf --- nerfstudio/configs/method_configs.py | 5 +++-- .../data/datamanagers/full_images_datamanager.py | 6 ++++-- nerfstudio/models/splatfacto.py | 16 +++++++++++----- 3 files changed, 18 insertions(+), 9 deletions(-) diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 90a816f4c..039138c63 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -27,7 +27,7 @@ from nerfstudio.cameras.camera_optimizers import CameraOptimizerConfig from nerfstudio.configs.base_config import ViewerConfig from nerfstudio.configs.external_methods import ExternalMethodDummyTrainerConfig, get_external_methods from nerfstudio.data.datamanagers.base_datamanager import VanillaDataManager, VanillaDataManagerConfig -from nerfstudio.data.datamanagers.full_images_datamanager import FullImageDatamanagerConfig +from nerfstudio.data.datamanagers.full_images_datamanager import FullImageDatamanagerConfig,FullImageDatamanager from nerfstudio.data.datamanagers.parallel_datamanager import ParallelDataManagerConfig from nerfstudio.data.datamanagers.random_cameras_datamanager import RandomCamerasDataManagerConfig from nerfstudio.data.dataparsers.blender_dataparser import BlenderDataParserConfig @@ -601,7 +601,8 @@ method_configs["splatfacto"] = TrainerConfig( mixed_precision=False, pipeline=VanillaPipelineConfig( datamanager=FullImageDatamanagerConfig( - dataparser=NerfstudioDataParserConfig(load_3D_points=True), + _target=FullImageDatamanager[DepthDataset], + dataparser=NerfstudioDataParserConfig(load_3D_points=True,depth_unit_scale_factor=1.0), cache_images_type="uint8", ), model=SplatfactoModelConfig(), diff --git a/nerfstudio/data/datamanagers/full_images_datamanager.py b/nerfstudio/data/datamanagers/full_images_datamanager.py index 5ea830f12..1440e5636 100644 --- a/nerfstudio/data/datamanagers/full_images_datamanager.py +++ b/nerfstudio/data/datamanagers/full_images_datamanager.py @@ -200,14 +200,16 @@ class FullImageDatamanager(DataManager, Generic[TDataset]): cache["image"] = cache["image"].to(self.device) if "mask" in cache: cache["mask"] = cache["mask"].to(self.device) - if "depth" in cache: - cache["depth"] = cache["depth"].to(self.device) + if "depth_image" in cache: + cache["depth_image"] = cache["depth_image"].to(self.device) self.train_cameras = self.train_dataset.cameras.to(self.device) elif cache_images_device == "cpu": for cache in undistorted_images: cache["image"] = cache["image"].pin_memory() if "mask" in cache: cache["mask"] = cache["mask"].pin_memory() + if "depth_image" in cache: + cache["depth_image"] = cache["depth_image"].to(self.device) self.train_cameras = self.train_dataset.cameras else: assert_never(cache_images_device) diff --git a/nerfstudio/models/splatfacto.py b/nerfstudio/models/splatfacto.py index 534c23ef7..14f26988e 100644 --- a/nerfstudio/models/splatfacto.py +++ b/nerfstudio/models/splatfacto.py @@ -98,7 +98,7 @@ def resize_image(image: torch.Tensor, d: int): return tf.conv2d(image.permute(2, 0, 1)[:, None, ...], weight, stride=d).squeeze(1).permute(1, 2, 0) -@torch.compile(dynamic=False) +@torch.compile() def get_viewmat(optimized_camera_to_world): """ function that converts c2w to gsplat world2camera matrix, using compile for some speed @@ -171,7 +171,7 @@ class SplatfactoModelConfig(ModelConfig): """threshold of ratio of gaussian max to min scale before applying regularization loss from the PhysGaussian paper """ - output_depth_during_training: bool = False + output_depth_during_training: bool = True """If True, output depth during training. Otherwise, only output depth during evaluation.""" rasterize_mode: Literal["classic", "antialiased"] = "classic" """ @@ -723,7 +723,7 @@ class SplatfactoModel(Model): raise ValueError("Unknown rasterize_mode: %s", self.config.rasterize_mode) if self.config.output_depth_during_training or not self.training: - render_mode = "RGB+D" + render_mode = "RGB+ED" else: render_mode = "RGB" @@ -761,7 +761,7 @@ class SplatfactoModel(Model): alpha = alpha[0, ...] rgb = render[0, ..., :3] + (1 - alpha) * background rgb = torch.clamp(rgb, 0.0, 1.0) - if render_mode == "RGB+D": + if render_mode == "RGB+ED": depth_im = render[0, ..., 3:4] depth_im = torch.where(alpha > 0, depth_im, depth_im.detach().max()) else: @@ -845,10 +845,16 @@ class SplatfactoModel(Model): scale_reg = 0.1 * scale_reg.mean() else: scale_reg = torch.tensor(0.0).to(self.device) - + if 'depth_image' in batch: + gt_depth = self.get_gt_img(batch['depth_image']) + ignores = gt_depth == 0 + depth_loss = (gt_depth - outputs['depth'])[~ignores].abs().mean() + else: + depth_loss = torch.tensor(0.0).to(self.device) loss_dict = { "main_loss": (1 - self.config.ssim_lambda) * Ll1 + self.config.ssim_lambda * simloss, "scale_reg": scale_reg, + 'depth_loss':depth_loss } if self.training: