tink/nerfstudio
4
0
Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 02:24:38 +08:00
Code

first try

Browse Source
This commit is contained in:
Justin Kerr 2022-11-06 20:58:11 -08:00
parent b86a21caf2
commit f3e31e2e14
4 changed files with 66 additions and 33 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
nerfstudio/cameras/camera_utils.py
View File

@ -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.

7
nerfstudio/cameras/cameras.py
View File

@ -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):

84
nerfstudio/data/datamanagers/base_datamanager.py
View File

@ -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]:

6
nerfstudio/model_components/ray_generators.py
View File

@ -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