stuff

nerfstudio/cameras/camera_optimizers.py

@@ -53,7 +53,7 @@ class CameraOptimizerConfig(cfg.InstantiateConfig):
     optimizer: AdamOptimizerConfig = AdamOptimizerConfig(lr=6e-4, eps=1e-15)
     """ADAM parameters for camera optimization."""
 
-    scheduler: SchedulerConfig = SchedulerConfig(max_steps=10000)
+    scheduler: SchedulerConfig = SchedulerConfig(max_steps=10000, lr_final=1e-5)
     """Learning rate scheduler for camera optimizer.."""
 
     param_group: tyro.conf.Suppress[str] = "camera_opt"
@@ -77,6 +77,7 @@ class CameraOptimizer(nn.Module):
         self.config = config
         self.num_cameras = num_cameras
         self.device = device
+        self.i = 0
 
         # Initialize learnable parameters.
         if self.config.mode == "off":
@@ -134,10 +135,15 @@ class CameraOptimizer(nn.Module):
         """
         Samples the blur parameters and returns transformation matrices for the origins and directions
         """
+        self.i += 1
+        # blur_adjs = self.blur_adjustment[indices, :] * torch.rand((indices.shape[0], 1), device=self.device)
         blur_adjs = self.blur_adjustment[indices, :] * torch.normal(
             torch.zeros((indices.shape[0], 1), device=self.device),  # mean
             torch.ones((indices.shape[0], 1), device=self.device),  # std
         )
+        # if self.i % 500 == 0:
+        #     # print the directions
+        #     print(self.blur_adjustment)
         # print(
         #     "trans,rot norm",
         #     torch.mean(torch.linalg.norm(blur_adjs[:, :3], dim=0)),

nerfstudio/configs/method_configs.py

@@ -63,10 +63,13 @@ method_configs["nerfacto"] = Config(
     pipeline=VanillaPipelineConfig(
         datamanager=VanillaDataManagerConfig(
             dataparser=NerfstudioDataParserConfig(),
-            train_num_rays_per_batch=4096,
+            train_num_rays_per_batch=1500,
+            num_blur_samples=5,
             eval_num_rays_per_batch=4096,
             camera_optimizer=CameraOptimizerConfig(
-                mode="SO3xR3", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)
+                mode="SO3xR3",
+                optimizer=AdamOptimizerConfig(lr=1e-2, eps=1e-8, weight_decay=0),
+                scheduler=SchedulerConfig(max_steps=10000, lr_final=1e-3),
             ),
         ),
         model=NerfactoModelConfig(eval_num_rays_per_chunk=1 << 15),

nerfstudio/data/datamanagers/base_datamanager.py

@@ -264,6 +264,7 @@ class VanillaDataManagerConfig(InstantiateConfig):
     camera_optimizer: CameraOptimizerConfig = CameraOptimizerConfig()
     """Specifies the camera pose optimizer used during training. Helpful if poses are noisy, such as for data from
     Record3D."""
+    num_blur_samples: int = 1
 
 
 class VanillaDataManager(DataManager):  # pylint: disable=abstract-method
@@ -365,12 +366,19 @@ class VanillaDataManager(DataManager):  # pylint: disable=abstract-method
         self.train_count += 1
         image_batch = next(self.iter_train_image_dataloader)
         batch = self.train_pixel_sampler.sample(image_batch)
-        ray_indices = batch["indices"]
-        ray_bundle = self.train_ray_generator(ray_indices)
+        ray_bundle = self.train_ray_generator(batch["indices"])
+        if self.config.num_blur_samples > 1:
+            # duplicate ray_bundle num_blur_samples times and stack
+            os = ray_bundle.origins.repeat(self.config.num_blur_samples, 1)
+            ds = ray_bundle.directions.repeat(self.config.num_blur_samples, 1)
+            ps = ray_bundle.pixel_area.repeat(self.config.num_blur_samples, 1)
+            cs = ray_bundle.camera_indices.repeat(self.config.num_blur_samples, 1)
+            ray_bundle = RayBundle(origins=os, directions=ds, pixel_area=ps, camera_indices=cs)
 
+        camera_ids = ray_bundle.camera_indices[:, 0]
         # Multiply in the camera optimization deltas to the rays
-        cam_deltas = to4x4(self.train_camera_optimizer(ray_indices[:, 0]))  # B x 3 x 4
-        blur_deltas = self.train_camera_optimizer.sample_blur_correction(ray_indices[:, 0])
+        cam_deltas = to4x4(self.train_camera_optimizer(camera_ids))  # B x 3 x 4
+        blur_deltas = self.train_camera_optimizer.sample_blur_correction(camera_ids)
         # TODO(Justin) is this allowed mathematically
         cam_deltas = torch.bmm(blur_deltas, cam_deltas)
         origins_homogeneous = torch.cat(

nerfstudio/models/nerfacto.py

@@ -216,6 +216,11 @@ class NerfactoModel(Model):
         ray_samples_list.append(ray_samples)
 
         rgb = self.renderer_rgb(rgb=field_outputs[FieldHeadNames.RGB], weights=weights)
+        s = 1500
+        if rgb.shape[0] == s * 5:
+            rgb = (1 / 5.0) * (
+                rgb[:s, :] + rgb[s : 2 * s, :] + rgb[2 * s : 3 * s, :] + rgb[3 * s : 4 * s, :] + rgb[4 * s : 5 * s, :]
+            )
         depth = self.renderer_depth(weights=weights, ray_samples=ray_samples)
         accumulation = self.renderer_accumulation(weights=weights)
 

nerfstudio/utils/tensor_dataclass.py

@@ -197,7 +197,7 @@ class TensorDataclass:
         """
         return self._apply_fn_to_fields(lambda x: x.broadcast_to((*shape, x.shape[-1])))
 
-    def to(self: TensorDataclassT, device) -> TensorDataclassT:
+    def to(self: TensorDataclassT, *args, **kwargs) -> TensorDataclassT:
         """Returns a new TensorDataclass with the same data but on the specified device.
 
         Args:
@@ -206,7 +206,18 @@ class TensorDataclass:
         Returns:
             A new TensorDataclass with the same data but on the specified device.
         """
-        return self._apply_fn_to_fields(lambda x: x.to(device))
+        return self._apply_fn_to_fields(lambda x: x.to(*args, **kwargs))
+
+    def pin_memory(self: TensorDataclassT) -> TensorDataclassT:
+        """Returns a new TensorDataclass with the same data but on the specified device.
+
+        Args:
+            device: The device to place the tensor dataclass.
+
+        Returns:
+            A new TensorDataclass with the same data but on the specified device.
+        """
+        return self._apply_fn_to_fields(lambda x: x.pin_memory())
 
     def _apply_fn_to_fields(
         self: TensorDataclassT, fn: Callable, dataclass_fn: Optional[Callable] = None