diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 46cde2cfc..ae3d78957 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -350,14 +350,23 @@ method_configs["dreamfusion"] = DreamfusionTrainerConfig( train_num_rays_per_batch=4096, eval_num_rays_per_batch=4096, ), - model=DreamFusionModelConfig(eval_num_rays_per_chunk=1 << 15, distortion_loss_mult=10), - alphas_loss_mult=1.5, + model=DreamFusionModelConfig( + eval_num_rays_per_chunk=1 << 15, + distortion_loss_mult=10, + sphere_collider=True, + initialize_density=False, + interlevel_loss_mult=1.0, + proposal_warmup=500, + proposal_update_every=5, + proposal_weights_anneal_max_num_iters=100, + ), + alphas_loss_mult=1.0, opacity_loss_mult=1.0, guidance_scale=100, ), optimizers={ "proposal_networks": { - "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15), + "optimizer": AdamOptimizerConfig(lr=5e-2, eps=1e-15), "scheduler": None, }, "fields": { diff --git a/nerfstudio/data/datamanagers/dreamfusion_datamanager.py b/nerfstudio/data/datamanagers/dreamfusion_datamanager.py index bbc315c67..153913d42 100644 --- a/nerfstudio/data/datamanagers/dreamfusion_datamanager.py +++ b/nerfstudio/data/datamanagers/dreamfusion_datamanager.py @@ -60,23 +60,28 @@ class TrivialDataset(InputDataset): def random_train_pose( - size, - resolution, - device, - radius_mean=1.0, - radius_std=0.1, - central_rotation_range=(0, 360), - vertical_rotation_range=(-90, 10), - focal_range=(0.75, 1.35), - jitter_std=0.01, + size: int, + resolution: int, + device: Union[torch.device, str], + radius_mean: float = 1.0, + radius_std: float = 0.1, + central_rotation_range: Tuple[float, float] = (0, 360), + vertical_rotation_range: Tuple[float, float] = (-80, 20), + focal_range: Tuple[float, float] = (0.75, 1.35), + jitter_std: float = 0.01, + center: Tuple[float, float, float] = (0, 0, 0), ): """generate random poses from an orbit camera Args: size: batch size of generated poses. device: where to allocate the output. - radius: camera radius + radius_mean: mean radius of the orbit camera. + radius_std: standard deviation of the radius of the orbit camera. central_rotation_range: amount that we rotate around the center of the object vertical_rotation_range: amount that we allow the cameras to pan up and down from horizontal + focal_range: focal length range + jitter_std: standard deviation of the jitter added to the camera position + center: center of the object Return: poses: [size, 4, 4] """ @@ -120,7 +125,11 @@ def random_train_pose( origins = torch.stack([torch.tensor([0, 0, 1])] * size, dim=0) origins = (origins * radius_mean) + (origins * (torch.randn((origins.shape)) * radius_std)) R = torch.bmm(rot_z, rot_x) # Want to have Rx @ Ry @ origin - t = torch.bmm(R, origins.unsqueeze(-1)) + torch.randn((size, 3, 1)) * jitter_std + t = ( + torch.bmm(R, origins.unsqueeze(-1)) + + torch.randn((size, 3, 1)) * jitter_std + + torch.tensor(center)[None, :, None] + ) camera_to_worlds = torch.cat([R, t], dim=-1) focals = torch.rand(size) * (focal_range[1] - focal_range[0]) + focal_range[0] @@ -154,16 +163,17 @@ class DreamFusionDataManagerConfig(VanillaDataManagerConfig): """Number of images per batch for training""" eval_images_per_batch: int = 1 """Number of images per batch for evaluation""" - radius_mean: float = 1.6 + radius_mean: float = 2.2 """Mean radius of camera orbit""" - radius_std: float = 0.2 + radius_std: float = 0.1 """Std of radius of camera orbit""" - focal_range: Tuple[float, float] = (0.75, 1.35) + focal_range: Tuple[float, float] = (0.6, 1.2) """Range of focal length""" - vertical_rotation_range: Tuple[float, float] = (-90, 30) + vertical_rotation_range: Tuple[float, float] = (-80, 30) """Range of vertical rotation""" jitter_std: float = 0.05 """Std of camera direction jitter, so we don't just point the cameras towards the center every time""" + center: Tuple[float, float, float] = (0, 0, 0) class DreamFusionDataManager(VanillaDataManager): # pylint: disable=abstract-method @@ -209,6 +219,7 @@ class DreamFusionDataManager(VanillaDataManager): # pylint: disable=abstract-me focal_range=self.config.focal_range, vertical_rotation_range=self.config.vertical_rotation_range, jitter_std=self.config.jitter_std, + center=self.config.center, ) self.train_dataset = TrivialDataset(cameras) @@ -243,6 +254,7 @@ class DreamFusionDataManager(VanillaDataManager): # pylint: disable=abstract-me focal_range=self.config.focal_range, vertical_rotation_range=self.config.vertical_rotation_range, jitter_std=self.config.jitter_std, + center=self.config.center, ) ray_bundle = cameras.generate_rays(torch.tensor(list(range(self.config.train_images_per_batch)))).flatten() @@ -264,6 +276,7 @@ class DreamFusionDataManager(VanillaDataManager): # pylint: disable=abstract-me focal_range=self.config.focal_range, vertical_rotation_range=self.config.vertical_rotation_range, jitter_std=self.config.jitter_std, + center=self.config.center, ) ray_bundle = cameras.generate_rays( torch.tensor([[i] for i in range(self.config.train_images_per_batch)]) diff --git a/nerfstudio/models/dreamfusion.py b/nerfstudio/models/dreamfusion.py index 7d85e0e11..98b07c9be 100644 --- a/nerfstudio/models/dreamfusion.py +++ b/nerfstudio/models/dreamfusion.py @@ -41,10 +41,7 @@ from nerfstudio.model_components.losses import ( orientation_loss, pred_normal_loss, ) -from nerfstudio.model_components.ray_samplers import ( - ProposalNetworkSampler, - UniformSampler, -) +from nerfstudio.model_components.ray_samplers import ProposalNetworkSampler from nerfstudio.model_components.renderers import ( AccumulationRenderer, DepthRenderer, @@ -115,6 +112,10 @@ class DreamFusionModelConfig(ModelConfig): """Proposal loss multiplier.""" distortion_loss_mult: float = 1.0 """Distortion loss multiplier.""" + start_normals_training: int = 1000 + """Start training normals after this many iterations""" + start_lambertian_training: int = 1000 + """start training with lambertian shading after this many iterations""" class DreamFusionModel(Model): @@ -226,11 +227,9 @@ class DreamFusionModel(Model): ) # anneal the weights of the proposal network before doing PDF sampling - N = self.config.proposal_weights_anneal_max_num_iters - def set_anneal(step): # https://arxiv.org/pdf/2111.12077.pdf eq. 18 - train_frac = np.clip(step / N, 0, 1) + train_frac = np.clip(step / self.config.proposal_weights_anneal_max_num_iters, 0, 1) bias = lambda x, b: (b * x) / ((b - 1) * x + 1) anneal = bias(train_frac, self.config.proposal_weights_anneal_slope) self.proposal_sampler.set_anneal(anneal) @@ -244,13 +243,13 @@ class DreamFusionModel(Model): ), TrainingCallback( where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION], - iters=(600,), + iters=(self.config.start_normals_training,), func=start_training_normals, args=[self, training_callback_attributes], ), TrainingCallback( where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION], - iters=(600,), + iters=(self.config.start_lambertian_training,), func=start_shaded_training, args=[self, training_callback_attributes], ), @@ -275,6 +274,7 @@ class DreamFusionModel(Model): def get_param_groups(self) -> Dict[str, List[Parameter]]: param_groups = {} + param_groups["proposal_networks"] = list(self.proposal_networks.parameters()) param_groups["fields"] = list(self.field.parameters()) return param_groups @@ -307,7 +307,7 @@ class DreamFusionModel(Model): background = accum_mask_inv * background_rgb outputs = { - "rgb": rgb, + "rgb_only": rgb, "background_rgb": background_rgb, "background": background, "accumulation": accum_mask, @@ -345,7 +345,7 @@ class DreamFusionModel(Model): outputs["pred_normals"] = self.shader_normals(pred_normals, weights=accum_mask) outputs["shaded"] = accum_mask * shaded + torch.ones_like(shaded) * accum_mask_inv outputs["shaded_albedo"] = shaded_albedo - outputs["render"] = accum_mask * rgb + background + outputs["rgb"] = accum_mask * rgb + background if shading_weight < 0: if np.random.random_sample() > 0.5: @@ -353,7 +353,7 @@ class DreamFusionModel(Model): else: outputs["train_output"] = accum_mask * shaded_albedo + background else: - outputs["train_output"] = outputs["render"] + outputs["train_output"] = outputs["rgb"] if self.training or True: outputs["rendered_orientation_loss"] = orientation_loss( @@ -409,7 +409,21 @@ class DreamFusionModel(Model): outputs["depth"], accumulation=outputs["accumulation"], ) + prop_depth_0 = colormaps.apply_depth_colormap( + outputs["prop_depth_0"], + accumulation=outputs["accumulation"], + ) + prop_depth_1 = colormaps.apply_depth_colormap( + outputs["prop_depth_1"], + accumulation=outputs["accumulation"], + ) metrics_dict = {} - images_dict = {"img": outputs["rgb"], "accumulation": acc, "depth": depth} + images_dict = { + "img": outputs["rgb"], + "accumulation": acc, + "depth": depth, + "prop_depth_0": prop_depth_0, + "prop_depth_1": prop_depth_1, + } return metrics_dict, images_dict diff --git a/nerfstudio/models/nerfacto.py b/nerfstudio/models/nerfacto.py index 168804138..7442bfe8b 100644 --- a/nerfstudio/models/nerfacto.py +++ b/nerfstudio/models/nerfacto.py @@ -237,9 +237,6 @@ class NerfactoModel(Model): return callbacks def get_outputs(self, ray_bundle: RayBundle): - print(ray_bundle.shape) - if not self.training: - raise ray_samples, weights_list, ray_samples_list = self.proposal_sampler(ray_bundle, density_fns=self.density_fns) field_outputs = self.field(ray_samples, compute_normals=self.config.predict_normals) weights = ray_samples.get_weights(field_outputs[FieldHeadNames.DENSITY])