diff --git a/nerfstudio/models/dreamfusion.py b/nerfstudio/models/dreamfusion.py index f5869c5b9..62b9a3bf7 100644 --- a/nerfstudio/models/dreamfusion.py +++ b/nerfstudio/models/dreamfusion.py @@ -29,6 +29,11 @@ from torchmetrics.functional import structural_similarity_index_measure from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity from nerfstudio.cameras.rays import RayBundle +from nerfstudio.engine.callbacks import ( + TrainingCallback, + TrainingCallbackAttributes, + TrainingCallbackLocation, +) # from nerfstudio.field_components.encodings import NeRFEncoding, from nerfstudio.field_components.field_heads import FieldHeadNames @@ -63,7 +68,7 @@ class DreamFusionModelConfig(ModelConfig): """Orientation loss multipier on computed normals.""" pred_normal_loss_mult: float = 0.001 """Predicted normal loss multiplier.""" - random_light_source: bool = True + random_light_source: bool = False """Randomizes light source per output.""" @@ -82,8 +87,29 @@ class DreamFusionModel(Model): **kwargs, ) -> None: self.num_samples = config.num_samples + self.initialize_density = True super().__init__(config=config, **kwargs) + def get_training_callbacks( + self, training_callback_attributes: TrainingCallbackAttributes + ) -> List[TrainingCallback]: + + # the callback that we want to run every X iterations after the training iteration + def stop_initialize_density( + self, training_callback_attributes: TrainingCallbackAttributes, step: int # pylint: disable=unused-argument + ): + self.initialize_density = False + + callbacks = [ + TrainingCallback( + where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION], + iters=(1000,), + func=stop_initialize_density, + args=[self, training_callback_attributes], + ) + ] + return callbacks + def populate_modules(self): """Set the fields and modules""" super().populate_modules() @@ -123,7 +149,14 @@ class DreamFusionModel(Model): # uniform sampling ray_samples_uniform = self.sampler_uniform(ray_bundle) field_outputs = self.field(ray_samples_uniform, compute_normals=True) - weights = ray_samples_uniform.get_weights(field_outputs[FieldHeadNames.DENSITY]) + + density = field_outputs[FieldHeadNames.DENSITY] + if self.initialize_density: + pos = ray_samples_uniform.frustums.get_positions() + density_blob = 10 * torch.exp(-torch.norm(pos, dim=-1) / (2 * 0.04))[..., None] + density += density_blob + + weights = ray_samples_uniform.get_weights(density) accumulation = self.renderer_accumulation(weights) depth = self.renderer_depth(weights, ray_samples_uniform)