From 4120d7945dbb5837173581122a8910a807408258 Mon Sep 17 00:00:00 2001 From: Kamyar Salahi Date: Thu, 3 Nov 2022 19:14:15 -0700 Subject: [PATCH] preliminary equirectangular commit --- nerfstudio/cameras/cameras.py | 13 +++++ nerfstudio/data/datamanagers.py | 4 +- nerfstudio/data/pixel_samplers.py | 92 +++++++++++++++++++++++++++++++ scripts/process_data.py | 32 ++++++----- 4 files changed, 126 insertions(+), 15 deletions(-) diff --git a/nerfstudio/cameras/cameras.py b/nerfstudio/cameras/cameras.py index 8d66e3ae4..4cfb89c1d 100644 --- a/nerfstudio/cameras/cameras.py +++ b/nerfstudio/cameras/cameras.py @@ -36,6 +36,7 @@ class CameraType(Enum): PERSPECTIVE = auto() FISHEYE = auto() + EQUIRECTANGULAR = auto() CAMERA_MODEL_TO_TYPE = { @@ -45,6 +46,7 @@ CAMERA_MODEL_TO_TYPE = { "RADIAL": CameraType.PERSPECTIVE, "OPENCV": CameraType.PERSPECTIVE, "OPENCV_FISHEYE": CameraType.FISHEYE, + "EQUIRECTANGULAR": CameraType.EQUIRECTANGULAR } @@ -312,6 +314,17 @@ class Cameras: ], dim=-1, ) + elif self.camera_type[0] == CameraType.EQUIRECTANGULAR.value: + phi = y / self.image_height * torch.pi + theta = -x / self.image_width * 2 * torch.pi + directions = torch.stack( + [ + torch.cos(theta) * torch.sin(phi), + torch.sin(theta) * torch.sin(phi), + torch.cos(phi) + ], + dim=-1, + ) else: raise ValueError(f"Camera type {CameraType(self.camera_type[0])} not supported.") diff --git a/nerfstudio/data/datamanagers.py b/nerfstudio/data/datamanagers.py index c3cf367c3..b98d5d48c 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 PixelSampler +from nerfstudio.data.pixel_samplers import PixelSampler, EquirectangularSampler from nerfstudio.data.utils.dataloaders import ( CacheDataloader, FixedIndicesEvalDataloader, @@ -302,7 +302,7 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method 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_pixel_sampler = EquirectangularSampler(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 ) diff --git a/nerfstudio/data/pixel_samplers.py b/nerfstudio/data/pixel_samplers.py index 8ba19ba63..b2e3122d8 100644 --- a/nerfstudio/data/pixel_samplers.py +++ b/nerfstudio/data/pixel_samplers.py @@ -109,3 +109,95 @@ class PixelSampler: # pylint: disable=too-few-public-methods image_batch, self.num_rays_per_batch, keep_full_image=self.keep_full_image ) 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 + + # 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() + + number_of_images = torch.rand(num_rays_per_batch, device=device) + phi = torch.acos(1 - 2 * torch.rand(num_rays_per_batch, device=device)) / 2 * 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() + + 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 + + # 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 + + # if keep_full_image: + # collated_batch["full_image"] = batch["image"] + + return collated_batch + +class EquirectangularSampler: # pylint: disable=too-few-public-methods + """Samples 'pixel_batch's from 'image_batch's. + + 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 + """ + + 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 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 + """ + 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 \ No newline at end of file diff --git a/scripts/process_data.py b/scripts/process_data.py index c0f457760..82774abf5 100755 --- a/scripts/process_data.py +++ b/scripts/process_data.py @@ -31,11 +31,13 @@ class CameraModel(Enum): OPENCV = "OPENCV" OPENCV_FISHEYE = "OPENCV_FISHEYE" + EQUIRECTANGULAR = "EQUIRECTANGULAR" CAMERA_MODELS = { "perspective": CameraModel.OPENCV, "fisheye": CameraModel.OPENCV_FISHEYE, + "equirectangular": CameraModel.EQUIRECTANGULAR, } @@ -225,7 +227,8 @@ def convert_insta360_to_images( else: CONSOLE.print("[bold red]Can't satify requested number of frames. Extracting all frames.") - vf_cmds.append(f"crop=iw*({crop_percentage}):ih*({crop_percentage})") + # vf_cmds.append(f"crop=iw*({crop_percentage}):ih*({crop_percentage})") + vf_cmds.append("v360=dfisheye:equirect:ih_fov=190:iv_fov=190:yaw=-90") front_vf_cmds = vf_cmds + ["transpose=2"] back_vf_cmds = vf_cmds + ["transpose=1"] @@ -510,6 +513,8 @@ def run_opensfm( proj_type = "brown" elif camera_model == CameraModel.OPENCV_FISHEYE: proj_type = "fisheye_opencv" + elif camera_model == CameraModel.EQUIRECTANGULAR: + proj_type = "equirectangular" cam_overrides = { "all": { "projection_type": proj_type, @@ -675,7 +680,7 @@ class ProcessImages: """Path the data, either a video file or a directory of images.""" output_dir: Path """Path to the output directory.""" - camera_type: Literal["perspective", "fisheye"] = "perspective" + camera_type: Literal["perspective", "fisheye", "equirectangular"] = "perspective" """Camera model to use.""" sfm_method: Literal["colmap", "OpenSfM"] = "OpenSfM" """Which sfm solver to use""" @@ -782,7 +787,7 @@ class ProcessVideo: """Path to the output directory.""" num_frames_target: int = 150 """Target number of frames to use for the dataset, results may not be exact.""" - camera_type: Literal["perspective", "fisheye"] = "perspective" + camera_type: Literal["perspective", "fisheye", "equirectangular"] = "perspective" """Camera model to use.""" sfm_method: Literal["colmap", "OpenSfM"] = "OpenSfM" """Which sfm solver to use""" @@ -930,27 +935,28 @@ class ProcessInsta360: summary_log.append(downscale_images(image_dir, self.num_downscales, verbose=self.verbose)) # Run Colmap - colmap_dir = self.output_dir / "colmap" + opensfm_dir = self.output_dir / "opensfm" if not self.skip_colmap: - colmap_dir.mkdir(parents=True, exist_ok=True) + opensfm_dir.mkdir(parents=True, exist_ok=True) - run_colmap( + run_opensfm( image_dir=image_dir, - colmap_dir=colmap_dir, - camera_model=CAMERA_MODELS["fisheye"], + opensfm_dir=opensfm_dir, + camera_model=CAMERA_MODELS["equirectangular"], gpu=self.gpu, verbose=self.verbose, matching_method=self.matching_method, + opensfm_install= ) # Save transforms.json - if (colmap_dir / "sparse" / "0" / "cameras.bin").exists(): + if (opensfm_dir / "sparse" / "0" / "cameras.bin").exists(): with CONSOLE.status("[bold yellow]Saving results to transforms.json", spinner="balloon"): - num_matched_frames = colmap_to_json( - cameras_path=colmap_dir / "sparse" / "0" / "cameras.bin", - images_path=colmap_dir / "sparse" / "0" / "images.bin", + num_matched_frames = opensfm_to_json( + cameras_path=opensfm_dir / "sparse" / "0" / "cameras.bin", + images_path=opensfm_dir / "sparse" / "0" / "images.bin", output_dir=self.output_dir, - camera_model=CAMERA_MODELS["fisheye"], + camera_model=CAMERA_MODELS["equirectangular"], ) summary_log.append(f"Colmap matched {num_matched_frames} images") else: