Add Support for Equirectangular Images (#966)

nerfstudio/cameras/cameras.py

@@ -38,6 +38,7 @@ class CameraType(Enum):
 
     PERSPECTIVE = auto()
     FISHEYE = auto()
+    EQUIRECTANGULAR = auto()
 
 
 CAMERA_MODEL_TO_TYPE = {
@@ -47,6 +48,7 @@ CAMERA_MODEL_TO_TYPE = {
     "RADIAL": CameraType.PERSPECTIVE,
     "OPENCV": CameraType.PERSPECTIVE,
     "OPENCV_FISHEYE": CameraType.FISHEYE,
+    "EQUIRECTANGULAR": CameraType.EQUIRECTANGULAR,
 }
 
 
@@ -299,7 +301,7 @@ class Cameras(TensorDataclass):
             image_coords = torch.stack(image_coords, dim=-1) + pixel_offset  # stored as (y, x) coordinates
         return image_coords
 
-    def generate_rays(
+    def generate_rays(  # pylint: disable=too-many-statements
         self,
         camera_indices: Union[TensorType["num_rays":..., "num_cameras_batch_dims"], int],
         coords: Optional[TensorType["num_rays":..., 2]] = None,
@@ -588,8 +590,15 @@ class Cameras(TensorDataclass):
         elif distortion_params_delta is not None:
             distortion_params = distortion_params_delta
 
+        # Do not apply distortion for equirectangular images
         if distortion_params is not None:
-            coord_stack = camera_utils.radial_and_tangential_undistort(coord_stack, distortion_params)
+            mask = (self.camera_type[true_indices] != CameraType.EQUIRECTANGULAR.value).squeeze(-1)  # (num_rays)
+            coord_mask = torch.stack([mask, mask, mask], dim=0)
+            if mask.any():
+                coord_stack[coord_mask, :] = camera_utils.radial_and_tangential_undistort(
+                    coord_stack[coord_mask, :].reshape(3, -1, 2),
+                    distortion_params[mask, :],
+                ).reshape(-1, 2)
 
         # Make sure after we have undistorted our images, the shapes are still correct
         assert coord_stack.shape == (3,) + num_rays_shape + (2,)
@@ -621,8 +630,21 @@ class Cameras(TensorDataclass):
             directions_stack[..., 1][mask] = torch.masked_select(coord_stack[..., 1] * sin_theta / theta, mask).float()
             directions_stack[..., 2][mask] = -torch.masked_select(torch.cos(theta), mask)
 
+        if CameraType.EQUIRECTANGULAR.value in cam_types:
+            mask = (self.camera_type[true_indices] == CameraType.EQUIRECTANGULAR.value).squeeze(-1)  # (num_rays)
+            mask = torch.stack([mask, mask, mask], dim=0)
+
+            # For equirect, fx = fy = height = width/2
+            # Then coord[..., 0] goes from -1 to 1 and coord[..., 1] goes from -1/2 to 1/2
+            theta = -torch.pi * coord_stack[..., 0]  # minus sign for right-handed
+            phi = torch.pi * (0.5 - coord_stack[..., 1])
+            # use spherical in local camera coordinates (+y up, x=0 and z<0 is theta=0)
+            directions_stack[..., 0][mask] = torch.masked_select(-torch.sin(theta) * torch.sin(phi), mask).float()
+            directions_stack[..., 1][mask] = torch.masked_select(torch.cos(phi), mask).float()
+            directions_stack[..., 2][mask] = torch.masked_select(-torch.cos(theta) * torch.sin(phi), mask).float()
+
         for value in cam_types:
-            if value not in [CameraType.PERSPECTIVE.value, CameraType.FISHEYE.value]:
+            if value not in [CameraType.PERSPECTIVE.value, CameraType.FISHEYE.value, CameraType.EQUIRECTANGULAR.value]:
                 raise ValueError(f"Camera type {value} not supported.")
 
         assert directions_stack.shape == (3,) + num_rays_shape + (3,)

nerfstudio/data/datamanagers/base_datamanager.py

@@ -20,7 +20,7 @@ from __future__ import annotations
 
 from abc import abstractmethod
 from dataclasses import dataclass, field
-from typing import Dict, List, Optional, Tuple, Type, Union
+from typing import Any, Dict, List, Optional, Tuple, Type, Union
 
 import torch
 import tyro
@@ -32,6 +32,7 @@ from torch.utils.data.distributed import DistributedSampler
 from typing_extensions import Literal
 
 from nerfstudio.cameras.camera_optimizers import CameraOptimizerConfig
+from nerfstudio.cameras.cameras import CameraType
 from nerfstudio.cameras.rays import RayBundle
 from nerfstudio.configs.base_config import InstantiateConfig
 from nerfstudio.data.dataparsers.blender_dataparser import BlenderDataParserConfig
@@ -47,7 +48,7 @@ from nerfstudio.data.dataparsers.phototourism_dataparser import (
 )
 from nerfstudio.data.dataparsers.record3d_dataparser import Record3DDataParserConfig
 from nerfstudio.data.datasets.base_dataset import InputDataset
-from nerfstudio.data.pixel_samplers import PixelSampler
+from nerfstudio.data.pixel_samplers import EquirectangularPixelSampler, PixelSampler
 from nerfstudio.data.utils.dataloaders import (
     CacheDataloader,
     FixedIndicesEvalDataloader,
@@ -332,6 +333,19 @@ class VanillaDataManager(DataManager):  # pylint: disable=abstract-method
             scale_factor=self.config.camera_res_scale_factor,
         )
 
+    def _get_pixel_sampler(  # pylint: disable=no-self-use
+        self, dataset: InputDataset, *args: Any, **kwargs: Any
+    ) -> PixelSampler:
+        """Infer pixel sampler to use."""
+        # If all images are equirectangular, use equirectangular pixel sampler
+        is_equirectangular = dataset.cameras.camera_type == CameraType.EQUIRECTANGULAR.value
+        if is_equirectangular.all():
+            return EquirectangularPixelSampler(*args, **kwargs)
+        # Otherwise, use the default pixel sampler
+        if is_equirectangular.any():
+            CONSOLE.print("[bold yellow]Warning: Some cameras are equirectangular, but using default pixel sampler.")
+        return PixelSampler(*args, **kwargs)
+
     def setup_train(self):
         """Sets up the data loaders for training"""
         assert self.train_dataset is not None
@@ -346,7 +360,7 @@ class VanillaDataManager(DataManager):  # pylint: disable=abstract-method
             collate_fn=self.config.collate_fn,
         )
         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 = self._get_pixel_sampler(self.train_dataset, self.config.train_num_rays_per_batch)
         self.train_camera_optimizer = self.config.camera_optimizer.setup(
             num_cameras=self.train_dataset.cameras.size, device=self.device
         )
@@ -369,7 +383,7 @@ class VanillaDataManager(DataManager):  # pylint: disable=abstract-method
             collate_fn=self.config.collate_fn,
         )
         self.iter_eval_image_dataloader = iter(self.eval_image_dataloader)
-        self.eval_pixel_sampler = PixelSampler(self.config.eval_num_rays_per_batch)
+        self.eval_pixel_sampler = self._get_pixel_sampler(self.eval_dataset, self.config.eval_num_rays_per_batch)
         self.eval_ray_generator = RayGenerator(
             self.eval_dataset.cameras.to(self.device),
             self.train_camera_optimizer,  # should be shared between train and eval.

nerfstudio/data/pixel_samplers.py

@@ -171,3 +171,69 @@ class PixelSampler:  # pylint: disable=too-few-public-methods
         else:
             raise ValueError("image_batch['image'] must be a list or torch.Tensor")
         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 equirectangular images and samples pixels to use for
+    generating rays. Rays will be generated uniformly on the sphere.
+    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
+    """
+    # TODO(kevinddchen): make more DRY
+    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:
+        # TODO(kevinddchen): implement this
+        raise NotImplementedError("Masking not implemented for equirectangular images.")
+
+    # We sample theta uniformly in [0, 2*pi]
+    # We sample phi in [0, pi] according to the PDF f(phi) = sin(phi) / 2.
+    # This is done by inverse transform sampling.
+    # http://corysimon.github.io/articles/uniformdistn-on-sphere/
+    num_images_rand = torch.rand(num_rays_per_batch, device=device)
+    phi_rand = torch.acos(1 - 2 * torch.rand(num_rays_per_batch, device=device)) / torch.pi
+    theta_rand = torch.rand(num_rays_per_batch, device=device)
+    indices = torch.floor(
+        torch.stack((num_images_rand, phi_rand, theta_rand), 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))
+    collated_batch = {key: value[c, y, x] for key, value in batch.items() if key != "image_idx" and value is not None}
+
+    assert collated_batch["image"].shape == (num_rays_per_batch, 3), collated_batch["image"].shape
+
+    # Needed to correct the random indices to their actual camera idx locations.
+    indices[:, 0] = batch["image_idx"][c]
+    collated_batch["indices"] = indices  # with the abs camera indices
+
+    if keep_full_image:
+        collated_batch["full_image"] = batch["image"]
+
+    return collated_batch
+
+
+class EquirectangularPixelSampler(PixelSampler):  # pylint: disable=too-few-public-methods
+    """Samples 'pixel_batch's from 'image_batch's. Assumes images are
+    equirectangular and the sampling is done uniformly on the sphere.
+
+    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
+    """
+
+    # overrides base method
+    def sample(self, image_batch: Dict):
+
+        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/viewer/view_dataset.py

@@ -13,7 +13,7 @@ from pathlib import Path
 import torch
 import tyro
 
-from nerfstudio.configs import method_configs as cfg
+from nerfstudio.configs.base_config import ViewerConfig
 from nerfstudio.data.datamanagers import AnnotatedDataParserUnion
 from nerfstudio.data.datasets.base_dataset import InputDataset
 from nerfstudio.viewer.server import viewer_utils
@@ -26,7 +26,7 @@ torch.backends.cudnn.benchmark = True  # type: ignore
 
 def main(
     dataparser: AnnotatedDataParserUnion,
-    viewer: cfg.ViewerConfig,
+    viewer: ViewerConfig,
     log_base_dir: Path = Path("/tmp/nerfstudio_viewer_logs"),
 ) -> None:
     """Main function."""
@@ -34,7 +34,7 @@ def main(
         viewer,
         log_filename=log_base_dir / viewer.relative_log_filename,
     )
-    dataset = InputDataset(dataparser.setup().get_dataset_inputs(split="train"))
+    dataset = InputDataset(dataparser.setup().get_dataparser_outputs(split="train"))
     viewer_state.init_scene(dataset=dataset, start_train=False)
     CONSOLE.log("Please refresh and load page at: %s", viewer_state.viewer_url)
     time.sleep(30)  # allowing time to refresh page

tests/cameras/test_cameras.py

@@ -6,7 +6,7 @@ from itertools import product
 
 import torch
 
-from nerfstudio.cameras.cameras import Cameras
+from nerfstudio.cameras.cameras import Cameras, CameraType
 from nerfstudio.cameras.rays import RayBundle
 
 BATCH_SIZE = 2
@@ -121,6 +121,44 @@ def test_pinhole_camera():
     pinhole_camera.generate_rays(camera_indices=0, coords=coords)
 
 
+def test_equirectangular_camera():
+    """Test that the equirectangular camera model works."""
+    height = 100  # width is twice the height
+    c2w = torch.eye(4)[None, :3, :]
+    equirectangular_camera = Cameras(
+        cx=float(height),
+        cy=0.5 * float(height),
+        fx=float(height),
+        fy=float(height),
+        camera_to_worlds=c2w,
+        camera_type=CameraType.EQUIRECTANGULAR,
+    )
+    camera_ray_bundle = equirectangular_camera.generate_rays(camera_indices=0)
+    assert isinstance(camera_ray_bundle, RayBundle)
+    assert torch.allclose(camera_ray_bundle.origins[0], torch.tensor([0.0, 0.0, 0.0]))
+
+    # Check that the directions are mostly correct in local camera coordinates (+y is up, -z is forward)
+    directions = camera_ray_bundle.directions
+    threshold = 0.9
+    x = torch.tensor([1.0, 0.0, 0.0])
+    y = torch.tensor([0.0, 1.0, 0.0])
+    z = torch.tensor([0.0, 0.0, 1.0])
+    # top pixels point up
+    assert directions[0, 0] @ y > threshold
+    assert directions[0, height] @ y > threshold
+    assert directions[0, -1] @ y > threshold
+    # middle pixels point horizontally; middle of image is camera forwards
+    assert directions[height // 2, 0] @ z > threshold
+    assert directions[height // 2, height // 2] @ -x > threshold
+    assert directions[height // 2, height] @ -z > threshold
+    assert directions[height // 2, 3 * height // 2] @ x > threshold
+    assert directions[height // 2, -1] @ z > threshold
+    # bottom pixels point down
+    assert directions[-1, 0] @ -y > threshold
+    assert directions[-1, height] @ -y > threshold
+    assert directions[-1, -1] @ -y > threshold
+
+
 def test_camera_as_tensordataclass():
     """Test that the camera class move to Tensordataclass works."""
     _ = C2_DIST[torch.tensor([0]), torch.tensor([0])]
@@ -265,4 +303,5 @@ def _check_cam_shapes(cam: Cameras, _batch_size):
 
 if __name__ == "__main__":
     test_pinhole_camera()
+    test_equirectangular_camera()
     test_camera_as_tensordataclass()