diff --git a/nerfstudio/cameras/camera_optimizers.py b/nerfstudio/cameras/camera_optimizers.py index ca3f90742..cc914bd2b 100644 --- a/nerfstudio/cameras/camera_optimizers.py +++ b/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 diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 6c6d787d7..0c6095d44 100644 --- a/nerfstudio/configs/method_configs.py +++ b/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), diff --git a/nerfstudio/models/nerfacto.py b/nerfstudio/models/nerfacto.py index 862abd170..ca40ed3aa 100644 --- a/nerfstudio/models/nerfacto.py +++ b/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)