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

add option to use random bg color + bg mlp

Browse Source
This commit is contained in:
terrance 2023-02-13 21:05:55 +00:00
parent 7c40594c37
commit 69b98459a5
2 changed files with 14 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
nerfstudio/configs/method_configs.py
View File

@ -356,6 +356,7 @@ method_configs["dreamfusion"] = DreamfusionTrainerConfig(
distortion_loss_mult=10,
sphere_collider=True,
initialize_density=False,
random_background=True,
interlevel_loss_mult=1.0,
proposal_warmup=500,
proposal_update_every=5,

15
nerfstudio/models/dreamfusion.py
View File

@ -73,6 +73,8 @@ class DreamFusionModelConfig(ModelConfig):
"""Initial strength of center density"""
sphere_collider: bool = True
"""Use spherical collider instead of box"""
random_background: bool = True
"""Randomly choose between using background mlp and random color for background"""
target_transmittance_start: float = 0.4
"""target transmittance for opacity penalty. This is the percent of the scene that is
background when rendered at the start of training"""
@ -133,6 +135,7 @@ class DreamFusionModel(Model):
self.initialize_density = config.initialize_density
self.train_normals = False
self.train_shaded = False
self.random_background = config.random_background
self.density_strength = 1.0
self.target_transmittance = config.target_transmittance_start
super().__init__(config=config, **kwargs)
@ -336,7 +339,7 @@ class DreamFusionModel(Model):
shading_weight = 0.0
shaded, shaded_albedo = self.shader_lambertian(
rgb=rgb, normals=pred_normals, light_direction=light_d, shading_weight=shading_weight, detach_normals=False
rgb=rgb, normals=pred_normals, light_direction=light_d, shading_weight=shading_weight, detach_normals=True
)
outputs["normals"] = self.shader_normals(normals, weights=accum_mask)
@ -347,8 +350,16 @@ class DreamFusionModel(Model):
outputs["rgb"] = accum_mask * rgb + background
if shading_weight > 0:
if np.random.random_sample() > 0.5 and not self.training:
samp = np.random.random_sample()
if samp > 0.5 and not self.training:
outputs["train_output"] = outputs["shaded"]
elif samp < 0.2 and self.random_background:
rand_bg = torch.zeros(background.shape)
rand_bg.index_fill_(1, torch.tensor([0]), np.random.random_sample())
rand_bg.index_fill_(1, torch.tensor([1]), np.random.random_sample())
rand_bg.index_fill_(1, torch.tensor([2]), np.random.random_sample())
rand_bg = rand_bg.to(self.device)
outputs["train_output"] = accum_mask * shaded_albedo + rand_bg * accum_mask_inv
else:
outputs["train_output"] = accum_mask * shaded_albedo + background
else: