deblur stuff

nerfstudio/cameras/camera_optimizers.py

@@ -20,7 +20,7 @@ from __future__ import annotations
 
 import functools
 from dataclasses import dataclass, field
-from typing import Literal, Optional, Type, Union
+from typing import Literal, Optional, Type, Union, Dict, Tuple
 
 import torch
 from jaxtyping import Float, Int
@@ -42,7 +42,7 @@ class CameraOptimizerConfig(InstantiateConfig):
     mode: Literal["off", "SO3xR3", "SE3"] = "off"
     """Pose optimization strategy to use. If enabled, we recommend SO3xR3."""
 
-    trans_l2_penalty: float = 1e-2
+    trans_l2_penalty: float = 1e-3
     """L2 penalty on translation parameters."""
 
     rot_l2_penalty: float = 1e-3
@@ -110,12 +110,14 @@ class CameraOptimizer(nn.Module):
             return torch.eye(4, device=self.device)[None, :3, :4].tile(indices.shape[0], 1, 1)
         return functools.reduce(pose_utils.multiply, outputs)
 
-    def apply_to_raybundle(self, raybundle: RayBundle) -> None:
+    def apply_to_raybundle(self, raybundle: RayBundle) -> RayBundle:
         """Apply the pose correction to the raybundle"""
+        assert raybundle.camera_indices is not None
         if self.config.mode != "off":
             correction_matrices = self(raybundle.camera_indices.squeeze())  # type: ignore
             raybundle.origins = raybundle.origins + correction_matrices[:, :3, 3]
             raybundle.directions = torch.bmm(correction_matrices[:, :3, :3], raybundle.directions[..., None]).squeeze()
+        return raybundle
 
     def get_loss_dict(self, loss_dict: dict) -> None:
         """Add regularization"""
@@ -143,3 +145,76 @@ class CameraOptimizer(nn.Module):
             param_groups["camera_opt"] = camera_opt_params
         else:
             assert len(camera_opt_params) == 0
+
+
+@dataclass
+class DeblurCameraOptimizerConfig(CameraOptimizerConfig):
+    """Configuration of optimization for camera poses."""
+
+    _target: Type = field(default_factory=lambda: DeblurCameraOptimizer)
+    """The target class to be instantiated."""
+
+    num_samples: int = 7
+    """The number of samples to use for deblurring."""
+
+    blur_l2_penalty: float = 1e-4
+    """The L2 penalty for the blur parameters."""
+
+
+class DeblurCameraOptimizer(CameraOptimizer):
+    config: DeblurCameraOptimizerConfig
+
+    def __init__(self, config: DeblurCameraOptimizerConfig, *args, **kwargs) -> None:
+        super().__init__(config, *args, **kwargs)
+        self.config = config
+        self.blur_adjustment = torch.nn.Parameter(torch.zeros((self.num_cameras, 6), device=self.device))
+
+    def sample_blur_correction(self, indices: Int[Tensor, "num_cameras"]):
+        """
+        Samples the blur parameters and returns transformation matrices for the origins and directions
+        """
+        times = (
+            torch.linspace(0, 1, self.config.num_samples, device=indices.device)
+            .unsqueeze(1)
+            .repeat(indices.shape[0] // self.config.num_samples, 1)
+        )
+        # add a small random jittter to each sample time to avoid aliasing
+        # w = 1 / (self.config.num_samples - 1)
+        # times += torch.rand_like(times) * w - w / 2
+        blur_adjs = self.blur_adjustment[indices, :] * times
+        # normal seems to work worse
+        # blur_adjs = self.blur_adjustment[indices, :] * torch.normal(
+        #     torch.zeros((indices.shape[0], 1), device=indices.device),  # mean
+        #     torch.ones((indices.shape[0], 1), device=indices.device),  # std
+        # )
+        return exp_map_SO3xR3(blur_adjs)
+
+    def apply_to_raybundle(self, ray_bundle: RayBundle) -> RayBundle:
+        assert ray_bundle.camera_indices is not None
+        if self.blur_adjustment.device != ray_bundle.origins.device:
+            self.blur_adjustment = self.blur_adjustment.to(ray_bundle.origins.device)
+
+        optimized_bundle = super().apply_to_raybundle(ray_bundle)
+        if self.config.num_samples == 1:
+            return optimized_bundle
+
+        # duplicate optimized_bundle num_samples times and stack
+        def repeatfn(x):
+            return x.repeat_interleave(self.config.num_samples, 0)
+
+        ray_bundle = optimized_bundle._apply_fn_to_fields(repeatfn)
+        camera_ids = ray_bundle.camera_indices.squeeze()
+        # Multiply in the camera optimization deltas to the rays
+        blur_deltas = self.sample_blur_correction(camera_ids)
+        ray_bundle.origins = ray_bundle.origins + blur_deltas[:, :, 3]
+        ray_bundle.directions = torch.bmm(
+            blur_deltas[:, :, :3], ray_bundle.directions.view(*ray_bundle.directions.shape, 1)
+        )[..., 0]
+        return ray_bundle
+
+    def get_loss_dict(self, loss_dict: dict) -> None:
+        """Add regularization"""
+        super().get_loss_dict(loss_dict)
+        if self.config.mode != "off":
+            # add a loss penalizing norm of blur adjustment
+            loss_dict["blur_regularizer"] = self.blur_adjustment.norm(dim=-1).mean() * self.config.blur_l2_penalty

nerfstudio/configs/method_configs.py

@@ -102,7 +102,6 @@ method_configs["nerfacto"] = TrainerConfig(
         ),
         model=NerfactoModelConfig(
             eval_num_rays_per_chunk=1 << 15,
-            camera_optimizer=CameraOptimizerConfig(mode="SO3xR3"),
         ),
     ),
     optimizers={
@@ -115,8 +114,8 @@ method_configs["nerfacto"] = TrainerConfig(
             "scheduler": ExponentialDecaySchedulerConfig(lr_final=0.0001, max_steps=200000),
         },
         "camera_opt": {
-            "optimizer": AdamOptimizerConfig(lr=1e-3, eps=1e-15),
-            "scheduler": ExponentialDecaySchedulerConfig(lr_final=1e-4, max_steps=5000),
+            "optimizer": AdamOptimizerConfig(lr=1e-4, eps=1e-15),
+            "scheduler": ExponentialDecaySchedulerConfig(lr_final=1e-5, max_steps=30000),
         },
     },
     viewer=ViewerConfig(num_rays_per_chunk=1 << 15),
@@ -145,7 +144,6 @@ method_configs["nerfacto-big"] = TrainerConfig(
             max_res=4096,
             proposal_weights_anneal_max_num_iters=5000,
             log2_hashmap_size=21,
-            camera_optimizer=CameraOptimizerConfig(mode="SO3xR3"),
         ),
     ),
     optimizers={
@@ -183,7 +181,7 @@ method_configs["nerfacto-huge"] = TrainerConfig(
             num_nerf_samples_per_ray=64,
             num_proposal_samples_per_ray=(512, 512),
             proposal_net_args_list=[
-                {"hidden_dim": 16, "log2_hashmap_size": 17, "num_levels": 5, "max_res": 512, "use_linear": False},
+                {"hidden_dim": 16, "log2_hashmap_size": 17, "num_levels": 6, "max_res": 512, "use_linear": False},
                 {"hidden_dim": 16, "log2_hashmap_size": 17, "num_levels": 7, "max_res": 2048, "use_linear": False},
             ],
             hidden_dim=256,
@@ -192,7 +190,6 @@ method_configs["nerfacto-huge"] = TrainerConfig(
             max_res=8192,
             proposal_weights_anneal_max_num_iters=5000,
             log2_hashmap_size=21,
-            camera_optimizer=CameraOptimizerConfig(mode="SO3xR3"),
         ),
     ),
     optimizers={
@@ -205,8 +202,8 @@ method_configs["nerfacto-huge"] = TrainerConfig(
             "scheduler": ExponentialDecaySchedulerConfig(lr_final=1e-4, max_steps=50000),
         },
         "camera_opt": {
-            "optimizer": AdamOptimizerConfig(lr=1e-3, eps=1e-15),
-            "scheduler": ExponentialDecaySchedulerConfig(lr_final=1e-4, max_steps=5000),
+            "optimizer": AdamOptimizerConfig(lr=1e-4, eps=1e-15),
+            "scheduler": ExponentialDecaySchedulerConfig(lr_final=1e-5, max_steps=30000),
         },
     },
     viewer=ViewerConfig(num_rays_per_chunk=1 << 15),

nerfstudio/models/nerfacto.py

@@ -30,7 +30,12 @@ from torchmetrics.functional import structural_similarity_index_measure
 from torchmetrics.image import PeakSignalNoiseRatio
 from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity
 
-from nerfstudio.cameras.camera_optimizers import CameraOptimizer, CameraOptimizerConfig
+from nerfstudio.cameras.camera_optimizers import (
+    CameraOptimizer,
+    CameraOptimizerConfig,
+    DeblurCameraOptimizerConfig,
+    DeblurCameraOptimizer,
+)
 from nerfstudio.cameras.rays import RayBundle, RaySamples
 from nerfstudio.engine.callbacks import TrainingCallback, TrainingCallbackAttributes, TrainingCallbackLocation
 from nerfstudio.field_components.field_heads import FieldHeadNames
@@ -129,7 +134,7 @@ class NerfactoModelConfig(ModelConfig):
     """Which implementation to use for the model."""
     appearance_embed_dim: int = 32
     """Dimension of the appearance embedding."""
-    camera_optimizer: CameraOptimizerConfig = CameraOptimizerConfig(mode="SO3xR3")
+    camera_optimizer: CameraOptimizerConfig = DeblurCameraOptimizerConfig(mode="SO3xR3")
     """Config of the camera optimizer to use"""
 
 
@@ -287,7 +292,7 @@ class NerfactoModel(Model):
     def get_outputs(self, ray_bundle: RayBundle):
         # apply the camera optimizer pose tweaks
         if self.training:
-            self.camera_optimizer.apply_to_raybundle(ray_bundle)
+            ray_bundle = self.camera_optimizer.apply_to_raybundle(ray_bundle)
         ray_samples: RaySamples
         ray_samples, weights_list, ray_samples_list = self.proposal_sampler(ray_bundle, density_fns=self.density_fns)
         field_outputs = self.field.forward(ray_samples, compute_normals=self.config.predict_normals)
@@ -299,6 +304,10 @@ class NerfactoModel(Model):
         ray_samples_list.append(ray_samples)
 
         rgb = self.renderer_rgb(rgb=field_outputs[FieldHeadNames.RGB], weights=weights)
+        if self.training and isinstance(self.camera_optimizer, DeblurCameraOptimizer):
+            n = self.camera_optimizer.config.num_samples
+            s = ray_bundle.origins.shape[0] // n
+            rgb = rgb.view(s, n, 3).mean(dim=1)
         depth = self.renderer_depth(weights=weights, ray_samples=ray_samples)
         accumulation = self.renderer_accumulation(weights=weights)