messing w/ focal optimization

nerfstudio/cameras/camera_optimizers.py

@@ -29,6 +29,7 @@ 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.cameras.cameras import Cameras
 from nerfstudio.configs.base_config import InstantiateConfig
 from nerfstudio.utils import poses as pose_utils
 
@@ -39,7 +40,12 @@ class CameraOptimizerConfig(InstantiateConfig):
 
     _target: Type = field(default_factory=lambda: CameraOptimizer)
 
-    mode: Literal["off", "SO3xR3", "SE3"] = "off"
+    mode: Literal[
+        "off",
+        "SO3xR3",
+        "SE3",
+        "SO3xR3xFocal",
+    ] = "off"
     """Pose optimization strategy to use. If enabled, we recommend SO3xR3."""
 
     trans_l2_penalty: float = 1e-2
@@ -48,6 +54,13 @@ class CameraOptimizerConfig(InstantiateConfig):
     rot_l2_penalty: float = 1e-3
     """L2 penalty on rotation parameters."""
 
+    focal_std_penalty: torch.Tensor = torch.tensor([0.1, 0.1, 0.01, 0.01])
+
+    optimize_focal: bool = False
+    """Whether to optimize focal length."""
+
+    principal_point_penalty: float = 1e-2
+
 
 class CameraOptimizer(nn.Module):
     """Layer that modifies camera poses to be optimized as well as the field during training."""
@@ -60,11 +73,14 @@ class CameraOptimizer(nn.Module):
         num_cameras: int,
         device: Union[torch.device, str],
         non_trainable_camera_indices: Optional[Int[Tensor, "num_non_trainable_cameras"]] = None,
+        cameras: Optional[Cameras] = None,
         **kwargs,
     ) -> None:
         super().__init__()
         self.config = config
         self.num_cameras = num_cameras
+        self.cameras = cameras
+        self.c2ws = self.cameras.camera_to_worlds
         self.device = device
         self.non_trainable_camera_indices = non_trainable_camera_indices
 
@@ -76,6 +92,13 @@ class CameraOptimizer(nn.Module):
         else:
             assert_never(self.config.mode)
 
+        if self.config.optimize_focal:
+            if self.cameras is None:
+                raise ValueError("cameras must be provided to optimize focal length")
+            # optimize fx, fy, cx, cy
+            # self.K_adjustment = torch.nn.Parameter(torch.zeros((num_cameras, 4), device=device))
+            self.K_adjustment = torch.nn.Parameter(torch.zeros((4), device=device))
+
     def forward(
         self,
         indices: Int[Tensor, "camera_indices"],
@@ -108,14 +131,49 @@ class CameraOptimizer(nn.Module):
         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)
+        pose_corrections = functools.reduce(pose_utils.multiply, outputs)
+
+        # next multiply in the focal corrections if provided
+        K_cor = torch.eye(3, device=self.pose_adjustment.device)
+        K_cors = K_cor[None, ...].repeat(indices.shape[0], 1, 1)
+        # optimize them separately
+        # if self.config.optimize_focal:
+        #     K_cors[:, 0, 0] += self.K_adjustment[indices, 0]
+        #     K_cors[:, 1, 1] += self.K_adjustment[indices, 1]
+        #     K_cors[:, 0, 2] += self.K_adjustment[indices, 2]
+        #     K_cors[:, 1, 2] += self.K_adjustment[indices, 3]
+        # optimize all as one
+        if self.config.optimize_focal:
+            K_cor[0, 0] += self.K_adjustment[0]
+            K_cor[1, 1] += self.K_adjustment[1]
+            K_cor[0, 2] += self.K_adjustment[2]
+            K_cor[1, 2] += self.K_adjustment[3]
+        K_cors = K_cor[None, ...].repeat(indices.shape[0], 1, 1)
+        pose_corrections = torch.bmm(K_cors, pose_corrections)  # 1x3x3 * Nx3x4
+        return pose_corrections
 
     def apply_to_raybundle(self, raybundle: RayBundle) -> None:
         """Apply the pose correction to the raybundle"""
         if self.config.mode != "off":
+            # print("focal corr", self.K_adjustment.mean(dim=0), self.K_adjustment.std(dim=0))
+            # print("focal corr", self.K_adjustment)
             correction_matrices = self(raybundle.camera_indices.squeeze())  # type: ignore
+            if self.cameras is not None:
+                # Convert the correction matrix C to H * C * H^-1
+                if self.c2ws.device != correction_matrices.device:
+                    self.c2ws = self.c2ws.to(correction_matrices.device)
+                Hs = self.c2ws[raybundle.camera_indices.squeeze()]
+                # concat a row of 0 0 0 1
+                rows = torch.tensor([0, 0, 0, 1], dtype=Hs.dtype, device=Hs.device)[None, :].repeat(Hs.shape[0], 1, 1)
+                Hs = torch.cat([Hs, rows], dim=1)
+                Hsinv = torch.inverse(Hs)
+                correction_matrices = torch.cat([correction_matrices, rows], dim=1)
+                correction_matrices = torch.bmm(Hs, torch.bmm(correction_matrices, Hsinv))[:, :3, :]
+
             raybundle.origins = raybundle.origins + correction_matrices[:, :3, 3]
             raybundle.directions = torch.bmm(correction_matrices[:, :3, :3], raybundle.directions[..., None]).squeeze()
+            # re-normalize
+            raybundle.directions = raybundle.directions / raybundle.directions.norm(dim=-1, keepdim=True)
 
     def get_loss_dict(self, loss_dict: dict) -> None:
         """Add regularization"""
@@ -124,6 +182,18 @@ class CameraOptimizer(nn.Module):
                 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
             )
+            if self.config.optimize_focal:
+                if self.config.focal_std_penalty.device != self.K_adjustment.device:
+                    self.config.focal_std_penalty = self.config.focal_std_penalty.to(self.K_adjustment.device)
+                # loss_dict["camera_opt_std_regularizer"] = (
+                #     self.K_adjustment.std(dim=0) * self.config.focal_std_penalty
+                # ).sum()
+                # loss_dict["camera_opt_principal_point_regularizer"] = (
+                #     self.K_adjustment.norm(dim=0)[2:].mean() * self.config.principal_point_penalty
+                # )
+                loss_dict["camera_opt_principal_point_regularizer"] = (
+                    self.K_adjustment[2:].norm() * self.config.principal_point_penalty
+                )
 
     def get_correction_matrices(self):
         """Get optimized pose correction matrices"""

nerfstudio/configs/method_configs.py

@@ -108,15 +108,15 @@ method_configs["nerfacto"] = TrainerConfig(
     optimizers={
         "proposal_networks": {
             "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
-            "scheduler": ExponentialDecaySchedulerConfig(lr_final=0.0001, max_steps=200000),
+            "scheduler": ExponentialDecaySchedulerConfig(lr_final=0.0001, max_steps=30000),
         },
         "fields": {
             "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
-            "scheduler": ExponentialDecaySchedulerConfig(lr_final=0.0001, max_steps=200000),
+            "scheduler": ExponentialDecaySchedulerConfig(lr_final=0.0001, max_steps=30000),
         },
         "camera_opt": {
             "optimizer": AdamOptimizerConfig(lr=1e-3, eps=1e-15),
-            "scheduler": ExponentialDecaySchedulerConfig(lr_final=1e-4, max_steps=5000),
+            "scheduler": ExponentialDecaySchedulerConfig(lr_final=1e-4, max_steps=10000),
         },
     },
     viewer=ViewerConfig(num_rays_per_chunk=1 << 15),

nerfstudio/models/nerfacto.py

@@ -29,6 +29,7 @@ from torch.nn import Parameter
 from torchmetrics.functional import structural_similarity_index_measure
 from torchmetrics.image import PeakSignalNoiseRatio
 from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity
+from nerfstudio.cameras.cameras import Cameras
 
 from nerfstudio.cameras.camera_optimizers import CameraOptimizer, CameraOptimizerConfig
 from nerfstudio.cameras.rays import RayBundle, RaySamples
@@ -145,6 +146,8 @@ class NerfactoModel(Model):
     def populate_modules(self):
         """Set the fields and modules."""
         super().populate_modules()
+        if "train_cameras" in self.kwargs:
+            self.train_cameras: Cameras = self.kwargs["train_cameras"]
 
         if self.config.disable_scene_contraction:
             scene_contraction = None
@@ -171,7 +174,7 @@ class NerfactoModel(Model):
         )
 
         self.camera_optimizer: CameraOptimizer = self.config.camera_optimizer.setup(
-            num_cameras=self.num_train_data, device="cpu"
+            num_cameras=self.num_train_data, device="cpu", cameras=self.train_cameras
         )
         self.density_fns = []
         num_prop_nets = self.config.num_proposal_iterations

nerfstudio/pipelines/base_pipeline.py

@@ -273,6 +273,7 @@ class VanillaPipeline(Pipeline):
             metadata=self.datamanager.train_dataset.metadata,
             device=device,
             grad_scaler=grad_scaler,
+            train_cameras=self.datamanager.train_dataset.cameras,
         )
         self.model.to(device)