From e0405439eb41ef638a5cf6de97d3f93288fea71f Mon Sep 17 00:00:00 2001 From: Justin Kerr Date: Sun, 13 Aug 2023 21:35:50 -0700 Subject: [PATCH] messing w/ focal optimization --- nerfstudio/cameras/camera_optimizers.py | 74 ++++++++++++++++++++++++- nerfstudio/configs/method_configs.py | 6 +- nerfstudio/models/nerfacto.py | 5 +- nerfstudio/pipelines/base_pipeline.py | 1 + 4 files changed, 80 insertions(+), 6 deletions(-) diff --git a/nerfstudio/cameras/camera_optimizers.py b/nerfstudio/cameras/camera_optimizers.py index 40f0a82f9..faaf5a54d 100644 --- a/nerfstudio/cameras/camera_optimizers.py +++ b/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""" diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index d5afb0654..379a060a0 100644 --- a/nerfstudio/configs/method_configs.py +++ b/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), diff --git a/nerfstudio/models/nerfacto.py b/nerfstudio/models/nerfacto.py index 862abd170..5e75c9250 100644 --- a/nerfstudio/models/nerfacto.py +++ b/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 diff --git a/nerfstudio/pipelines/base_pipeline.py b/nerfstudio/pipelines/base_pipeline.py index 030cfcb32..61e8eb79c 100644 --- a/nerfstudio/pipelines/base_pipeline.py +++ b/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)