merging

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": {

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)])

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

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])