From f3e31e2e14d5fced345ffedcf484849320e6c76b Mon Sep 17 00:00:00 2001 From: Justin Kerr Date: Sun, 6 Nov 2022 20:58:11 -0800 Subject: [PATCH] first try --- nerfstudio/cameras/camera_utils.py | 2 +- nerfstudio/cameras/cameras.py | 7 +- .../data/datamanagers/base_datamanager.py | 84 ++++++++++++++----- nerfstudio/model_components/ray_generators.py | 6 +- 4 files changed, 66 insertions(+), 33 deletions(-) diff --git a/nerfstudio/cameras/camera_utils.py b/nerfstudio/cameras/camera_utils.py index b0de30a9b..30d755596 100644 --- a/nerfstudio/cameras/camera_utils.py +++ b/nerfstudio/cameras/camera_utils.py @@ -342,7 +342,7 @@ def _compute_residual_and_jacobian( def radial_and_tangential_undistort( coords: torch.Tensor, distortion_params: torch.Tensor, - eps: float = 1e-9, + eps: float = 1e-3, max_iterations: int = 10, ) -> torch.Tensor: """Computes undistorted coords given opencv distortion parameters. diff --git a/nerfstudio/cameras/cameras.py b/nerfstudio/cameras/cameras.py index 7cf107f2f..2b547b2e7 100644 --- a/nerfstudio/cameras/cameras.py +++ b/nerfstudio/cameras/cameras.py @@ -252,7 +252,6 @@ class Cameras: self, camera_indices: Union[TensorType["num_rays":...], int], coords: Optional[TensorType["num_rays":..., 2]] = 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. @@ -316,8 +315,6 @@ class Cameras: raise ValueError(f"Camera type {CameraType(self.camera_type[0])} not supported.") c2w = self.camera_to_worlds[camera_indices] - 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 @@ -328,8 +325,8 @@ class Cameras: origins = c2w[..., :3, 3] # (..., 3) directions = directions_stack[0] - dx = torch.sqrt(torch.sum((directions - directions_stack[1]) ** 2, dim=-1)) - dy = torch.sqrt(torch.sum((directions - directions_stack[2]) ** 2, dim=-1)) + dx = torch.norm((directions - directions_stack[1]), dim=-1) + dy = torch.norm((directions - directions_stack[2]), dim=-1) pixel_area = (dx * dy)[..., None] if not isinstance(camera_indices, torch.Tensor): diff --git a/nerfstudio/data/datamanagers/base_datamanager.py b/nerfstudio/data/datamanagers/base_datamanager.py index 82cf4a0d5..9d82eb2ab 100644 --- a/nerfstudio/data/datamanagers/base_datamanager.py +++ b/nerfstudio/data/datamanagers/base_datamanager.py @@ -264,7 +264,39 @@ class VanillaDataManagerConfig(InstantiateConfig): """Specifies the camera pose optimizer used during training. Helpful if poses are noisy, such as for data from Record3D.""" - +import torch.multiprocessing as mp +MP_CON = mp.get_context('spawn') +class Runner(MP_CON.Process): + def __init__(self, bundle_queue, train_dataset, device, config): + super().__init__() + self.train_dataset = train_dataset + self.device = device + self.bundle_queue = bundle_queue + self.config = config + self.daemon=True + self.start() + + def run(self): + self.train_image_dataloader = CacheDataloader( + self.train_dataset, + num_images_to_sample_from=self.config.train_num_images_to_sample_from, + num_times_to_repeat_images=self.config.train_num_times_to_repeat_images, + device='cpu', + num_workers=1, + pin_memory=True, + ) + 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.dataparser_outputs.cameras.to('cpu') + ) + while True: + image_batch = next(self.iter_train_image_dataloader) + batch = self.train_pixel_sampler.sample(image_batch) + ray_indices = batch["indices"] + ray_bundle = self.train_ray_generator(ray_indices) + self.bundle_queue.put((ray_bundle,batch),block=True) + class VanillaDataManager(DataManager): # pylint: disable=abstract-method """Basic stored data manager implementation. @@ -300,7 +332,9 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method self.train_dataset = self.create_train_dataset() self.eval_dataset = self.create_eval_dataset() - super().__init__() + self.bundle_queue = MP_CON.Queue(5) + runners = [Runner(self.bundle_queue, self.train_dataset, self.device,self.config) for i in range(3)] + DataManager.__init__(self) def create_train_dataset(self) -> InputDataset: """Sets up the data loaders for training""" @@ -311,28 +345,14 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method return InputDataset( self.config.dataparser.setup().get_dataparser_outputs(split="val" if not self.test_mode else "test") ) - + def setup_train(self): """Sets up the data loaders for training""" assert self.train_dataset is not None CONSOLE.print("Setting up training dataset...") - self.train_image_dataloader = CacheDataloader( - self.train_dataset, - num_images_to_sample_from=self.config.train_num_images_to_sample_from, - num_times_to_repeat_images=self.config.train_num_times_to_repeat_images, - device=self.device, - num_workers=self.world_size * 4, - pin_memory=True, - ) - self.iter_train_image_dataloader = iter(self.train_image_dataloader) - self.train_pixel_sampler = PixelSampler(self.config.train_num_rays_per_batch) 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, - ) def setup_eval(self): """Sets up the data loader for evaluation""" @@ -349,8 +369,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_dataset.dataparser_outputs.cameras.to(self.device) ) # for loading full images self.fixed_indices_eval_dataloader = FixedIndicesEvalDataloader( @@ -368,10 +387,21 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method def next_train(self, step: int) -> Tuple[RayBundle, Dict]: """Returns the next batch of data from the train dataloader.""" self.train_count += 1 - image_batch = next(self.iter_train_image_dataloader) - batch = self.train_pixel_sampler.sample(image_batch) - ray_indices = batch["indices"] - ray_bundle = self.train_ray_generator(ray_indices) + # image_batch = next(self.iter_train_image_dataloader) + # batch = self.train_pixel_sampler.sample(image_batch) + # ray_indices = batch["indices"] + # ray_bundle = self.train_ray_generator(ray_indices) + ray_bundle,batch = self.bundle_queue.get(block=True) + ray_bundle = ray_bundle.to(self.device) + ray_indices = ray_bundle.camera_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(self, step: int) -> Tuple[RayBundle, Dict]: @@ -381,6 +411,14 @@ 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..b9658466a 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: @@ -49,11 +49,9 @@ 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