From 1a74c6a21256ac52311c4574e7792a76643da581 Mon Sep 17 00:00:00 2001 From: Kevin Chen <79341521+kevinddchen@users.noreply.github.com> Date: Wed, 16 Nov 2022 05:44:29 +0000 Subject: [PATCH] Cleanup --- nerfstudio/cameras/cameras.py | 4 +- nerfstudio/data/datamanagers.py | 2 +- nerfstudio/data/pixel_samplers.py | 155 +++++++++++++++--------------- scripts/process_data.py | 7 +- 4 files changed, 87 insertions(+), 81 deletions(-) diff --git a/nerfstudio/cameras/cameras.py b/nerfstudio/cameras/cameras.py index 858745a49..bd5b4e13b 100644 --- a/nerfstudio/cameras/cameras.py +++ b/nerfstudio/cameras/cameras.py @@ -275,8 +275,8 @@ class Cameras: coords = self.get_image_coords().to(self.device) assert coords is not None - y = coords[..., 0] # (..., 1) - x = coords[..., 1] # (..., 1) + y = coords[..., 0] + x = coords[..., 1] fx, fy = self.fx[camera_indices], self.fy[camera_indices] cx, cy = self.cx[camera_indices], self.cy[camera_indices] diff --git a/nerfstudio/data/datamanagers.py b/nerfstudio/data/datamanagers.py index c8b5a40da..c3cf367c3 100644 --- a/nerfstudio/data/datamanagers.py +++ b/nerfstudio/data/datamanagers.py @@ -38,7 +38,7 @@ from nerfstudio.data.dataparsers.instant_ngp_dataparser import ( ) from nerfstudio.data.dataparsers.nerfstudio_dataparser import NerfstudioDataParserConfig from nerfstudio.data.dataparsers.record3d_dataparser import Record3DDataParserConfig -from nerfstudio.data.pixel_samplers import EquirectangularSampler, PixelSampler +from nerfstudio.data.pixel_samplers import PixelSampler from nerfstudio.data.utils.dataloaders import ( CacheDataloader, FixedIndicesEvalDataloader, diff --git a/nerfstudio/data/pixel_samplers.py b/nerfstudio/data/pixel_samplers.py index 0ec2ccdf9..432a8210f 100644 --- a/nerfstudio/data/pixel_samplers.py +++ b/nerfstudio/data/pixel_samplers.py @@ -110,93 +110,96 @@ class PixelSampler: # pylint: disable=too-few-public-methods ) return pixel_batch -def collate_image_dataset_batch_equirectangular(batch: Dict, num_rays_per_batch: int, keep_full_image: bool = False): - """ - Operates on a batch of images and samples pixels to use for generating rays. - Returns a collated batch which is input to the Graph. - It will sample only within the valid 'mask' if it's specified. - Args: - batch: batch of images to sample from - num_rays_per_batch: number of rays to sample per batch - keep_full_image: whether or not to include a reference to the full image in returned batch - """ - device = batch["image"].device - num_images, image_height, image_width, _ = batch["image"].shape +# TODO(kevinddchen): Implement equirect pixel sampler - # only sample within the mask, if the mask is in the batch - # if "mask" in batch: - # nonzero_indices = torch.nonzero(batch["mask"][..., 0], as_tuple=False) - # chosen_indices = random.sample(range(len(nonzero_indices)), k=num_rays_per_batch) - # indices = nonzero_indices[chosen_indices] - # else: - # indices = torch.floor( - # torch.rand((num_rays_per_batch, 3), device=device) - # * torch.tensor([num_images, image_height, image_width], device=device) - # ).long() +# def collate_image_dataset_batch_equirectangular(batch: Dict, num_rays_per_batch: int, keep_full_image: bool = False): +# """ +# Operates on a batch of images and samples pixels to use for generating rays. +# Returns a collated batch which is input to the Graph. +# It will sample only within the valid 'mask' if it's specified. - number_of_images = torch.rand(num_rays_per_batch, device=device) - phi = torch.acos(1 - 2 * torch.rand(num_rays_per_batch, device=device)) / torch.pi - theta = torch.rand(num_rays_per_batch, device=device) - indices = torch.floor( - torch.stack((number_of_images, phi, theta), dim=-1) - * torch.tensor([num_images, image_height, image_width], device=device) - ).long() +# Args: +# batch: batch of images to sample from +# num_rays_per_batch: number of rays to sample per batch +# keep_full_image: whether or not to include a reference to the full image in returned batch +# """ +# device = batch["image"].device +# num_images, image_height, image_width, _ = batch["image"].shape - c, y, x = (i.flatten() for i in torch.split(indices, 1, dim=-1)) - image = batch["image"][c, y, x] - # mask, semantics = None, None - # if "mask" in batch: - # mask = batch["mask"][c, y, x] - # if "semantics" in batch: - # semantics = batch["semantics"][c, y, x] - # assert image.shape == (num_rays_per_batch, 3), image.shape +# # only sample within the mask, if the mask is in the batch +# # if "mask" in batch: +# # nonzero_indices = torch.nonzero(batch["mask"][..., 0], as_tuple=False) +# # chosen_indices = random.sample(range(len(nonzero_indices)), k=num_rays_per_batch) +# # indices = nonzero_indices[chosen_indices] +# # else: +# # indices = torch.floor( +# # torch.rand((num_rays_per_batch, 3), device=device) +# # * torch.tensor([num_images, image_height, image_width], device=device) +# # ).long() - # Needed to correct the random indices to their actual camera idx locations. - local_indices = indices.clone() - indices[:, 0] = batch["image_idx"][c] - collated_batch = { - "local_indices": local_indices, # local to the batch returned - "indices": indices, # with the abs camera indices - "image": image, - } - # if mask is not None: - # collated_batch["mask"] = mask - # if semantics is not None: - # collated_batch["semantics"] = semantics +# number_of_images = torch.rand(num_rays_per_batch, device=device) +# phi = torch.acos(1 - 2 * torch.rand(num_rays_per_batch, device=device)) / torch.pi +# theta = torch.rand(num_rays_per_batch, device=device) +# indices = torch.floor( +# torch.stack((number_of_images, phi, theta), dim=-1) +# * torch.tensor([num_images, image_height, image_width], device=device) +# ).long() - # if keep_full_image: - # collated_batch["full_image"] = batch["image"] +# c, y, x = (i.flatten() for i in torch.split(indices, 1, dim=-1)) +# image = batch["image"][c, y, x] +# # mask, semantics = None, None +# # if "mask" in batch: +# # mask = batch["mask"][c, y, x] +# # if "semantics" in batch: +# # semantics = batch["semantics"][c, y, x] +# # assert image.shape == (num_rays_per_batch, 3), image.shape - return collated_batch +# # Needed to correct the random indices to their actual camera idx locations. +# local_indices = indices.clone() +# indices[:, 0] = batch["image_idx"][c] +# collated_batch = { +# "local_indices": local_indices, # local to the batch returned +# "indices": indices, # with the abs camera indices +# "image": image, +# } +# # if mask is not None: +# # collated_batch["mask"] = mask +# # if semantics is not None: +# # collated_batch["semantics"] = semantics -class EquirectangularSampler: # pylint: disable=too-few-public-methods - """Samples 'pixel_batch's from 'image_batch's. +# # if keep_full_image: +# # collated_batch["full_image"] = batch["image"] - Args: - num_rays_per_batch: number of rays to sample per batch - keep_full_image: whether or not to include a reference to the full image in returned batch - """ +# return collated_batch - def __init__(self, num_rays_per_batch: int, keep_full_image: bool = False) -> None: - self.num_rays_per_batch = num_rays_per_batch - self.keep_full_image = keep_full_image +# class EquirectangularSampler: # pylint: disable=too-few-public-methods +# """Samples 'pixel_batch's from 'image_batch's. - def set_num_rays_per_batch(self, num_rays_per_batch: int): - """Set the number of rays to sample per batch. +# Args: +# num_rays_per_batch: number of rays to sample per batch +# keep_full_image: whether or not to include a reference to the full image in returned batch +# """ - Args: - num_rays_per_batch: number of rays to sample per batch - """ - self.num_rays_per_batch = num_rays_per_batch +# def __init__(self, num_rays_per_batch: int, keep_full_image: bool = False) -> None: +# self.num_rays_per_batch = num_rays_per_batch +# self.keep_full_image = keep_full_image - def sample(self, image_batch: Dict): - """Sample an image batch and return a pixel batch. +# def set_num_rays_per_batch(self, num_rays_per_batch: int): +# """Set the number of rays to sample per batch. - Args: - image_batch: batch of images to sample from - """ - pixel_batch = collate_image_dataset_batch_equirectangular( - image_batch, self.num_rays_per_batch, keep_full_image=self.keep_full_image - ) - return pixel_batch \ No newline at end of file +# Args: +# num_rays_per_batch: number of rays to sample per batch +# """ +# self.num_rays_per_batch = num_rays_per_batch + +# def sample(self, image_batch: Dict): +# """Sample an image batch and return a pixel batch. + +# Args: +# image_batch: batch of images to sample from +# """ +# pixel_batch = collate_image_dataset_batch_equirectangular( +# image_batch, self.num_rays_per_batch, keep_full_image=self.keep_full_image +# ) +# return pixel_batch diff --git a/scripts/process_data.py b/scripts/process_data.py index 2810efefb..ee6af4b91 100755 --- a/scripts/process_data.py +++ b/scripts/process_data.py @@ -260,7 +260,7 @@ def copy_images(data: Path, image_dir: Path, verbose: bool, rename: bool = True) data: Path to the directory of images. image_dir: Path to the output directory. verbose: If True, print extra logging. - rename: If True, rename the images to be sequential. + rename: If True, rename the images sequentially to `frame_{num}`. Returns: The number of images copied. """ @@ -578,7 +578,7 @@ def opensfm_to_json(reconstruction_path: Path, output_dir: Path, camera_model: C aa_vec = np.array(shot_info["rotation"]) # 3D axis angle repr angle = np.linalg.norm(aa_vec) if angle > 1e-8: - # normalilze the axis-angle repr if angle is large enough + # normalize the axis-angle repr if angle is large enough aa_vec /= angle qx = aa_vec[0] * np.sin(angle / 2) qy = aa_vec[1] * np.sin(angle / 2) @@ -594,6 +594,9 @@ def opensfm_to_json(reconstruction_path: Path, output_dir: Path, camera_model: C # Convert camera +z forwards (OpenSfM) convention to -z forwards (NeRF) convention # Equivalent to right-multiplication by diag([1, -1, -1, 1]) c2w[0:3, 1:3] *= -1 + # Rotation around global z-axis by -90 degrees + c2w = c2w[np.array([1, 0, 2, 3]), :] + c2w[2, :] *= -1 name = Path(f"./images/{shot_name}")