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

first stab at it...

Browse Source
This commit is contained in:
Justin 2023-02-03 22:58:24 +00:00
parent 0b45c8f994
commit 4157e21e28
4 changed files with 24 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
nerfstudio/configs/method_configs.py
View File

@ -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)

2
nerfstudio/data/datamanagers/base_datamanager.py
View File

@ -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"""

20
nerfstudio/model_components/losses.py
View File

@ -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 = (residuals<med_residual).float()
#equation 9
blurred_w = (torch.nn.functional.conv2d(weight,(1/9.)*torch.ones((1,1,3,3),device=device),padding='same')>0.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)

7
nerfstudio/models/nerfacto.py
View File

@ -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