preliminary equirectangular commit

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.")
 

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
         )

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

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: