Camera Extrinsics Optimization (#610)

* init commit.

* cleanup.

* docstring cleanup.

nerfstudio/cameras/camera_optimizers.py

@@ -0,0 +1,150 @@
+# Copyright 2022 The Plenoptix Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Pose and Intrinsics Optimizers
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Optional, Type, Union
+
+import torch
+from torch import nn
+from torchtyping import TensorType
+
+from nerfstudio.configs import base_config as cfg
+from nerfstudio.utils import poses as pose_utils
+
+
+# Camera Optimization related configs
+@dataclass
+class CameraOptimizerConfig(cfg.InstantiateConfig):
+    """Default camera optimizer config. Note: This is a no-op class and will not optimize cameras."""
+
+    _target: Type = field(default_factory=lambda: CameraOptimizer)
+
+
+@dataclass
+class BARFPoseOptimizerConfig(CameraOptimizerConfig):
+    """BARF camera optimizer."""
+
+    _target: Type = field(default_factory=lambda: BARFOptimizer)
+    noise_variance: Optional[float] = None
+    """Optional additional noise to add to poses. Useful for debugging pose optimization."""
+
+
+class CameraOptimizer(nn.Module):
+    """Layer that modifies camera poses to be optimized as well as the field during training."""
+
+    config: CameraOptimizerConfig
+
+    def __init__(
+        self,
+        config: CameraOptimizerConfig,
+        num_cameras: int,
+        device: Union[torch.device, str],
+        **kwargs,  # pylint: disable=unused-argument
+    ) -> None:
+        super().__init__()
+        self.config = config
+        self.num_cameras = num_cameras
+        self.device = device
+
+    def forward(
+        self,
+        indices: TensorType["num_cameras"],
+    ) -> TensorType["num_cameras", 3, 4]:
+        """Indexing into camera adjustments.
+        Args:
+            indices: indices of Cameras to optimize.
+        Returns:
+            Tranformation matrices from optimized camera coordinates coordinates
+            to given camera coordinates.
+        """
+        return torch.eye(4, device=self.device).repeat(indices.shape[0], 1, 1)[:, :3, :4]  # no-op (Identity Transform)
+
+
+class BARFOptimizer(CameraOptimizer):
+    """Bundle-Adjusting NeRF (BARF) Pose Optimization"""
+
+    config: BARFPoseOptimizerConfig
+
+    def __init__(self, config: BARFPoseOptimizerConfig, num_cameras: int, device: Union[torch.device, str]) -> None:
+        super().__init__(config, num_cameras, device)
+        if self.config.noise_variance is not None:
+            pose_noise = torch.normal(torch.zeros(self.num_cameras, 6), self.config.noise_variance)
+            self.pose_noise = self.exp_map(pose_noise).detach().to(device)
+        self.pose_adjustment = nn.Embedding(self.num_cameras, 6, device=device)
+        nn.init.zeros_(self.pose_adjustment.weight)
+
+    @classmethod
+    def exp_map(cls, tangent_vector: TensorType["num_cameras", 6]) -> TensorType["num_cameras", 3, 4]:
+        """Convert SE3 vector into [R|t] transformation matrix.
+        Args:
+            tangent_vector: SE3 vector
+        Returns:
+            Respective [R|t] tranformation matrices.
+        """
+
+        tangent_vector_lin = tangent_vector[:, :3].view(-1, 3, 1)
+        tangent_vector_ang = tangent_vector[:, 3:].view(-1, 3, 1)
+
+        theta = torch.linalg.norm(tangent_vector_ang, dim=1).unsqueeze(1)
+        theta2 = theta**2
+        theta3 = theta**3
+
+        near_zero = theta < 1e-2
+        non_zero = torch.ones(1, dtype=tangent_vector.dtype, device=tangent_vector.device)
+        theta_nz = torch.where(near_zero, non_zero, theta)
+        theta2_nz = torch.where(near_zero, non_zero, theta2)
+        theta3_nz = torch.where(near_zero, non_zero, theta3)
+
+        # Compute the rotation
+        sine = theta.sin()
+        cosine = torch.where(near_zero, 8 / (4 + theta2) - 1, theta.cos())
+        sine_by_theta = torch.where(near_zero, 0.5 * cosine + 0.5, sine / theta_nz)
+        one_minus_cosine_by_theta2 = torch.where(near_zero, 0.5 * sine_by_theta, (1 - cosine) / theta2_nz)
+        ret = torch.zeros(tangent_vector.shape[0], 3, 4).to(dtype=tangent_vector.dtype, device=tangent_vector.device)
+        ret[:, :3, :3] = one_minus_cosine_by_theta2 * tangent_vector_ang @ tangent_vector_ang.transpose(1, 2)
+
+        ret[:, 0, 0] += cosine.view(-1)
+        ret[:, 1, 1] += cosine.view(-1)
+        ret[:, 2, 2] += cosine.view(-1)
+        temp = sine_by_theta.view(-1, 1) * tangent_vector_ang.view(-1, 3)
+        ret[:, 0, 1] -= temp[:, 2]
+        ret[:, 1, 0] += temp[:, 2]
+        ret[:, 0, 2] += temp[:, 1]
+        ret[:, 2, 0] -= temp[:, 1]
+        ret[:, 1, 2] -= temp[:, 0]
+        ret[:, 2, 1] += temp[:, 0]
+
+        # Compute the translation
+        sine_by_theta = torch.where(near_zero, 1 - theta2 / 6, sine_by_theta)
+        one_minus_cosine_by_theta2 = torch.where(near_zero, 0.5 - theta2 / 24, one_minus_cosine_by_theta2)
+        theta_minus_sine_by_theta3_t = torch.where(near_zero, 1.0 / 6 - theta2 / 120, (theta - sine) / theta3_nz)
+
+        ret[:, :, 3:] = sine_by_theta * tangent_vector_lin
+        ret[:, :, 3:] += one_minus_cosine_by_theta2 * torch.cross(tangent_vector_ang, tangent_vector_lin, dim=1)
+        ret[:, :, 3:] += theta_minus_sine_by_theta3_t * (
+            tangent_vector_ang @ (tangent_vector_ang.transpose(1, 2) @ tangent_vector_lin)
+        )
+        return ret
+
+    def forward(self, indices: TensorType["num_cameras"]) -> TensorType["num_cameras", 3, 4]:
+        c_opt2c = self.exp_map(self.pose_adjustment.weight[indices])
+        if self.config.noise_variance is not None:
+            c_opt2c = pose_utils.multiply(c_opt2c, self.pose_noise[indices])
+        return c_opt2c

nerfstudio/cameras/cameras.py

@@ -26,6 +26,7 @@ import torchvision
 from torch.nn.functional import normalize
 from torchtyping import TensorType
 
+import nerfstudio.utils.poses as pose_utils
 from nerfstudio.cameras import camera_utils
 from nerfstudio.cameras.rays import RayBundle
 
@@ -247,7 +248,7 @@ class Cameras:
         self,
         camera_indices: Union[TensorType["num_rays":...], int],
         coords: Optional[TensorType["num_rays":..., 2]] = None,
-        camera_to_world_delta: Optional[TensorType["num_rays":..., 3, 4]] = None,
+        camera_opt_to_camera: Optional[TensorType["num_rays":..., 3, 4]] = None,
         distortion_params_delta: Optional[TensorType["num_rays":..., 6]] = None,
     ) -> RayBundle:
         """Generates rays for the given camera indices.
@@ -258,7 +259,7 @@ class Cameras:
         Args:
             camera_indices: Indices of the cameras to generate rays for.
             coords: Coordinates of the pixels to generate rays for. If None, the full image will be rendered.
-            camera_to_world_delta: Optional delta for the camera to world matrices.
+            camera_opt_to_camera: Optional transform for the camera to world matrices.
             distortion_params_delta: Optional delta for the distortion parameters.
 
         Returns:
@@ -315,8 +316,8 @@ class Cameras:
             raise ValueError(f"Camera type {CameraType(self.camera_type[0])} not supported.")
 
         c2w = self.camera_to_worlds[camera_indices]
-        if camera_to_world_delta is not None:
-            c2w = c2w + camera_to_world_delta
+        if camera_opt_to_camera is not None:
+            c2w = pose_utils.multiply(c2w, camera_opt_to_camera)
         rotation = c2w[..., :3, :3]  # (..., 3, 3)
         directions_stack = torch.sum(
             directions_stack[..., None, :] * rotation, dim=-1

nerfstudio/data/datamanagers.py

@@ -28,6 +28,7 @@ from torch.nn import Parameter
 from torch.utils.data import Dataset
 from torch.utils.data.distributed import DistributedSampler
 
+from nerfstudio.cameras.camera_optimizers import CameraOptimizerConfig
 from nerfstudio.cameras.rays import RayBundle
 from nerfstudio.configs.base_config import InstantiateConfig
 from nerfstudio.data.dataparsers.blender_dataparser import BlenderDataParserConfig
@@ -247,6 +248,8 @@ class VanillaDataManagerConfig(InstantiateConfig):
     """number of images to sample during eval iteration"""
     eval_image_indices: Optional[Tuple[int, ...]] = (0,)
     """specifies the image indices to use during eval; if None, uses all"""
+    train_camera_optimizer: CameraOptimizerConfig = CameraOptimizerConfig()
+    """specifies the camera pose optimizer used during training"""
 
 
 class VanillaDataManager(DataManager):  # pylint: disable=abstract-method
@@ -306,7 +309,13 @@ class VanillaDataManager(DataManager):  # pylint: disable=abstract-method
         )
         self.iter_train_image_dataloader = iter(self.train_image_dataloader)
         self.train_pixel_sampler = PixelSampler(self.config.train_num_rays_per_batch)
-        self.train_ray_generator = RayGenerator(self.train_dataset.dataset_inputs.cameras.to(self.device))
+        self.train_camera_optimizer = self.config.train_camera_optimizer.setup(
+            num_cameras=self.train_dataset.dataset_inputs.cameras.size, device=self.device
+        )
+        self.train_ray_generator = RayGenerator(
+            self.train_dataset.dataset_inputs.cameras.to(self.device),
+            self.train_camera_optimizer,
+        )
 
     def setup_eval(self):
         """Sets up the data loader for evaluation"""
@@ -327,7 +336,10 @@ class VanillaDataManager(DataManager):  # pylint: disable=abstract-method
         )
         self.iter_eval_image_dataloader = iter(self.eval_image_dataloader)
         self.eval_pixel_sampler = PixelSampler(self.config.eval_num_rays_per_batch)
-        self.eval_ray_generator = RayGenerator(self.eval_dataset.dataset_inputs.cameras.to(self.device))
+        self.eval_ray_generator = RayGenerator(
+            self.eval_dataset.dataset_inputs.cameras.to(self.device),
+            self.train_camera_optimizer,  # should be shared between train and eval.
+        )
         # for loading full images
         self.fixed_indices_eval_dataloader = FixedIndicesEvalDataloader(
             input_dataset=self.eval_dataset,
@@ -368,3 +380,15 @@ class VanillaDataManager(DataManager):  # pylint: disable=abstract-method
             image_idx = int(camera_ray_bundle.camera_indices[0, 0, 0])
             return image_idx, camera_ray_bundle, batch
         raise ValueError("No more eval images")
+
+    def get_param_groups(self) -> Dict[str, List[Parameter]]:  # pylint: disable=no-self-use
+        """Get the param groups for the data manager.
+        Returns:
+            A list of dictionaries containing the data manager's param groups.
+        """
+        param_groups = {}
+
+        camera_opt_params = list(self.train_camera_optimizer.parameters())
+        if len(camera_opt_params) > 0:
+            param_groups["camera_opt"] = list(camera_opt_params)
+        return param_groups

nerfstudio/engine/schedulers.py

@@ -15,10 +15,10 @@
 """Scheduler Classes"""
 
 from dataclasses import dataclass, field
-from typing import Any, Type
+from typing import Any, Optional, Type
 
 import numpy as np
-from torch.optim.lr_scheduler import LambdaLR
+from torch.optim import Optimizer, lr_scheduler
 
 from nerfstudio.configs.base_config import InstantiateConfig
 
@@ -38,7 +38,7 @@ class SchedulerConfig(InstantiateConfig):
         return self._target(optimizer, lr_init, self.lr_final, self.max_steps)
 
 
-class ExponentialDecaySchedule(LambdaLR):
+class ExponentialDecaySchedule(lr_scheduler.LambdaLR):
     """Exponential learning rate decay function.
     See https://github.com/google-research/google-research/blob/
     fd2cea8cdd86b3ed2c640cbe5561707639e682f3/jaxnerf/nerf/utils.py#L360
@@ -71,3 +71,46 @@ class ExponentialDecaySchedule(LambdaLR):
             return delay_rate * multiplier
 
         super().__init__(optimizer, lr_lambda=func)
+
+
+class DelayerScheduler(lr_scheduler.LambdaLR):
+    """Starts with a flat lr schedule until it reaches N epochs then applies a given scheduler"""
+
+    def __init__(
+        self,
+        optimizer: Optimizer,
+        lr_init,  # pylint: disable=unused-argument
+        lr_final,  # pylint: disable=unused-argument
+        max_steps,  # pylint: disable=unused-argument
+        delay_epochs: int = 500,
+        after_scheduler: Optional[lr_scheduler.LambdaLR] = None,
+    ) -> None:
+        def func(step):
+            if step > delay_epochs:
+                if after_scheduler is not None:
+                    multiplier = after_scheduler.lr_lambdas[0](step - delay_epochs)  # type: ignore
+                    return multiplier
+                return 1.0
+            return 0.0
+
+        super().__init__(optimizer, lr_lambda=func)
+
+
+class DelayedExponentialScheduler(DelayerScheduler):
+    """Delayer Scheduler with an Exponential Scheduler initialized afterwards."""
+
+    def __init__(
+        self,
+        optimizer: Optimizer,
+        lr_init,
+        lr_final,
+        max_steps,
+        delay_epochs: int = 200,
+    ):
+        after_scheduler = ExponentialDecaySchedule(
+            optimizer,
+            lr_init,
+            lr_final,
+            max_steps,
+        )
+        super().__init__(optimizer, lr_init, lr_final, max_steps, delay_epochs, after_scheduler=after_scheduler)

nerfstudio/field_components/spatial_distortions.py

@@ -67,9 +67,10 @@ class SceneContraction(SpatialDistortion):
         def contract(x):
             mag = torch.linalg.norm(x, ord=self.order, dim=-1)
             mask = mag >= 1
-            x[mask] = (2 - (1 / mag[mask][..., None])) * (x[mask] / mag[mask][..., None])
+            x_new = x.clone()
+            x_new[mask] = (2 - (1 / mag[mask][..., None])) * (x[mask] / mag[mask][..., None])
 
-            return x
+            return x_new
 
         if isinstance(positions, Gaussians):
             means = contract(positions.mean.clone())

nerfstudio/model_components/ray_generators.py

@@ -18,21 +18,24 @@ Ray generator.
 from torch import nn
 from torchtyping import TensorType
 
+from nerfstudio.cameras.camera_optimizers import CameraOptimizer
 from nerfstudio.cameras.cameras import Cameras
 from nerfstudio.cameras.rays import RayBundle
 
 
 class RayGenerator(nn.Module):
     """torch.nn Module for generating rays.
-    This class will store the intrinsics and extrinsics parameters of the cameras.
+    This class is the interface between the scene's cameras/camera optimizer and the ray sampler.
 
     Args:
-        cameras: Camera objects containing camera info
+        cameras: Camera objects containing camera info.
+        pose_optimizer: pose optimization module, for optimizing noisy camera intrisics/extrinsics.
     """
 
-    def __init__(self, cameras: Cameras) -> None:
+    def __init__(self, cameras: Cameras, pose_optimizer: CameraOptimizer) -> None:
         super().__init__()
         self.cameras = cameras
+        self.pose_optimizer = pose_optimizer
         self.image_coords = nn.Parameter(cameras.get_image_coords(), requires_grad=False)
 
     def forward(self, ray_indices: TensorType["num_rays", 3]) -> RayBundle:
@@ -46,8 +49,11 @@ class RayGenerator(nn.Module):
         x = ray_indices[:, 2]  # col indices
         coords = self.image_coords[y, x]
 
+        camera_opt_to_camera = self.pose_optimizer(c)
+
         ray_bundle = self.cameras.generate_rays(
             camera_indices=c,
             coords=coords,
+            camera_opt_to_camera=camera_opt_to_camera,
         )
         return ray_bundle

nerfstudio/model_components/ray_samplers.py

@@ -75,7 +75,6 @@ class SpacedSampler(Sampler):
         self.spacing_fn = spacing_fn
         self.spacing_fn_inv = spacing_fn_inv
 
-    @torch.no_grad()
     def generate_ray_samples(
         self,
         ray_bundle: Optional[RayBundle] = None,
@@ -265,7 +264,6 @@ class PDFSampler(Sampler):
         self.include_original = include_original
         self.histogram_padding = histogram_padding
 
-    @torch.no_grad()
     def generate_ray_samples(
         self,
         ray_bundle: Optional[RayBundle] = None,

nerfstudio/pipelines/base_pipeline.py

@@ -250,6 +250,14 @@ class VanillaPipeline(Pipeline):
             valid_mask = model_outputs["valid_mask"]
             batch = get_masked_dict(batch, valid_mask)
         metrics_dict = self.model.get_metrics_dict(model_outputs, batch)
+        if "camera_opt" in self.datamanager.get_param_groups():
+            # Report the camera optimization metrics
+            metrics_dict["camera_opt_translation"] = (
+                self.datamanager.get_param_groups()["camera_opt"][0].data[:, :3].norm()
+            )
+            metrics_dict["camera_opt_rotation"] = (
+                self.datamanager.get_param_groups()["camera_opt"][0].data[:, 3:].norm()
+            )
         loss_dict = self.model.get_loss_dict(model_outputs, batch, metrics_dict)
 
         return model_outputs, loss_dict, metrics_dict