Move camera optimization out of datamanager and parallelize dataloading

nerfstudio/cameras/camera_optimizers.py

@@ -23,19 +23,16 @@
 from typing import Literal, Optional, Type, Union
 
 import torch
-import tyro
 from jaxtyping import Float, Int
 from torch import Tensor, nn
 from typing_extensions import assert_never
 
 from nerfstudio.cameras.lie_groups import exp_map_SE3, exp_map_SO3xR3
+from nerfstudio.cameras.rays import RayBundle
 from nerfstudio.configs.base_config import InstantiateConfig
-from nerfstudio.engine.optimizers import AdamOptimizerConfig, OptimizerConfig
-from nerfstudio.engine.schedulers import (
-    ExponentialDecaySchedulerConfig,
-    SchedulerConfig,
-)
 from nerfstudio.utils import poses as pose_utils
+from nerfstudio.engine.optimizers import OptimizerConfig
+from nerfstudio.engine.schedulers import SchedulerConfig
 
 
 @dataclass
@@ -47,21 +44,38 @@ class CameraOptimizerConfig(InstantiateConfig):
     mode: Literal["off", "SO3xR3", "SE3"] = "off"
     """Pose optimization strategy to use. If enabled, we recommend SO3xR3."""
 
-    position_noise_std: float = 0.0
-    """Noise to add to initial positions. Useful for debugging."""
+    trans_l2_penalty: float = 1e-2
+    """L2 penalty on translation parameters."""
 
-    orientation_noise_std: float = 0.0
-    """Noise to add to initial orientations. Useful for debugging."""
+    rot_l2_penalty: float = 1e-3
+    """L2 penalty on rotation parameters."""
 
-    optimizer: OptimizerConfig = field(default_factory=lambda: AdamOptimizerConfig(lr=6e-4, eps=1e-15))
-    """ADAM parameters for camera optimization."""
+    optimizer: Optional[OptimizerConfig] = field(default=None)
+    """Deprecated, now specified inside the optimizers dict"""
 
-    scheduler: SchedulerConfig = field(default_factory=lambda: ExponentialDecaySchedulerConfig(max_steps=10000))
-    """Learning rate scheduler for camera optimizer.."""
+    scheduler: Optional[SchedulerConfig] = field(default=None)
+    """Deprecated, now specified inside the optimizers dict"""
 
-    param_group: tyro.conf.Suppress[str] = "camera_opt"
-    """Name of the parameter group used for pose optimization. Can be any string that doesn't conflict with other
-    groups."""
+    def __post_init__(self):
+        if self.optimizer is not None:
+            import warnings
+            from nerfstudio.utils.rich_utils import CONSOLE
+
+            CONSOLE.print(
+                "\noptimizer is no longer specified in the CameraOptimizerConfig, it is now defined with the rest of the param groups inside the config file under the name 'camera_opt'\n",
+                style="bold yellow",
+            )
+            warnings.warn("above message coming from", FutureWarning, stacklevel=3)
+
+        if self.scheduler is not None:
+            import warnings
+            from nerfstudio.utils.rich_utils import CONSOLE
+
+            CONSOLE.print(
+                "\nscheduler is no longer specified in the CameraOptimizerConfig, it is now defined with the rest of the param groups inside the config file under the name 'camera_opt'\n",
+                style="bold yellow",
+            )
+            warnings.warn("above message coming from", FutureWarning, stacklevel=3)
 
 
 class CameraOptimizer(nn.Module):
@@ -91,16 +105,6 @@ def __init__(
         else:
             assert_never(self.config.mode)
 
-        # Initialize pose noise; useful for debugging.
-        if config.position_noise_std != 0.0 or config.orientation_noise_std != 0.0:
-            assert config.position_noise_std >= 0.0 and config.orientation_noise_std >= 0.0
-            std_vector = torch.tensor(
-                [config.position_noise_std] * 3 + [config.orientation_noise_std] * 3, device=device
-            )
-            self.pose_noise = exp_map_SE3(torch.normal(torch.zeros((num_cameras, 6), device=device), std_vector))
-        else:
-            self.pose_noise = None
-
     def forward(
         self,
         indices: Int[Tensor, "camera_indices"],
@@ -125,13 +129,46 @@ def forward(
             assert_never(self.config.mode)
         # Detach non-trainable indices by setting to identity transform
         if self.non_trainable_camera_indices is not None:
-            outputs[0][self.non_trainable_camera_indices] = torch.eye(4, device=self.device)[:3, :4]
+            if self.non_trainable_camera_indices.device != self.pose_adjustment.device:
+                self.non_trainable_camera_indices = self.non_trainable_camera_indices.to(self.pose_adjustment.device)
+            outputs[0][self.non_trainable_camera_indices] = torch.eye(4, device=self.pose_adjustment.device)[:3, :4]
 
-        # Apply initial pose noise.
-        if self.pose_noise is not None:
-            outputs.append(self.pose_noise[indices, :, :])
         # Return: identity if no transforms are needed, otherwise multiply transforms together.
         if len(outputs) == 0:
             # Note that using repeat() instead of tile() here would result in unnecessary copies.
             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:
+        """Apply the pose correction to the raybundle"""
+        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()
+
+    def get_loss_dict(self, loss_dict: dict) -> None:
+        """Add regularization"""
+        if self.config.mode != "off":
+            loss_dict["camera_opt_regularizer"] = (
+                self.pose_adjustment[:, :3].norm(dim=-1).mean() * self.config.trans_l2_penalty
+                + self.pose_adjustment[:, 3:].norm(dim=-1).mean() * self.config.rot_l2_penalty
+            )
+
+    def get_correction_matrices(self):
+        """Get optimized pose correction matrices"""
+        return self(torch.arange(0, self.num_cameras).long())
+
+    def get_metrics_dict(self, metrics_dict: dict) -> None:
+        """Get camera optimizer metrics"""
+        if self.config.mode != "off":
+            metrics_dict["camera_opt_translation"] = self.pose_adjustment[:, :3].norm()
+            metrics_dict["camera_opt_rotation"] = self.pose_adjustment[:, 3:].norm()
+
+    def get_param_groups(self, param_groups: dict) -> None:
+        """Get camera optimizer parameters"""
+        camera_opt_params = list(self.parameters())
+        if self.config.mode != "off":
+            assert len(camera_opt_params) > 0
+            param_groups["camera_opt"] = camera_opt_params
+        else:
+            assert len(camera_opt_params) == 0