tink/nerfstudio
4
0
Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 18:44:21 +08:00
Code

Camera pose optimization for Splatfacto (#2891)

Browse Source 
* Add pose optimization to Splatfacto

* Disable Splatfacto pose optimization by default

* Improve apply_to_camera for Gaussian Splatting pose optimization

Do not chain modifications to camera_to_worlds to improve
numerical stability and enable L2 rot/trans penalties.

* Add separate mean and max rot/trans metrics to camera-opt

* Tweak pose optimization hyperparameters

Parameters used in the Gaussian Splatting on the Move paper v1

* Unit test fix for new cameara_optimizer training metrics

* Adjust splatfacto-big camera pose optimization parameters

Same parameters as in normal Splatfacto

---------

Co-authored-by: jh-surh <jh.surh@bucketplace.net>
This commit is contained in:
Otto Seiskari 2024-04-11 21:23:17 +03:00 committed by GitHub
parent 2d9bbe5091
commit eba72db22a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 54 additions and 19 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

32
nerfstudio/cameras/camera_optimizers.py
View File

@ -22,6 +22,7 @@ import functools
from dataclasses import dataclass, field
from typing import Literal, Optional, Type, Union
import numpy
import torch
import tyro
from jaxtyping import Float, Int
@ -151,15 +152,20 @@ class CameraOptimizer(nn.Module):
raybundle.origins = raybundle.origins + correction_matrices[:, :3, 3]
raybundle.directions = torch.bmm(correction_matrices[:, :3, :3], raybundle.directions[..., None]).squeeze()
def apply_to_camera(self, camera: Cameras) -> None:
"""Apply the pose correction to the raybundle"""
if self.config.mode != "off":
assert camera.metadata is not None, "Must provide id of camera in its metadata"
assert "cam_idx" in camera.metadata, "Must provide id of camera in its metadata"
camera_idx = camera.metadata["cam_idx"]
adj = self(torch.tensor([camera_idx], dtype=torch.long, device=camera.device)) # type: ignore
adj = torch.cat([adj, torch.Tensor([0, 0, 0, 1])[None, None].to(adj)], dim=1)
camera.camera_to_worlds = torch.bmm(camera.camera_to_worlds, adj)
def apply_to_camera(self, camera: Cameras) -> torch.Tensor:
"""Apply the pose correction to the world-to-camera matrix in a Camera object"""
if self.config.mode == "off":
return camera.camera_to_worlds
assert camera.metadata is not None, "Must provide id of camera in its metadata"
if "cam_idx" not in camera.metadata:
# Evalutaion cams?
return camera.camera_to_worlds
camera_idx = camera.metadata["cam_idx"]
adj = self(torch.tensor([camera_idx], dtype=torch.long, device=camera.device)) # type: ignore
adj = torch.cat([adj, torch.Tensor([0, 0, 0, 1])[None, None].to(adj)], dim=1)
return torch.bmm(camera.camera_to_worlds, adj)
def get_loss_dict(self, loss_dict: dict) -> None:
"""Add regularization"""
@ -176,8 +182,12 @@ class CameraOptimizer(nn.Module):
def get_metrics_dict(self, metrics_dict: dict) -> None:
"""Get camera optimizer metrics"""
if self.config.mode != "off":
metrics_dict["camera_opt_translation"] = self.pose_adjustment[:, :3].norm()
metrics_dict["camera_opt_rotation"] = self.pose_adjustment[:, 3:].norm()
trans = self.pose_adjustment[:, :3].detach().norm(dim=-1)
rot = self.pose_adjustment[:, 3:].detach().norm(dim=-1)
metrics_dict["camera_opt_translation_max"] = trans.max()
metrics_dict["camera_opt_translation_mean"] = trans.mean()
metrics_dict["camera_opt_rotation_mean"] = numpy.rad2deg(rot.mean().cpu())
metrics_dict["camera_opt_rotation_max"] = numpy.rad2deg(rot.max().cpu())
def get_param_groups(self, param_groups: dict) -> None:
"""Get camera optimizer parameters"""

12
nerfstudio/configs/method_configs.py
View File

@ -632,8 +632,10 @@ method_configs["splatfacto"] = TrainerConfig(
},
"quats": {"optimizer": AdamOptimizerConfig(lr=0.001, eps=1e-15), "scheduler": None},
"camera_opt": {
"optimizer": AdamOptimizerConfig(lr=1e-3, eps=1e-15),
"scheduler": ExponentialDecaySchedulerConfig(lr_final=5e-5, max_steps=30000),
"optimizer": AdamOptimizerConfig(lr=1e-4, eps=1e-15),
"scheduler": ExponentialDecaySchedulerConfig(
lr_final=5e-7, max_steps=30000, warmup_steps=1000, lr_pre_warmup=0
),
},
},
viewer=ViewerConfig(num_rays_per_chunk=1 << 15),
@ -684,8 +686,10 @@ method_configs["splatfacto-big"] = TrainerConfig(
},
"quats": {"optimizer": AdamOptimizerConfig(lr=0.001, eps=1e-15), "scheduler": None},
"camera_opt": {
"optimizer": AdamOptimizerConfig(lr=1e-3, eps=1e-15),
"scheduler": ExponentialDecaySchedulerConfig(lr_final=5e-5, max_steps=30000),
"optimizer": AdamOptimizerConfig(lr=1e-4, eps=1e-15),
"scheduler": ExponentialDecaySchedulerConfig(
lr_final=5e-7, max_steps=30000, warmup_steps=1000, lr_pre_warmup=0
),
},
},
viewer=ViewerConfig(num_rays_per_chunk=1 << 15),

29
nerfstudio/models/splatfacto.py
View File

@ -33,6 +33,7 @@ from pytorch_msssim import SSIM
from torch.nn import Parameter
from typing_extensions import Literal
from nerfstudio.cameras.camera_optimizers import CameraOptimizer, CameraOptimizerConfig
from nerfstudio.cameras.cameras import Cameras
from nerfstudio.data.scene_box import OrientedBox
from nerfstudio.engine.callbacks import TrainingCallback, TrainingCallbackAttributes, TrainingCallbackLocation
@ -146,6 +147,8 @@ class SplatfactoModelConfig(ModelConfig):
However, PLY exported with antialiased rasterize mode is not compatible with classic mode. Thus many web viewers that
were implemented for classic mode can not render antialiased mode PLY properly without modifications.
"""
camera_optimizer: CameraOptimizerConfig = field(default_factory=lambda: CameraOptimizerConfig(mode="off"))
"""Config of the camera optimizer to use"""
class SplatfactoModel(Model):
@ -213,6 +216,10 @@ class SplatfactoModel(Model):
}
)
self.camera_optimizer: CameraOptimizer = self.config.camera_optimizer.setup(
num_cameras=self.num_train_data, device="cpu"
)
# metrics
from torchmetrics.image import PeakSignalNoiseRatio
from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity
@ -609,6 +616,7 @@ class SplatfactoModel(Model):
Mapping of different parameter groups
"""
gps = self.get_gaussian_param_groups()
self.camera_optimizer.get_param_groups(param_groups=gps)
return gps
def _get_downscale_factor(self):
@ -648,6 +656,8 @@ class SplatfactoModel(Model):
# get the background color
if self.training:
optimized_camera_to_world = self.camera_optimizer.apply_to_camera(camera)[0, ...]
if self.config.background_color == "random":
background = torch.rand(3, device=self.device)
elif self.config.background_color == "white":
@ -657,6 +667,8 @@ class SplatfactoModel(Model):
else:
background = self.background_color.to(self.device)
else:
optimized_camera_to_world = camera.camera_to_worlds[0, ...]
if renderers.BACKGROUND_COLOR_OVERRIDE is not None:
background = renderers.BACKGROUND_COLOR_OVERRIDE.to(self.device)
else:
@ -674,8 +686,9 @@ class SplatfactoModel(Model):
camera_downscale = self._get_downscale_factor()
camera.rescale_output_resolution(1 / camera_downscale)
# shift the camera to center of scene looking at center
R = camera.camera_to_worlds[0, :3, :3] # 3 x 3
T = camera.camera_to_worlds[0, :3, 3:4] # 3 x 1
R = optimized_camera_to_world[:3, :3] # 3 x 3
T = optimized_camera_to_world[:3, 3:4] # 3 x 1
# flip the z and y axes to align with gsplat conventions
R_edit = torch.diag(torch.tensor([1, -1, -1], device=self.device, dtype=R.dtype))
R = R @ R_edit
@ -738,7 +751,7 @@ class SplatfactoModel(Model):
self.xys.retain_grad()
if self.config.sh_degree > 0:
viewdirs = means_crop.detach() - camera.camera_to_worlds.detach()[..., :3, 3] # (N, 3)
viewdirs = means_crop.detach() - optimized_camera_to_world.detach()[:3, 3] # (N, 3)
viewdirs = viewdirs / viewdirs.norm(dim=-1, keepdim=True)
n = min(self.step // self.config.sh_degree_interval, self.config.sh_degree)
rgbs = spherical_harmonics(n, viewdirs, colors_crop)
@ -829,6 +842,8 @@ class SplatfactoModel(Model):
metrics_dict["psnr"] = self.psnr(predicted_rgb, gt_rgb)
metrics_dict["gaussian_count"] = self.num_points
self.camera_optimizer.get_metrics_dict(metrics_dict)
return metrics_dict
def get_loss_dict(self, outputs, batch, metrics_dict=None) -> Dict[str, torch.Tensor]:
@ -867,11 +882,17 @@ class SplatfactoModel(Model):
else:
scale_reg = torch.tensor(0.0).to(self.device)
return {
loss_dict = {
"main_loss": (1 - self.config.ssim_lambda) * Ll1 + self.config.ssim_lambda * simloss,
"scale_reg": scale_reg,
}
if self.training:
# Add loss from camera optimizer
self.camera_optimizer.get_loss_dict(loss_dict)
return loss_dict
@torch.no_grad()
def get_outputs_for_camera(self, camera: Cameras, obb_box: Optional[OrientedBox] = None) -> Dict[str, torch.Tensor]:
"""Takes in a camera, generates the raybundle, and computes the output of the model.