From 60bf31844fdbfc80589642cc788a5a4b3d2b80a4 Mon Sep 17 00:00:00 2001 From: Alexander Kristoffersen Date: Sun, 2 Oct 2022 17:56:25 -0700 Subject: [PATCH] Camera Extrinsics Optimization (#610) * init commit. * cleanup. * docstring cleanup. --- nerfstudio/cameras/camera_optimizers.py | 150 ++++++++++++++++++ nerfstudio/cameras/cameras.py | 9 +- nerfstudio/data/datamanagers.py | 28 +++- nerfstudio/engine/schedulers.py | 49 +++++- .../field_components/spatial_distortions.py | 5 +- nerfstudio/model_components/ray_generators.py | 12 +- nerfstudio/model_components/ray_samplers.py | 2 - nerfstudio/pipelines/base_pipeline.py | 8 + 8 files changed, 247 insertions(+), 16 deletions(-) create mode 100644 nerfstudio/cameras/camera_optimizers.py diff --git a/nerfstudio/cameras/camera_optimizers.py b/nerfstudio/cameras/camera_optimizers.py new file mode 100644 index 000000000..e9b7badca --- /dev/null +++ b/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 diff --git a/nerfstudio/cameras/cameras.py b/nerfstudio/cameras/cameras.py index 5ee225f52..4aa962c9a 100644 --- a/nerfstudio/cameras/cameras.py +++ b/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 diff --git a/nerfstudio/data/datamanagers.py b/nerfstudio/data/datamanagers.py index 6b4fdb453..b84b536cb 100644 --- a/nerfstudio/data/datamanagers.py +++ b/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 diff --git a/nerfstudio/engine/schedulers.py b/nerfstudio/engine/schedulers.py index 25effd438..6b61ed1dc 100644 --- a/nerfstudio/engine/schedulers.py +++ b/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) diff --git a/nerfstudio/field_components/spatial_distortions.py b/nerfstudio/field_components/spatial_distortions.py index b4b1d7f82..a38348d39 100644 --- a/nerfstudio/field_components/spatial_distortions.py +++ b/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()) diff --git a/nerfstudio/model_components/ray_generators.py b/nerfstudio/model_components/ray_generators.py index e4cdc4efb..3ebb69dd2 100644 --- a/nerfstudio/model_components/ray_generators.py +++ b/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 diff --git a/nerfstudio/model_components/ray_samplers.py b/nerfstudio/model_components/ray_samplers.py index cfd9e88ea..ee474ff7b 100644 --- a/nerfstudio/model_components/ray_samplers.py +++ b/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, diff --git a/nerfstudio/pipelines/base_pipeline.py b/nerfstudio/pipelines/base_pipeline.py index 0992a866d..56a62ba11 100644 --- a/nerfstudio/pipelines/base_pipeline.py +++ b/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