diff --git a/nerfstudio/cameras/camera_optimizers.py b/nerfstudio/cameras/camera_optimizers.py index 891cd272a..5e1897659 100644 --- a/nerfstudio/cameras/camera_optimizers.py +++ b/nerfstudio/cameras/camera_optimizers.py @@ -83,6 +83,7 @@ class CameraOptimizer(nn.Module): pass elif self.config.mode in ("SO3xR3", "SE3"): self.pose_adjustment = torch.nn.Parameter(torch.zeros((num_cameras, 6), device=device)) + self.blur_adjustment = torch.nn.Parameter(torch.zeros((num_cameras, 6), device=device)) else: assert_never(self.config.mode) @@ -128,3 +129,18 @@ class CameraOptimizer(nn.Module): # 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) + + def sample_blur_correction(self, indices: TensorType["num_cameras"]): + """ + Samples the blur parameters and returns transformation matrices for the origins and directions + """ + blur_adjs = self.blur_adjustment[indices, :] * torch.normal( + torch.zeros((indices.shape[0], 1), device=self.device), # mean + torch.ones((indices.shape[0], 1), device=self.device), # std + ) + # print( + # "trans,rot norm", + # torch.mean(torch.linalg.norm(blur_adjs[:, :3], dim=0)), + # torch.mean(torch.linalg.norm(blur_adjs[:, 3:], dim=0)), + # ) + return exp_map_SO3xR3(blur_adjs) diff --git a/nerfstudio/data/datamanagers/base_datamanager.py b/nerfstudio/data/datamanagers/base_datamanager.py index 82cf4a0d5..0e971d0c1 100644 --- a/nerfstudio/data/datamanagers/base_datamanager.py +++ b/nerfstudio/data/datamanagers/base_datamanager.py @@ -51,6 +51,7 @@ from nerfstudio.data.utils.dataloaders import ( from nerfstudio.engine.callbacks import TrainingCallback, TrainingCallbackAttributes from nerfstudio.model_components.ray_generators import RayGenerator from nerfstudio.utils.misc import IterableWrapper +from nerfstudio.utils.poses import to4x4 CONSOLE = Console(width=120) @@ -329,10 +330,7 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method self.train_camera_optimizer = self.config.camera_optimizer.setup( num_cameras=self.train_dataset.dataparser_outputs.cameras.size, device=self.device ) - self.train_ray_generator = RayGenerator( - self.train_dataset.dataparser_outputs.cameras.to(self.device), - self.train_camera_optimizer, - ) + self.train_ray_generator = RayGenerator(self.train_dataset.dataparser_outputs.cameras.to(self.device)) def setup_eval(self): """Sets up the data loader for evaluation""" @@ -348,10 +346,7 @@ 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.dataparser_outputs.cameras.to(self.device), - self.train_camera_optimizer, # should be shared between train and eval. - ) + self.eval_ray_generator = RayGenerator(self.eval_dataset.dataparser_outputs.cameras.to(self.device)) # for loading full images self.fixed_indices_eval_dataloader = FixedIndicesEvalDataloader( input_dataset=self.eval_dataset, @@ -372,6 +367,25 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method batch = self.train_pixel_sampler.sample(image_batch) ray_indices = batch["indices"] ray_bundle = self.train_ray_generator(ray_indices) + + # Multiply in the camera optimization deltas to the rays + cam_deltas = to4x4(self.train_camera_optimizer(ray_indices[:, 0])) # B x 3 x 4 + blur_deltas = self.train_camera_optimizer.sample_blur_correction(ray_indices[:, 0]) + # TODO(Justin) is this allowed mathematically + cam_deltas = torch.bmm(blur_deltas, cam_deltas) + origins_homogeneous = torch.cat( + [ + ray_bundle.origins, + torch.ones((*ray_bundle.shape, 1), dtype=ray_bundle.origins.dtype, device=ray_bundle.origins.device), + ], + dim=-1, + ) + origins_homogeneous = origins_homogeneous.view(*origins_homogeneous.shape, 1) + ray_bundle.origins = torch.bmm(cam_deltas, origins_homogeneous)[..., :3, 0] + ray_bundle.directions = torch.bmm( + cam_deltas[:, :, :3], ray_bundle.directions.view(*ray_bundle.directions.shape, 1) + )[..., 0] + return ray_bundle, batch def next_eval(self, step: int) -> Tuple[RayBundle, Dict]: @@ -381,6 +395,20 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method batch = self.eval_pixel_sampler.sample(image_batch) ray_indices = batch["indices"] ray_bundle = self.eval_ray_generator(ray_indices) + # Multiply in the camera optimization deltas to the rays + cam_deltas = self.train_camera_optimizer(ray_indices[:, 0]) # B x 3 x 4 + origins_homogeneous = torch.cat( + [ + ray_bundle.origins, + torch.ones((*ray_bundle.shape, 1), dtype=ray_bundle.origins.dtype, device=ray_bundle.origins.device), + ], + dim=-1, + ) + origins_homogeneous = origins_homogeneous.view(*origins_homogeneous.shape, 1) + ray_bundle.origins = torch.bmm(cam_deltas, origins_homogeneous)[..., :3, 0] + ray_bundle.directions = torch.bmm( + cam_deltas[:, :, :3], ray_bundle.directions.view(*ray_bundle.directions.shape, 1) + )[..., 0] return ray_bundle, batch def next_eval_image(self, step: int) -> Tuple[int, RayBundle, Dict]: diff --git a/nerfstudio/model_components/ray_generators.py b/nerfstudio/model_components/ray_generators.py index 6e79cd988..50070dfb2 100644 --- a/nerfstudio/model_components/ray_generators.py +++ b/nerfstudio/model_components/ray_generators.py @@ -15,6 +15,7 @@ """ Ray generator. """ +import torch from torch import nn from torchtyping import TensorType @@ -32,10 +33,9 @@ class RayGenerator(nn.Module): pose_optimizer: pose optimization module, for optimizing noisy camera intrisics/extrinsics. """ - def __init__(self, cameras: Cameras, pose_optimizer: CameraOptimizer) -> None: + def __init__(self, cameras: Cameras) -> 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: @@ -44,16 +44,11 @@ class RayGenerator(nn.Module): Args: ray_indices: Contains camera, row, and col indicies for target rays. """ - c = ray_indices[:, 0] # camera indices - y = ray_indices[:, 1] # row indices - x = ray_indices[:, 2] # col indices - coords = self.image_coords[y, x] + with torch.no_grad(): + c = ray_indices[:, 0] # camera indices + y = ray_indices[:, 1] # row indices + 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, - ) + ray_bundle = self.cameras.generate_rays(camera_indices=c, coords=coords) return ray_bundle