diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 947dad880..f5eec1086 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -83,7 +83,7 @@ method_configs["nerfacto"] = TrainerConfig( pipeline=VanillaPipelineConfig( datamanager=VanillaDataManagerConfig( dataparser=NerfstudioDataParserConfig(), - train_num_rays_per_batch=4096, + train_num_rays_per_batch=8192, eval_num_rays_per_batch=4096, camera_optimizer=CameraOptimizerConfig( mode="SO3xR3", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2) diff --git a/nerfstudio/data/datamanagers/base_datamanager.py b/nerfstudio/data/datamanagers/base_datamanager.py index 8114cf6cb..e9dd0f3fe 100644 --- a/nerfstudio/data/datamanagers/base_datamanager.py +++ b/nerfstudio/data/datamanagers/base_datamanager.py @@ -351,7 +351,7 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method # Otherwise, use the default pixel sampler if is_equirectangular.any(): CONSOLE.print("[bold yellow]Warning: Some cameras are equirectangular, but using default pixel sampler.") - return PatchPixelSampler(*args, **kwargs,patch_size=3) + return PatchPixelSampler(*args, **kwargs,patch_size=16) def setup_train(self): """Sets up the data loaders for training""" diff --git a/nerfstudio/model_components/losses.py b/nerfstudio/model_components/losses.py index 6a981d664..8473c46ce 100644 --- a/nerfstudio/model_components/losses.py +++ b/nerfstudio/model_components/losses.py @@ -302,3 +302,23 @@ def depth_loss( return urban_radiance_field_depth_loss(weights, termination_depth, predicted_depth, steps, sigma) raise NotImplementedError("Provided depth loss type not implemented.") + +def robust_rgb_loss(rgb_outputs,rgb_gt,patch_size=16): + out_patches = rgb_outputs.view(-1,1,patch_size,patch_size) + gt_patches = rgb_gt.view(-1,1,patch_size,patch_size) + device = rgb_outputs.device + batch_size = out_patches.shape[0] + residuals = (out_patches - gt_patches)**2 + + with torch.no_grad(): + med_residual = torch.median(residuals) + #equation 8 + weight = (residuals0.5).float() + expected_w = torch.mean(torch.mean(blurred_w,dim=-1),dim=-1) + weight_r8 = (expected_w > 0.6).float() + w_8x8 = weight_r8[...,None] * torch.ones((batch_size,8,8),device=device) + final_w = torch.zeros((batch_size,16,16),device=device) + final_w[:,4:12,4:12] = w_8x8 + return torch.mean(residuals*final_w) \ No newline at end of file diff --git a/nerfstudio/models/nerfacto.py b/nerfstudio/models/nerfacto.py index a0938455a..f4204840d 100644 --- a/nerfstudio/models/nerfacto.py +++ b/nerfstudio/models/nerfacto.py @@ -45,6 +45,7 @@ from nerfstudio.model_components.losses import ( interlevel_loss, orientation_loss, pred_normal_loss, + robust_rgb_loss ) from nerfstudio.model_components.ray_samplers import ProposalNetworkSampler from nerfstudio.model_components.renderers import ( @@ -303,7 +304,7 @@ 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"] = self.rgb_loss(image, outputs["rgb"]) + loss_dict["rgb_loss"] = robust_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"] @@ -320,10 +321,6 @@ class NerfactoModel(Model): loss_dict["pred_normal_loss"] = self.config.pred_normal_loss_mult * torch.mean( outputs["rendered_pred_normal_loss"] ) - #depth regularity loss - p=3 - patches = outputs['depth'].view(-1,p,p) #Bx3x3 - loss_dict['depth_reg'] = np.interp(self.step,[0,500],[0,1])*torch.mean(torch.sum((patches.view(-1,p*p) - patches[:,p//2,p//2][:,None])**2,dim=1)) return loss_dict