diff --git a/nerfstudio/models/dreamfusion.py b/nerfstudio/models/dreamfusion.py index 82ccb8c73..d8251b120 100644 --- a/nerfstudio/models/dreamfusion.py +++ b/nerfstudio/models/dreamfusion.py @@ -93,6 +93,7 @@ class DreamFusionModel(Model): self.num_samples = config.num_samples self.initialize_density = config.initialize_density self.train_normals = False + self.train_shaded = False super().__init__(config=config, **kwargs) def get_training_callbacks( @@ -110,6 +111,11 @@ class DreamFusionModel(Model): ): self.train_normals = True + def start_shaded_training( + self, training_callback_attributes: TrainingCallbackAttributes, step: int # pylint: disable=unused-argument + ): + self.train_shaded = True + callbacks = [ TrainingCallback( where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION], @@ -123,6 +129,12 @@ class DreamFusionModel(Model): func=start_training_normals, args=[self, training_callback_attributes], ), + TrainingCallback( + where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION], + iters=(1000,), + func=start_shaded_training, + args=[self, training_callback_attributes], + ), ] return callbacks @@ -154,7 +166,7 @@ class DreamFusionModel(Model): # colliders if self.config.sphere_collider: - self.collider = SphereCollider(torch.Tensor([0, 0, 0]), 0.8) + self.collider = SphereCollider(torch.Tensor([0, 0, 0]), 1.0) else: self.collider = AABBBoxCollider(scene_box=self.scene_box) @@ -206,7 +218,10 @@ class DreamFusionModel(Model): light_d = ray_bundle.origins[0] light_d = math.safe_normalize(light_d) - ratio = 0.3 + if self.train_shaded: + ratio = 0.1 + else: + ratio = 0.0 # with torch.no_grad(): lambertian = ratio + (1 - ratio) * (pred_normals.detach() @ light_d).clamp(min=0) shaded = lambertian.unsqueeze(-1).repeat(1, 3) @@ -218,6 +233,14 @@ class DreamFusionModel(Model): outputs["shaded_albedo"] = shaded_albedo outputs["render"] = accum_mask * shaded_albedo + background + if ratio > 0: + if np.random.random_sample() > 0.5: + outputs["train_output"] = accum_mask * shaded + background + else: + outputs["train_output"] = accum_mask * shaded_albedo + background + else: + outputs["train_output"] = accum_mask * rgb + background + # if self.train_normals: outputs["rendered_orientation_loss"] = orientation_loss( diff --git a/nerfstudio/pipelines/dreamfusion_pipeline.py b/nerfstudio/pipelines/dreamfusion_pipeline.py index 7a535015b..35cc49045 100644 --- a/nerfstudio/pipelines/dreamfusion_pipeline.py +++ b/nerfstudio/pipelines/dreamfusion_pipeline.py @@ -56,7 +56,7 @@ class DreamfusionPipeline(VanillaPipeline): model_outputs = self.model(ray_bundle) # Just uses albedo for now - albedo_output = model_outputs["render"].view(1, 64, 64, 3).permute(0, 3, 1, 2) + albedo_output = model_outputs["train_output"].view(1, 64, 64, 3).permute(0, 3, 1, 2) accumulation = model_outputs["accumulation"].view(64, 64).detach().cpu().numpy() accumulation = np.clip(accumulation, 0.0, 1.0) @@ -71,6 +71,9 @@ class DreamfusionPipeline(VanillaPipeline): plt.imsave("nerf_textureless.jpg", shaded) sds_loss, latents, grad = self.sd.sds_loss(self.text_embedding, albedo_output) + # TODO: opacity penalty using transmittance, not accumultation + accum_mean = np.mean(1.0 - accumulation) + sds_loss *= np.min((0.8, accum_mean)) grad_scaler.scale(latents).backward(gradient=grad, retain_graph=True) # optimizers.scheduler_step_all(step)