tink/nerfstudio
4
0
Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 10:34:29 +08:00
Code

merging

Browse Source
This commit is contained in:
terrance 2023-02-16 00:55:49 +00:00
commit ce0d97f811
5 changed files with 44 additions and 32 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
nerfstudio/configs/method_configs.py
View File

@ -341,7 +341,7 @@ method_configs["dreamfusion"] = TrainerConfig(
method_name="dreamfusion",
steps_per_eval_batch=50,
steps_per_eval_image=50,
steps_per_save=2000,
steps_per_save=200,
max_num_iterations=30000,
mixed_precision=True,
pipeline=DreamfusionPipelineConfig(
@ -352,7 +352,9 @@ method_configs["dreamfusion"] = TrainerConfig(
),
model=DreamFusionModelConfig(
eval_num_rays_per_chunk=1 << 15,
distortion_loss_mult=10,
distortion_loss_mult=10.0,
orientation_loss_mult=0.1,
max_res=256,
sphere_collider=True,
initialize_density=False,
random_background=True,
@ -362,9 +364,9 @@ method_configs["dreamfusion"] = TrainerConfig(
proposal_weights_anneal_max_num_iters=100,
start_lambertian_training=1000,
start_normals_training=500,
opacity_loss_mult=0.001,
),
alphas_loss_mult=1.0,
opacity_loss_mult=1.0,
interpolated_prompting=False,
guidance_scale=100,
),
optimizers={

6
nerfstudio/fields/dreamfusion_field.py
View File

@ -65,6 +65,9 @@ class DreamFusionField(Field):
num_layers: int = 2,
hidden_dim: int = 64,
geo_feat_dim: int = 15,
num_levels: int = 16,
max_res: int = 256,
log2_hashmap_size: int = 19,
num_layers_color: int = 3,
hidden_dim_color: int = 64,
hidden_dim_normal: int = 32,
@ -74,10 +77,7 @@ class DreamFusionField(Field):
self.aabb = Parameter(aabb, requires_grad=False)
self.geo_feat_dim = geo_feat_dim
num_levels = 16
max_res = 1024
base_res = 16
log2_hashmap_size = 19
features_per_level = 2
growth_factor = np.exp((np.log(max_res) - np.log(base_res)) / (num_levels - 1))

17
nerfstudio/model_components/ray_samplers.py
View File

@ -509,7 +509,16 @@ class VolumetricSampler(Sampler):
class ProposalNetworkSampler(Sampler):
"""Sampler that uses a proposal network to generate samples."""
"""Sampler that uses a proposal network to generate samples.
Args:
num_proposal_samples_per_ray: Number of samples to generate per ray for each proposal step.
num_nerf_samples_per_ray: Number of samples to generate per ray for the NERF model.
num_proposal_network_iterations: Number of proposal network iterations to run.
single_jitter: Use a same random jitter for all samples along a ray.
update_sched: A function that takes the iteration number of steps between updates.
initial_sampler: Sampler to use for the first iteration. Uses UniformLinDispPiecewise if not set.
"""
def __init__(
self,
@ -518,6 +527,7 @@ class ProposalNetworkSampler(Sampler):
num_proposal_network_iterations: int = 2,
single_jitter: bool = False,
update_sched: Callable = lambda x: 1,
initial_sampler: Optional[Sampler] = None,
) -> None:
super().__init__()
self.num_proposal_samples_per_ray = num_proposal_samples_per_ray
@ -528,7 +538,10 @@ class ProposalNetworkSampler(Sampler):
raise ValueError("num_proposal_network_iterations must be >= 1")
# samplers
self.initial_sampler = UniformLinDispPiecewiseSampler(single_jitter=single_jitter)
if initial_sampler is None:
self.initial_sampler = UniformLinDispPiecewiseSampler(single_jitter=single_jitter)
else:
self.initial_sampler = initial_sampler
self.pdf_sampler = PDFSampler(include_original=False, single_jitter=single_jitter)
self._anneal = 1.0

35
nerfstudio/models/dreamfusion.py
View File

@ -41,7 +41,10 @@ from nerfstudio.model_components.losses import (
orientation_loss,
pred_normal_loss,
)
from nerfstudio.model_components.ray_samplers import ProposalNetworkSampler
from nerfstudio.model_components.ray_samplers import (
ProposalNetworkSampler,
UniformSampler,
)
from nerfstudio.model_components.renderers import (
AccumulationRenderer,
DepthRenderer,
@ -118,6 +121,12 @@ class DreamFusionModelConfig(ModelConfig):
"""Start training normals after this many iterations"""
start_lambertian_training: int = 1000
"""start training with lambertian shading after this many iterations"""
opacity_penalty: bool = True
"""enables penalty to encourage sparse weights (penalizing for uniform density along ray)"""
opacity_loss_mult: float = 1
"""scale for opacity penalty"""
max_res: int = 256
"""Maximum resolution of the density field."""
class DreamFusionModel(Model):
@ -148,7 +157,7 @@ class DreamFusionModel(Model):
super().populate_modules()
# setting up fields
self.field = DreamFusionField(self.scene_box.aabb)
self.field = DreamFusionField(self.scene_box.aabb, max_res=self.config.max_res)
# samplers
self.density_fns = []
@ -176,10 +185,9 @@ class DreamFusionModel(Model):
num_proposal_network_iterations=self.config.num_proposal_iterations,
single_jitter=self.config.use_single_jitter,
update_sched=update_schedule,
initial_sampler=UniformSampler(single_jitter=self.config.use_single_jitter),
)
# self.sampler_uniform = UniformSampler(num_samples=self.num_samples, single_jitter=True)
# renderers
self.renderer_rgb = RGBRenderer(background_color=colors.WHITE)
self.renderer_accumulation = AccumulationRenderer()
@ -357,11 +365,7 @@ class DreamFusionModel(Model):
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)
rand_bg = torch.ones_like(background) * torch.rand(3, device=self.device)
outputs["train_output"] = accum_mask * shaded_albedo + rand_bg * accum_mask_inv
else:
outputs["train_output"] = accum_mask * shaded_albedo + background
@ -370,7 +374,7 @@ class DreamFusionModel(Model):
if self.training or True:
outputs["rendered_orientation_loss"] = orientation_loss(
weights, field_outputs[FieldHeadNames.NORMALS], ray_bundle.directions
weights.detach(), field_outputs[FieldHeadNames.NORMALS], ray_bundle.directions
)
outputs["rendered_pred_normal_loss"] = pred_normal_loss(
@ -380,7 +384,10 @@ class DreamFusionModel(Model):
)
assert weights.shape[-1] == 1
outputs["alphas_loss"] = torch.sqrt(torch.sum(weights, dim=-2) ** 2 + 0.01)
if self.config.opacity_penalty:
outputs["opacity_loss"] = (
torch.sqrt(torch.sum(weights, dim=-2) ** 2 + 0.01) * self.config.opacity_loss_mult
)
return outputs
@ -402,11 +409,9 @@ class DreamFusionModel(Model):
loss_dict["orientation_loss"] = 0
loss_dict["pred_normal_loss"] = 0
loss_dict["alphas_loss"] = outputs["alphas_loss"].mean()
if self.config.opacity_penalty:
loss_dict["opacity_loss"] = self.config.opacity_loss_mult * outputs["opacity_loss"].mean()
loss_dict["opacity_loss"] = -torch.minimum(
(1 - outputs["accumulation"]).mean(), torch.tensor(self.target_transmittance)
)
if self.training:
loss_dict["distortion_loss"] = self.config.distortion_loss_mult * distortion_loss(
outputs["weights_list"], outputs["ray_samples_list"]

8
nerfstudio/pipelines/dreamfusion_pipeline.py
View File

@ -51,14 +51,6 @@ class DreamfusionPipelineConfig(VanillaPipelineConfig):
"""appended to prompt for front view"""
back_prompt: str = ", back view"
"""appended to prompt for back view"""
alphas_penalty: bool = True
"""enables penalty to encourage sparse weights (penalizing for uniform density along ray)"""
alphas_loss_mult: float = 1
"""scale for alphas penalty"""
opacity_penalty: bool = True
"""enables penalty to encourage transparent scenes, as in "dreamfields" paper"""
opacity_loss_mult: float = 1
"""scale for opacity penalty"""
guidance_scale: float = 100
"""guidance scale for sds loss"""