diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index f0eeb80e2..1359b4c27 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -80,6 +80,40 @@ method_configs["nerfacto"] = TrainerConfig( steps_per_save=2000, max_num_iterations=30000, mixed_precision=True, + pipeline=VanillaPipelineConfig( + datamanager=VanillaDataManagerConfig( + dataparser=NerfstudioDataParserConfig(), + train_num_rays_per_batch=4096, + eval_num_rays_per_batch=4096, + camera_optimizer=CameraOptimizerConfig( + mode="SO3xR3", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2) + ), + ), + model=NerfactoModelConfig( + eval_num_rays_per_chunk=1 << 15, + robust=False, + ), + ), + optimizers={ + "proposal_networks": { + "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15), + "scheduler": None, + }, + "fields": { + "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15), + "scheduler": None, + }, + }, + viewer=ViewerConfig(num_rays_per_chunk=1 << 15), + vis="viewer", +) + +method_configs["nerfacto-robust"] = TrainerConfig( + method_name="nerfacto-robust", + steps_per_eval_batch=500, + steps_per_save=2000, + max_num_iterations=30000, + mixed_precision=True, pipeline=VanillaPipelineConfig( datamanager=VanillaDataManagerConfig( dataparser=NerfstudioDataParserConfig(), @@ -89,7 +123,11 @@ method_configs["nerfacto"] = TrainerConfig( mode="SO3xR3", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2) ), ), - model=NerfactoModelConfig(eval_num_rays_per_chunk=1 << 15), + model=NerfactoModelConfig( + eval_num_rays_per_chunk=1 << 15, + robust=True, + patch_size=16 + ), ), optimizers={ "proposal_networks": { diff --git a/nerfstudio/model_components/losses.py b/nerfstudio/model_components/losses.py index b65a361a0..792d21e58 100644 --- a/nerfstudio/model_components/losses.py +++ b/nerfstudio/model_components/losses.py @@ -303,7 +303,7 @@ def depth_loss( raise NotImplementedError("Provided depth loss type not implemented.") -def robust_rgb_loss(rgb_gt,rgb_outputs,patch_size=16): +def robust_rgb_loss(rgb_gt,rgb_outputs,patch_size): out_patches = rgb_outputs.view(-1,1,patch_size,patch_size,3) gt_patches = rgb_gt.view(-1,1,patch_size,patch_size,3) device = rgb_outputs.device diff --git a/nerfstudio/models/nerfacto.py b/nerfstudio/models/nerfacto.py index 896964bd1..2a69c0dab 100644 --- a/nerfstudio/models/nerfacto.py +++ b/nerfstudio/models/nerfacto.py @@ -116,6 +116,10 @@ class NerfactoModelConfig(ModelConfig): """Whether use single jitter or not for the proposal networks.""" predict_normals: bool = False """Whether to predict normals or not.""" + robust: bool = False + """Whether to use robust RGB loss or not.""" + patch_size: int = 16 + """Patch size to use for robust RGB loss.""" class NerfactoModel(Model): @@ -301,7 +305,10 @@ class NerfactoModel(Model): def get_loss_dict(self, outputs, batch, metrics_dict=None): loss_dict = {} image = batch["image"].to(self.device) - loss_dict["rgb_loss"] = robust_rgb_loss(image, outputs["rgb"]) + if self.config.robust: + loss_dict["rgb_loss"] = robust_rgb_loss(image, outputs["rgb"], self.config.patch_size) + else: + loss_dict["rgb_loss"] = self.rgb_loss(image, outputs["rgb"]) if self.training: loss_dict["interlevel_loss"] = self.config.interlevel_loss_mult * interlevel_loss( outputs["weights_list"], outputs["ray_samples_list"]