Cameras TensorDataclass and Phototourism (#961)

This branch gives us the flexibility to use Cameras as a tensordataclass, giving arbitrary batch dimensions, slicing, etc all for free. It tests this with the phototourism dataset.

Tensordataclass changes:
Originally, the tensordataclass would look at tensors in an object and take everything in front of the last dimension to be the batch dimensions, but in the case of the camera_to_world matrices, everything except for the last TWO are the important dimensions, so we needed to have a way for the tensordataclass to acknowledge these specific fields where there are multiple data dimensions behind the batch dimensions. This is important for things like reshaping where we need to know which dimensions get reshaped and which dimensions get left alone.

Cameras changes:
Now subclass of tensordataclass, which required a change to the tensordataclass itself.
The possibility of extra batch dimensions is the hairy part that requires the majority of these changes. The extra batch dimensions aren't a hassle in every function except in ray generation. The ray generation is where this operates different relative to the old cameras. The reason for this is before, for each ray we wanted to generate, we could specify a single index into our batch of cameras, but if you have multiple camera dimensions, for each ray, you need to now specify one index per camera batch dimension, meaning our camera_indices we specify need to have one extra dimension at the end, a dimension of size Cameras.ndim (number of batch dimensions in the camera object).
I also added some neat logic in generate_rays where if you want your output raybundle to be batched, if you specify batched camera_indices, you will receive the same batch size for the output raybundle, meaning you can easily generate raybundles corresponding to things like image patches and whatnot.

Phototourism changes:
The difficulty here is that images in the phototourism dataset are of different sizes, which is difficult because we can no longer batch the image tensors. We also need to better utilize caching in the CacheDataloader since the vanilladatamanager defaults don't really use it as well as they could, and without properly utilizing it, the rays / sec are maybe 25% of the standard nerfstudio dataset speeds (which are 10% of the size of the phototourism datasets). To do this, we specify a collate function in the existing datamanager configs, and we make a subclass of the vanilladatamanagerconfig named VariableResDataManagerConfig which will have a custom collate function storing the images in a list, and using better defaults for caching in the dataloader, caching ~40 images at a time. The slowest part of using a list is pixel sampling, where we added a custom if case for if our pixel sampler receives a batch of images that are in a list, and the lack of ability to vectorize operations over a list means that caching in the dataloader is even more important.

nerfstudio/cameras/camera_paths.py

@@ -69,6 +69,7 @@ def get_spiral_path(
     """
 
     assert radius is not None or radiuses is not None, "Either radius or radiuses must be specified."
+    assert camera.ndim == 1, "We assume only one batch dim here"
     if radius is not None and radiuses is None:
         rad = torch.tensor([radius] * 3, device=camera.device)
     elif radiuses is not None and radius is None:
@@ -99,7 +100,13 @@ def get_spiral_path(
         new_c2ws.append(c2wh[:3, :4])
     new_c2ws = torch.stack(new_c2ws, dim=0)
 
-    return Cameras(fx=camera.fx[0], fy=camera.fy[0], cx=camera.cx[0], cy=camera.cy[0], camera_to_worlds=new_c2ws)
+    return Cameras(
+        fx=camera.fx[0],
+        fy=camera.fy[0],
+        cx=camera.cx[0],
+        cy=camera.cy[0],
+        camera_to_worlds=new_c2ws,
+    )
 
 
 def get_path_from_json(camera_path: Dict[str, Any]) -> Cameras:

nerfstudio/cameras/cameras.py

@@ -17,8 +17,9 @@ Camera Models
 """
 import base64
 import math
+from dataclasses import dataclass
 from enum import Enum, auto
-from typing import Dict, List, Optional, Union
+from typing import Dict, List, Optional, Tuple, Union
 
 import cv2
 import torch
@@ -29,6 +30,7 @@ from torchtyping import TensorType
 import nerfstudio.utils.poses as pose_utils
 from nerfstudio.cameras import camera_utils
 from nerfstudio.cameras.rays import RayBundle
+from nerfstudio.utils.tensor_dataclass import TensorDataclass
 
 
 class CameraType(Enum):
@@ -48,85 +50,123 @@ CAMERA_MODEL_TO_TYPE = {
 }
 
 
-class Cameras:
+@dataclass(init=False)
+class Cameras(TensorDataclass):
     """Dataparser outputs for the image dataset and the ray generator.
 
     Note: currently only supports cameras with the same principal points and types. The reason we type
     the focal lengths, principal points, and image sizes as tensors is to allow for batched cameras
     down the line in cases where your batches of camera data don't come from the same cameras.
 
+     If a single value is provided, it is broadcasted to all cameras.
+
     Args:
-        camera_to_worlds: Tensor of per-image c2w matrices, in [R | t] format.
-        fx: Focal length x. If a single value is provided, it is broadcasted to all cameras.
-        fy: Focal length y. If a single value is provided, it is broadcasted to all cameras.
-        cx: Principal point x. If a single value is provided, it is broadcasted to all cameras.
-        cy: Principal point y. If a single value is provided, it is broadcasted to all cameras.
-        width: Image width. If a single value is provided, it is broadcasted to all cameras.
-        height: Image height. If a single value is provided, it is broadcasted to all cameras.
-        distortion_params: OpenCV 6 radial distortion coefficients.
-        camera_type: Type of camera model. If a single value is provided, it is broadcasted to
-            all cameras. This will be an int corresponding to the CameraType enum.
+        camera_to_worlds: Camera to world matrices. Tensor of per-image c2w matrices, in [R | t] format
+        fx: Focal length x
+        fy: Focal length y
+        cx: Principal point x
+        cy: Principal point y
+        width: Image width
+        height: Image height
+        distortion_params: OpenCV 6 radial distortion coefficients
+        camera_type: Type of camera model. This will be an int corresponding to the CameraType enum.
+        times: Timestamps for each camera
     """
 
+    camera_to_worlds: TensorType["num_cameras":..., 3, 4]
+    fx: TensorType["num_cameras":..., 1]
+    fy: TensorType["num_cameras":..., 1]
+    cx: TensorType["num_cameras":..., 1]
+    cy: TensorType["num_cameras":..., 1]
+    width: TensorType["num_cameras":..., 1]
+    height: TensorType["num_cameras":..., 1]
+    distortion_params: Optional[TensorType["num_cameras":..., 6]]
+    camera_type: TensorType["num_cameras":..., 1]
+    times: Optional[TensorType["num_cameras", 1]]
+
     def __init__(
         self,
-        camera_to_worlds: TensorType["num_cameras", 3, 4],
-        fx: Union[TensorType["num_cameras"], float],
-        fy: Union[TensorType["num_cameras"], float],
-        cx: Union[TensorType["num_cameras"], float],
-        cy: Union[TensorType["num_cameras"], float],
-        width: Optional[Union[TensorType["num_cameras"], int]] = None,
-        height: Optional[Union[TensorType["num_cameras"], int]] = None,
-        distortion_params: Optional[TensorType["num_cameras", 6]] = None,
+        camera_to_worlds: TensorType["batch_c2ws":..., 3, 4],
+        fx: Union[TensorType["batch_fxs":..., 1], float],
+        fy: Union[TensorType["batch_fys":..., 1], float],
+        cx: Union[TensorType["batch_cxs":..., 1], float],
+        cy: Union[TensorType["batch_cys":..., 1], float],
+        width: Optional[Union[TensorType["batch_ws":..., 1], int]] = None,
+        height: Optional[Union[TensorType["batch_hs":..., 1], int]] = None,
+        distortion_params: Optional[TensorType["batch_dist_params":..., 6]] = None,
         camera_type: Optional[
-            Union[TensorType["num_cameras"], int, List[CameraType], CameraType]
+            Union[
+                TensorType["batch_cam_types":..., 1],
+                int,
+                List[CameraType],
+                CameraType,
+            ]
         ] = CameraType.PERSPECTIVE,
         times: Optional[TensorType["num_cameras"]] = None,
     ):
-        self._num_cameras = camera_to_worlds.shape[0]
-        self.camera_to_worlds = camera_to_worlds  # This comes first since it determines @property self.device
+        """Initializes the Cameras object.
+
+        Note on Input Tensor Dimensions: All of these tensors have items of dimensions TensorType[3, 4]
+        (in the case of the c2w matrices), TensorType[6] (in the case of distortion params), or
+        TensorType[1] (in the case of the rest of the elements). The dimensions before that are
+        considered the batch dimension of that tensor (batch_c2ws, batch_fxs, etc.). We will broadcast
+        all the tensors to be the same batch dimension. This means you can use any combination of the
+        input types in the function signature and it won't break. Your batch size for all tensors
+        must be broadcastable to the same size, and the resulting number of batch dimensions will be
+        the batch dimension with the largest number of dimensions.
+        """
+
+        # This will notify the tensordataclass that we have a field with more than 1 dimension
+        self._field_custom_dimensions = {"camera_to_worlds": 2}
+
+        self.camera_to_worlds = camera_to_worlds
 
         # fx fy calculation
-        if not isinstance(fx, torch.Tensor):
-            fx = torch.Tensor([fx])
-        if not isinstance(fy, torch.Tensor):
-            fy = torch.Tensor([fy])
-        self.fx = fx.to(self.device).broadcast_to((self._num_cameras))
-        self.fy = fy.to(self.device).broadcast_to((self._num_cameras))
+        self.fx = self._init_get_fc_xy(fx, "fx")  # @dataclass's post_init will take care of broadcasting
+        self.fy = self._init_get_fc_xy(fy, "fy")  # @dataclass's post_init will take care of broadcasting
 
         # cx cy calculation
-        if not isinstance(cx, torch.Tensor):
-            cx = torch.Tensor([cx])
-        else:
-            assert torch.all(
-                cx == (cx[0] if cx.ndim > 0 else cx.item())
-            ), "Batched cameras of different types will be allowed in the future."
-        if not isinstance(cy, torch.Tensor):
-            cy = torch.Tensor([cy])
-        else:
-            assert torch.all(
-                cy == (cy[0] if cy.ndim > 0 else cy.item())
-            ), "Batched cameras of different types will be allowed in the future."
-        self.cx = cx.to(self.device).broadcast_to((self._num_cameras))
-        self.cy = cy.to(self.device).broadcast_to((self._num_cameras))
+        self.cx = self._init_get_fc_xy(cx, "cx")  # @dataclass's post_init will take care of broadcasting
+        self.cy = self._init_get_fc_xy(cy, "cy")  # @dataclass's post_init will take care of broadcasting
 
         # Distortion Params Calculation:
-        if distortion_params is not None:
-            self.distortion_params = distortion_params.broadcast_to((self._num_cameras, 6))
-        else:
-            self.distortion_params = None
-
-        self._image_heights = self._init_get_height_width(height, cy)
-
-        self._image_widths = self._init_get_height_width(width, cx)
+        self.distortion_params = distortion_params  # @dataclass's post_init will take care of broadcasting
 
+        # @dataclass's post_init will take care of broadcasting
+        self.height = self._init_get_height_width(height, self.cy)
+        self.width = self._init_get_height_width(width, self.cx)
         self.camera_type = self._init_get_camera_type(camera_type)
+        self.times = self._init_get_times(times)
 
-        self.times = times
+        self.__post_init__()  # This will do the dataclass post_init and broadcast all the tensors
+
+    def _init_get_fc_xy(self, fc_xy, name):
+        """
+        Parses the input focal length / principle point x or y and returns a tensor of the correct shape
+
+        Only needs to make sure that we a 1 in the last dimension if it is a tensor. If it is a float, we
+        just need to make it into a tensor and it will be broadcasted later in the __post_init__ function.
+
+        Args:
+            fc_xy: The focal length / principle point x or y
+            name: The name of the variable. Used for error messages
+        """
+        if isinstance(fc_xy, float):
+            fc_xy = torch.Tensor([fc_xy], device=self.device)
+        elif isinstance(fc_xy, torch.Tensor):
+            if fc_xy.ndim == 0 or fc_xy.shape[-1] != 1:
+                fc_xy = fc_xy.unsqueeze(-1)
+            fc_xy = fc_xy.to(self.device)
+        else:
+            raise ValueError(f"{name} must be a float or tensor, got {type(fc_xy)}")
+        return fc_xy
 
     def _init_get_camera_type(
-        self, camera_type: Union[TensorType["num_cameras"], int, List[CameraType], CameraType]
-    ) -> TensorType["num_cameras"]:
+        self,
+        camera_type: Union[
+            TensorType["batch_cam_types":..., 1], TensorType["batch_cam_types":...], int, List[CameraType], CameraType
+        ],
+    ) -> TensorType["num_cameras":..., 1]:
         """
         Parses the __init__() argument camera_type
 
@@ -140,17 +180,21 @@ class Cameras:
             camera_type: camera_type argument from __init__()
         """
         if isinstance(camera_type, CameraType):
-            camera_type = torch.tensor(camera_type.value, device=self.device).broadcast_to((self._num_cameras))
-        elif isinstance(camera_type, List):
-            camera_type = torch.tensor([c.value for c in camera_type], device=self.device).broadcast_to(
-                (self._num_cameras)
-            )
+            camera_type = torch.tensor([camera_type.value], device=self.device)
+        elif isinstance(camera_type, List) and isinstance(camera_type[0], CameraType):
+            camera_type = torch.tensor([[c.value] for c in camera_type], device=self.device)
         elif isinstance(camera_type, int):
-            camera_type = torch.tensor(camera_type, device=self.device).broadcast_to((self._num_cameras))
+            camera_type = torch.tensor([camera_type], device=self.device)
         elif isinstance(camera_type, torch.Tensor):
-            for cam_type in camera_type:
-                assert camera_type[0] == cam_type, "Batched cameras of different types will be allowed in the future."
-            camera_type = camera_type.to(self.device).broadcast_to((self._num_cameras))
+            assert not torch.is_floating_point(
+                camera_type
+            ), f"camera_type tensor must be of type int, not: {camera_type.dtype}"
+            camera_type = camera_type.to(self.device)
+            if camera_type.ndim == 0 or camera_type.shape[-1] != 1:
+                camera_type = camera_type.unsqueeze(-1)
+            # assert torch.all(
+            #     camera_type.view(-1)[0] == camera_type
+            # ), "Batched cameras of different camera_types will be allowed in the future."
         else:
             raise ValueError(
                 'Invalid camera_type. Must be CameraType, List[CameraType], int, or torch.Tensor["num_cameras"]. \
@@ -160,8 +204,10 @@ class Cameras:
         return camera_type
 
     def _init_get_height_width(
-        self, h_w: Union[TensorType["num_cameras"], int, None], c_x_y: TensorType["num_cameras"]
-    ) -> TensorType["num_cameras"]:
+        self,
+        h_w: Union[TensorType["batch_hws":..., 1], TensorType["batch_hws":...], int, None],
+        c_x_y: TensorType["batch_cxys":...],
+    ) -> TensorType["num_cameras":..., 1]:
         """
         Parses the __init__() argument for height or width
 
@@ -177,120 +223,366 @@ class Cameras:
         """
         if isinstance(h_w, int):
             h_w = torch.Tensor([h_w]).to(torch.int64).to(self.device)
-            h_w = h_w.broadcast_to((self._num_cameras))
         elif isinstance(h_w, torch.Tensor):
-            h_w = h_w.to(torch.int64).to(self.device).broadcast_to((self._num_cameras))
-            assert torch.all(h_w == h_w[0]), "Batched cameras of different types will be allowed in the future."
+            assert not torch.is_floating_point(h_w), f"height and width tensor must be of type int, not: {h_w.dtype}"
+            h_w = h_w.to(torch.int64).to(self.device)
+            if h_w.ndim == 0 or h_w.shape[-1] != 1:
+                h_w = h_w.unsqueeze(-1)
+        # assert torch.all(h_w == h_w.view(-1)[0]), "Batched cameras of different h, w will be allowed in the future."
         elif h_w is None:
-            h_w = torch.Tensor((c_x_y * 2).to(torch.int64).to(self.device)).broadcast_to((self._num_cameras))
+            h_w = torch.Tensor((c_x_y * 2).to(torch.int64).to(self.device))
         else:
             raise ValueError("Height must be an int, tensor, or None, received: " + str(type(h_w)))
         return h_w
 
+    def _init_get_times(self, times):
+        if times is None:
+            times = None
+        elif isinstance(times, torch.Tensor):
+            if times.ndim == 0 or times.shape[-1] != 1:
+                times = times.unsqueeze(-1).to(self.device)
+        else:
+            raise ValueError(f"times must be a tensor, got {type(times)}")
+
+        return times
+
     @property
     def device(self):
         """Returns the device that the camera is on."""
         return self.camera_to_worlds.device
 
     @property
-    def size(self) -> int:
-        """Returns the number of cameras."""
-        return self._num_cameras
+    def image_height(self) -> TensorType["num_cameras":..., 1]:
+        """Returns the height of the images."""
+        return self.height
 
     @property
-    def image_height(self) -> TensorType["num_cameras"]:
+    def image_width(self) -> TensorType["num_cameras":..., 1]:
         """Returns the height of the images."""
-        return self._image_heights
+        return self.width
 
     @property
-    def image_width(self) -> TensorType["num_cameras"]:
-        """Returns the height of the images."""
-        return self._image_widths
-
-    def to(self, device: Union[torch.device, str]) -> "Cameras":
+    def is_jagged(self):
         """
-        Args:
-            device: Device to move the camera onto
-
-        Returns:
-            Cameras on the specified device.
+        Returns whether or not the cameras are "jagged" (i.e. the height and widths are different, meaning that
+        you cannot concatenate the image coordinate maps together)
         """
-        distortion_params = self.distortion_params.to(device) if self.distortion_params is not None else None
-        return Cameras(
-            camera_to_worlds=self.camera_to_worlds.to(device),
-            fx=self.fx.to(device),
-            fy=self.fy.to(device),
-            cx=self.cx.to(device),
-            cy=self.cy.to(device),
-            width=self.image_width.to(device),
-            height=self.image_height.to(device),
-            distortion_params=distortion_params.to(device) if distortion_params is not None else None,
-            camera_type=self.camera_type.to(device),
-            times=self.times.to(device) if self.times is not None else None,
-        )
+        h_jagged = not torch.all(self.height == self.height.view(-1)[0])
+        w_jagged = not torch.all(self.width == self.width.view(-1)[0])
+        return h_jagged or w_jagged
 
-    def get_image_coords(self, pixel_offset: float = 0.5) -> TensorType["height", "width", 2]:
-        """This gets the image coordinates of one of the cameras in this object
+    def get_image_coords(
+        self, pixel_offset: float = 0.5, index: Optional[Tuple] = None
+    ) -> TensorType["height", "width", 2]:
+        """This gets the image coordinates of one of the cameras in this object.
 
-        Down the line we may support jagged images, allowing this to return multiple image coordinates of
-        different sizes, but for the time being since all cameras are constrained to be the same height and
-        width, this will return the same image coordinates for all cameras.
+        If no index is specified, it will return the maximum possible sized height / width image coordinate map,
+        by looking at the maximum height and width of all the cameras in this object.
 
         Args:
             pixel_offset: Offset for each pixel. Defaults to center of pixel (0.5)
+            index: Tuple of indices into the batch dimensions of the camera. Defaults to None, which returns the 0th
+                flattened camera
 
         Returns:
             Grid of image coordinates.
         """
-        image_height = self.image_height[0]
-        image_width = self.image_width[0]
-        image_coords = torch.meshgrid(torch.arange(image_height), torch.arange(image_width), indexing="ij")
-        image_coords = torch.stack(image_coords, dim=-1) + pixel_offset  # stored as (y, x) coordinates
+        if index is None:
+            image_height = torch.max(self.image_height.view(-1))
+            image_width = torch.max(self.image_width.view(-1))
+            image_coords = torch.meshgrid(torch.arange(image_height), torch.arange(image_width), indexing="ij")
+            image_coords = torch.stack(image_coords, dim=-1) + pixel_offset  # stored as (y, x) coordinates
+        else:
+            image_height = self.image_height[index].item()
+            image_width = self.image_width[index].item()
+            image_coords = torch.meshgrid(torch.arange(image_height), torch.arange(image_width), indexing="ij")
+            image_coords = torch.stack(image_coords, dim=-1) + pixel_offset  # stored as (y, x) coordinates
         return image_coords
 
     def generate_rays(
         self,
-        camera_indices: Union[TensorType["num_rays":...], int],
+        camera_indices: Union[TensorType["num_rays":..., "num_cameras_batch_dims"], 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,
+        keep_shape: Optional[bool] = None,
     ) -> RayBundle:
         """Generates rays for the given camera indices.
 
-        TODO - add support for intrinsics delta.
-        TODO - add support for distortions.
+        This function will standardize the input arguments and then call the _generate_rays_from_coords function
+        to generate the rays. Our goal is to parse the arguments and then get them into the right shape:
+            - camera_indices: (num_rays:..., num_cameras_batch_dims)
+            - coords: (num_rays:..., 2)
+            - camera_opt_to_camera: (num_rays:..., 3, 4) or None
+            - distortion_params_delta: (num_rays:..., 6) or None
+
+        Read the docstring for _generate_rays_from_coords for more information on how we generate the rays
+        after we have standardized the arguments.
+
+        We are only concerned about different combinations of camera_indices and coords matrices, and the following
+        are the 4 cases we have to deal with:
+            1. isinstance(camera_indices, int) and coords == None
+                - In this case we broadcast our camera_indices / coords shape (h, w, 1 / 2 respectively)
+            2. isinstance(camera_indices, int) and coords != None
+                - In this case, we broadcast camera_indices to the same batch dim as coords
+            3. not isinstance(camera_indices, int) and coords == None
+                - In this case, we will need to set coords so that it is of shape (h, w, num_rays, 2), and broadcast
+                    all our other args to match the new definition of num_rays := (h, w) + num_rays
+            4. not isinstance(camera_indices, int) and coords != None
+                - In this case, we have nothing to do, only check that the arguments are of the correct shape
+
+        There is one more edge case we need to be careful with: when we have "jagged cameras" (ie: different heights
+        and widths for each camera). This isn't problematic when we specify coords, since coords is already a tensor.
+        When coords == None (ie: when we render out the whole image associated with this camera), we run into problems
+        since there's no way to stack each coordinate map as all coordinate maps are all different shapes. In this case,
+        we will need to flatten each individual coordinate map and concatenate them, giving us only one batch dimension,
+        regaurdless of the number of prepended extra batch dimensions in the camera_indices tensor.
+
 
         Args:
-            camera_indices: Indices of the cameras to generate rays for.
+            camera_indices: Camera indices of the flattened cameras object to generate rays for.
             coords: Coordinates of the pixels to generate rays for. If None, the full image will be rendered.
             camera_opt_to_camera: Optional transform for the camera to world matrices.
             distortion_params_delta: Optional delta for the distortion parameters.
+            keep_shape: If None, then we default to the regular behavior of flattening if cameras is jagged, otherwise
+                keeping dimensions. If False, we flatten at the end. If True, then we keep the shape of the
+                camera_indices and coords tensors (if we can).
 
         Returns:
             Rays for the given camera indices and coords.
         """
+        # Check the argument types to make sure they're valid and all shaped correctly
+        assert isinstance(camera_indices, (torch.Tensor, int)), "camera_indices must be a tensor or int"
+        assert coords is None or isinstance(coords, torch.Tensor), "coords must be a tensor or None"
+        assert camera_opt_to_camera is None or isinstance(camera_opt_to_camera, torch.Tensor)
+        assert distortion_params_delta is None or isinstance(distortion_params_delta, torch.Tensor)
+        if isinstance(camera_indices, torch.Tensor) and isinstance(coords, torch.Tensor):
+            num_rays_shape = camera_indices.shape[:-1]
+            errormsg = "Batch dims of inputs must match when inputs are all tensors"
+            assert coords.shape[:-1] == num_rays_shape, errormsg
+            assert camera_opt_to_camera is None or camera_opt_to_camera.shape[:-2] == num_rays_shape, errormsg
+            assert distortion_params_delta is None or distortion_params_delta.shape[:-1] == num_rays_shape, errormsg
 
-        if isinstance(camera_indices, torch.Tensor):
-            camera_indices = camera_indices.to(self.device)
+        # If zero dimensional, we need to unsqueeze to get a batch dimension and then squeeze later
+        if not self.shape:
+            cameras = self.reshape((1,))
+            assert torch.all(
+                torch.tensor(camera_indices == 0) if isinstance(camera_indices, int) else camera_indices == 0
+            ), "Can only index into single camera with no batch dimensions if index is zero"
+        else:
+            cameras = self
 
+        # If the camera indices are an int, then we need to make sure that the camera batch is 1D
+        if isinstance(camera_indices, int):
+            assert (
+                len(cameras.shape) == 1
+            ), "camera_indices must be a tensor if cameras are batched with more than 1 batch dimension"
+            camera_indices = torch.tensor([camera_indices], device=cameras.device)
+
+        assert camera_indices.shape[-1] == len(
+            cameras.shape
+        ), "camera_indices must have shape (num_rays:..., num_cameras_batch_dims)"
+
+        # If keep_shape is True, then we need to make sure that the camera indices in question
+        # are all the same height and width and can actually be batched while maintaining the image
+        # shape
+        if keep_shape is True:
+            assert torch.all(cameras.height[camera_indices] == cameras.height[camera_indices[0]]) and torch.all(
+                cameras.width[camera_indices] == cameras.width[camera_indices[0]]
+            ), "Can only keep shape if all cameras have the same height and width"
+
+        # If the cameras don't all have same height / width, if coords is not none, we will need to generate
+        # a flat list of coords for each camera and then concatenate otherwise our rays will be jagged.
+        # Camera indices, camera_opt, and distortion will also need to be broadcasted accordingly which is non-trivial
+        if cameras.is_jagged and coords is None and (keep_shape is None or keep_shape is False):
+            index_dim = camera_indices.shape[-1]
+            camera_indices = camera_indices.reshape(-1, index_dim)
+            _coords = [cameras.get_image_coords(index=tuple(index)).reshape(-1, 2) for index in camera_indices]
+            camera_indices = torch.cat(
+                [index.unsqueeze(0).repeat(coords.shape[0], 1) for index, coords in zip(camera_indices, _coords)],
+            )
+            coords = torch.cat(_coords, dim=0)
+            assert coords.shape[0] == camera_indices.shape[0]
+            # Need to get the coords of each indexed camera and flatten all coordinate maps and concatenate them
+
+        # The case where we aren't jagged && keep_shape (since otherwise coords is already set) and coords
+        # is None. In this case we append (h, w) to the num_rays dimensions for all tensors. In this case,
+        # each image in camera_indices has to have the same shape since otherwise we would have error'd when
+        # we checked keep_shape is valid or we aren't jagged.
         if coords is None:
-            coords = self.get_image_coords().to(self.device)
+            index_dim = camera_indices.shape[-1]
+            index = camera_indices.reshape(-1, index_dim)[0]
+            coords: torch.Tensor = cameras.get_image_coords(index=tuple(index))  # (h, w, 2)
+            coords = coords.reshape(coords.shape[:2] + (1,) * len(camera_indices.shape[:-1]) + (2,))  # (h, w, 1..., 2)
+            coords = coords.expand(coords.shape[:2] + camera_indices.shape[:-1] + (2,))  # (h, w, num_rays, 2)
+            camera_opt_to_camera = (  # (h, w, num_rays, 3, 4) or None
+                camera_opt_to_camera.broadcast_to(coords.shape[:-1] + (3, 4))
+                if camera_opt_to_camera is not None
+                else None
+            )
+            distortion_params_delta = (  # (h, w, num_rays, 6) or None
+                distortion_params_delta.broadcast_to(coords.shape[:-1] + (6,))
+                if distortion_params_delta is not None
+                else None
+            )
 
-        assert coords is not None
-        y = coords[..., 0]  # (..., 1)
-        x = coords[..., 1]  # (..., 1)
-        fx, fy = self.fx[camera_indices], self.fy[camera_indices]
-        cx, cy = self.cx[camera_indices], self.cy[camera_indices]
+        # If camera indices was an int or coords was none, we need to broadcast our indices along batch dims
+        camera_indices = camera_indices.broadcast_to(coords.shape[:-1] + (len(cameras.shape),)).to(torch.long)
 
-        coord = torch.stack([(x - cx) / fx, -(y - cy) / fy], -1)
-        coord_x_offset = torch.stack([(x - cx + 1) / fx, -(y - cy) / fy], -1)
-        coord_y_offset = torch.stack([(x - cx) / fx, -(y - cy + 1) / fy], -1)
+        # Checking our tensors have been standardized
+        assert isinstance(coords, torch.Tensor) and isinstance(camera_indices, torch.Tensor)
+        assert camera_indices.shape[-1] == len(cameras.shape)
+        assert camera_opt_to_camera is None or camera_opt_to_camera.shape[:-2] == coords.shape[:-1]
+        assert distortion_params_delta is None or distortion_params_delta.shape[:-1] == coords.shape[:-1]
 
-        coord_stack = torch.stack([coord, coord_x_offset, coord_y_offset], dim=0)
+        # This will do the actual work of generating the rays now that we have standardized the inputs
+        # raybundle.shape == (num_rays) when done
+        # pylint: disable=protected-access
+        raybundle = cameras._generate_rays_from_coords(
+            camera_indices, coords, camera_opt_to_camera, distortion_params_delta
+        )
 
+        # If we have mandated that we don't keep the shape, then we flatten
+        if keep_shape is False:
+            raybundle = raybundle.flatten()
+
+        # TODO: We should have to squeeze the last dimension here if we started with zero batch dims, but never have to,
+        # so there might be a rogue squeeze happening somewhere, and this may cause some unintended behaviour
+        # that we haven't caught yet with tests
+        return raybundle
+
+    # pylint: disable=too-many-statements
+    def _generate_rays_from_coords(
+        self,
+        camera_indices: TensorType["num_rays":..., "num_cameras_batch_dims"],
+        coords: TensorType["num_rays":..., 2],
+        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 and coords where self isn't jagged
+
+        This is a fairly complex function, so let's break this down slowly.
+
+        Shapes involved:
+            - num_rays: This is your output raybundle shape. It dictates the number and shape of the rays generated
+            - num_cameras_batch_dims: This is the number of dimensions of our camera
+
+        Args:
+            camera_indices: Camera indices of the flattened cameras object to generate rays for.
+                The shape of this is such that indexing into camera_indices["num_rays":...] will return the
+                index into each batch dimension of the camera in order to get the correct camera specified by
+                "num_rays".
+
+                Example:
+                    >>> cameras = Cameras(...)
+                    >>> cameras.shape
+                        (2, 3, 4)
+
+                    >>> camera_indices = torch.tensor([0, 0, 0]) # We need an axis of length 3 since cameras.ndim == 3
+                    >>> camera_indices.shape
+                        (3,)
+                    >>> coords = torch.tensor([1,1])
+                    >>> coords.shape
+                        (2,)
+                    >>> out_rays = cameras.generate_rays(camera_indices=camera_indices, coords = coords)
+                        # This will generate a RayBundle with a single ray for the
+                        # camera at cameras[0,0,0] at image coordinates (1,1), so out_rays.shape == ()
+                    >>> out_rays.shape
+                        ()
+
+                    >>> camera_indices = torch.tensor([[0,0,0]])
+                    >>> camera_indices.shape
+                        (1, 3)
+                    >>> coords = torch.tensor([[1,1]])
+                    >>> coords.shape
+                        (1, 2)
+                    >>> out_rays = cameras.generate_rays(camera_indices=camera_indices, coords = coords)
+                        # This will generate a RayBundle with a single ray for the
+                        # camera at cameras[0,0,0] at point (1,1), so out_rays.shape == (1,)
+                        # since we added an extra dimension in front of camera_indices
+                    >>> out_rays.shape
+                        (1,)
+
+                If you want more examples, check tests/cameras/test_cameras and the function check_generate_rays_shape
+
+                The bottom line is that for camera_indices: (num_rays:..., num_cameras_batch_dims), num_rays is the
+                output shape and if you index into the output RayBundle with some indices [i:...], if you index into
+                camera_indices with camera_indices[i:...] as well, you will get a 1D tensor containing the batch
+                indices into the original cameras object corresponding to that ray (ie: you will get the camera
+                from our batched cameras corresponding to the ray at RayBundle[i:...]).
+
+            coords: Coordinates of the pixels to generate rays for. If None, the full image will be rendered, meaning
+                height and width get prepended to the num_rays dimensions. Indexing into coords with [i:...] will
+                get you the image coordinates [x, y] of that specific ray located at output RayBundle[i:...].
+
+            camera_opt_to_camera: Optional transform for the camera to world matrices.
+                In terms of shape, it follows the same rules as coords, but indexing into it with [i:...] gets you
+                the 2D camera to world transform matrix for the camera optimization at RayBundle[i:...].
+
+            distortion_params_delta: Optional delta for the distortion parameters.
+                In terms of shape, it follows the same rules as coords, but indexing into it with [i:...] gets you
+                the 1D tensor with the 6 distortion parameters for the camera optimization at RayBundle[i:...].
+
+        Returns:
+            Rays for the given camera indices and coords. RayBundle.shape == num_rays
+        """
+        # Make sure we're on the right devices
+        camera_indices = camera_indices.to(self.device)
+        coords = coords.to(self.device)
+
+        # Checking to make sure everything is of the right shape and type
+        num_rays_shape = camera_indices.shape[:-1]
+        assert camera_indices.shape == num_rays_shape + (self.ndim,)
+        assert coords.shape == num_rays_shape + (2,)
+        assert coords.shape[-1] == 2
+        assert camera_opt_to_camera is None or camera_opt_to_camera.shape == num_rays_shape + (3, 4)
+        assert distortion_params_delta is None or distortion_params_delta.shape == num_rays_shape + (6,)
+
+        # Here, we've broken our indices down along the num_cameras_batch_dims dimension allowing us to index by all
+        # of our output rays at each dimension of our cameras object
+        true_indices = [camera_indices[..., i] for i in range(camera_indices.shape[-1])]
+
+        # Get all our focal lengths, principal points and make sure they are the right shapes
+        y = coords[..., 0]  # (num_rays,) get rid of the last dimension
+        x = coords[..., 1]  # (num_rays,) get rid of the last dimension
+        fx, fy = self.fx[true_indices].squeeze(-1), self.fy[true_indices].squeeze(-1)  # (num_rays,)
+        cx, cy = self.cx[true_indices].squeeze(-1), self.cy[true_indices].squeeze(-1)  # (num_rays,)
+        assert (
+            y.shape == num_rays_shape
+            and x.shape == num_rays_shape
+            and fx.shape == num_rays_shape
+            and fy.shape == num_rays_shape
+            and cx.shape == num_rays_shape
+            and cy.shape == num_rays_shape
+        ), (
+            str(num_rays_shape)
+            + str(y.shape)
+            + str(x.shape)
+            + str(fx.shape)
+            + str(fy.shape)
+            + str(cx.shape)
+            + str(cy.shape)
+        )
+
+        # Get our image coordinates and image coordinates offset by 1 (offsets used for dx, dy calculations)
+        # Also make sure the shapes are correct
+        coord = torch.stack([(x - cx) / fx, -(y - cy) / fy], -1)  # (num_rays, 2)
+        coord_x_offset = torch.stack([(x - cx + 1) / fx, -(y - cy) / fy], -1)  # (num_rays, 2)
+        coord_y_offset = torch.stack([(x - cx) / fx, -(y - cy + 1) / fy], -1)  # (num_rays, 2)
+        assert (
+            coord.shape == num_rays_shape + (2,)
+            and coord_x_offset.shape == num_rays_shape + (2,)
+            and coord_y_offset.shape == num_rays_shape + (2,)
+        )
+
+        # Stack image coordinates and image coordinates offset by 1, check shapes too
+        coord_stack = torch.stack([coord, coord_x_offset, coord_y_offset], dim=0)  # (3, num_rays, 2)
+        assert coord_stack.shape == (3,) + num_rays_shape + (2,)
+
+        # Undistorts our images according to our distortion parameters
         distortion_params = None
         if self.distortion_params is not None:
-            distortion_params = self.distortion_params[camera_indices]
+            distortion_params = self.distortion_params[true_indices]
             if distortion_params_delta is not None:
                 distortion_params = distortion_params + distortion_params_delta
         elif distortion_params_delta is not None:
@@ -299,43 +591,69 @@ class Cameras:
         if distortion_params is not None:
             coord_stack = camera_utils.radial_and_tangential_undistort(coord_stack, distortion_params)
 
-        if self.camera_type[0] == CameraType.PERSPECTIVE.value:
-            directions_stack = torch.stack(
-                [coord_stack[..., 0], coord_stack[..., 1], -torch.ones_like(coord_stack[..., 1])], dim=-1
-            )
-        elif self.camera_type[0] == CameraType.FISHEYE.value:
+        # Make sure after we have undistorted our images, the shapes are still correct
+        assert coord_stack.shape == (3,) + num_rays_shape + (2,)
+
+        # Gets our directions for all our rays in camera coordinates and checks shapes at the end
+        # Here, directions_stack is of shape (3, num_rays, 3)
+        # directions_stack[0] is the direction for ray in camera coordinates
+        # directions_stack[1] is the direction for ray in camera coordinates offset by 1 in x
+        # directions_stack[2] is the direction for ray in camera coordinates offset by 1 in y
+        cam_types = torch.unique(self.camera_type, sorted=False)
+        directions_stack = torch.empty((3,) + num_rays_shape + (3,), device=self.device)
+        if CameraType.PERSPECTIVE.value in cam_types:
+            mask = (self.camera_type[true_indices] == CameraType.PERSPECTIVE.value).squeeze(-1)  # (num_rays)
+            mask = torch.stack([mask, mask, mask], dim=0)
+            directions_stack[..., 0][mask] = torch.masked_select(coord_stack[..., 0], mask).float()
+            directions_stack[..., 1][mask] = torch.masked_select(coord_stack[..., 1], mask).float()
+            directions_stack[..., 2][mask] = -1.0
+
+        if CameraType.FISHEYE.value in cam_types:
+            mask = (self.camera_type[true_indices] == CameraType.FISHEYE.value).squeeze(-1)  # (num_rays)
+            mask = torch.stack([mask, mask, mask], dim=0)
+
             theta = torch.sqrt(torch.sum(coord_stack**2, dim=-1))
             theta = torch.clip(theta, 0.0, math.pi)
 
             sin_theta = torch.sin(theta)
-            directions_stack = torch.stack(
-                [coord_stack[..., 0] * sin_theta / theta, coord_stack[..., 1] * sin_theta / theta, -torch.cos(theta)],
-                dim=-1,
-            )
-        else:
-            raise ValueError(f"Camera type {CameraType(self.camera_type[0])} not supported.")
 
-        c2w = self.camera_to_worlds[camera_indices]
+            directions_stack[..., 0][mask] = torch.masked_select(coord_stack[..., 0] * sin_theta / theta, mask).float()
+            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)
+
+        for value in cam_types:
+            if value not in [CameraType.PERSPECTIVE.value, CameraType.FISHEYE.value]:
+                raise ValueError(f"Camera type {value} not supported.")
+
+        assert directions_stack.shape == (3,) + num_rays_shape + (3,)
+
+        c2w = self.camera_to_worlds[true_indices]
+        assert c2w.shape == num_rays_shape + (3, 4)
+
         if camera_opt_to_camera is not None:
             c2w = pose_utils.multiply(c2w, camera_opt_to_camera)
         rotation = c2w[..., :3, :3]  # (..., 3, 3)
+        assert rotation.shape == num_rays_shape + (3, 3)
+
         directions_stack = torch.sum(
             directions_stack[..., None, :] * rotation, dim=-1
         )  # (..., 1, 3) * (..., 3, 3) -> (..., 3)
-
         directions_stack = normalize(directions_stack, dim=-1)
+        assert directions_stack.shape == (3,) + num_rays_shape + (3,)
 
         origins = c2w[..., :3, 3]  # (..., 3)
+        assert origins.shape == num_rays_shape + (3,)
+
         directions = directions_stack[0]
+        assert directions.shape == num_rays_shape + (3,)
 
-        dx = torch.sqrt(torch.sum((directions - directions_stack[1]) ** 2, dim=-1))
-        dy = torch.sqrt(torch.sum((directions - directions_stack[2]) ** 2, dim=-1))
-        pixel_area = (dx * dy)[..., None]
+        # norms of the vector going between adjacent coords, giving us dx and dy per output ray
+        dx = torch.sqrt(torch.sum((directions - directions_stack[1]) ** 2, dim=-1))  # ("num_rays":...,)
+        dy = torch.sqrt(torch.sum((directions - directions_stack[2]) ** 2, dim=-1))  # ("num_rays":...,)
+        assert dx.shape == num_rays_shape and dy.shape == num_rays_shape
 
-        if not isinstance(camera_indices, torch.Tensor):
-            ray_bundle_camera_indices = torch.Tensor([camera_indices]).broadcast_to(pixel_area.shape).to(self.device)
-        else:
-            ray_bundle_camera_indices = camera_indices.view(pixel_area.shape)
+        pixel_area = (dx * dy)[..., None]  # ("num_rays":..., 1)
+        assert pixel_area.shape == num_rays_shape + (1,)
 
         times = self.times[camera_indices, None] if self.times is not None else None
 
@@ -343,7 +661,7 @@ class Cameras:
             origins=origins,
             directions=directions,
             pixel_area=pixel_area,
-            camera_indices=ray_bundle_camera_indices,
+            camera_indices=camera_indices,
             times=times,
         )
 
@@ -360,15 +678,16 @@ class Cameras:
         Returns:
             A JSON representation of the camera
         """
+        flattened = self.flatten()
         json_ = {
             "type": "PinholeCamera",
-            "cx": self.cx[camera_idx].item(),
-            "cy": self.cy[camera_idx].item(),
-            "fx": self.fx[camera_idx].tolist(),
-            "fy": self.fy[camera_idx].tolist(),
+            "cx": flattened[camera_idx].cx.item(),
+            "cy": flattened[camera_idx].cy.item(),
+            "fx": flattened[camera_idx].fx.item(),
+            "fy": flattened[camera_idx].fy.item(),
             "camera_to_world": self.camera_to_worlds[camera_idx].tolist(),
             "camera_index": camera_idx,
-            "times": self.times[camera_idx].item() if self.times is not None else None,
+            "times": flattened[camera_idx].times.item() if self.times is not None else None,
         }
         if image is not None:
             image_uint8 = (image * 255).detach().type(torch.uint8)
@@ -381,63 +700,42 @@ class Cameras:
             json_["image"] = str("data:image/jpeg;base64," + base64.b64encode(data).decode("ascii"))
         return json_
 
-    def get_intrinsics_matrices(self) -> TensorType["num_cameras", 3, 3]:
+    def get_intrinsics_matrices(self) -> TensorType["num_cameras":..., 3, 3]:
         """Returns the intrinsic matrices for each camera.
 
         Returns:
             Pinhole camera intrinsics matrices
         """
-        K = torch.zeros((self.size, 3, 3), dtype=torch.float32)
-        K[:, 0, 0] = self.fx
-        K[:, 1, 1] = self.fy
-        K[:, 0, 2] = self.cx
-        K[:, 1, 2] = self.cy
-        K[:, 2, 2] = 1.0
+        K = torch.zeros((*self.shape, 3, 3), dtype=torch.float32)
+        K[..., 0, 0] = self.fx.squeeze(-1)
+        K[..., 1, 1] = self.fy.squeeze(-1)
+        K[..., 0, 2] = self.cx.squeeze(-1)
+        K[..., 1, 2] = self.cy.squeeze(-1)
+        K[..., 2, 2] = 1.0
         return K
 
-    def rescale_output_resolution(self, scaling_factor: Union[TensorType["num_cameras"], float]) -> None:
+    def rescale_output_resolution(
+        self, scaling_factor: Union[TensorType["num_cameras":...], TensorType["num_cameras":..., 1], float, int]
+    ) -> None:
         """Rescale the output resolution of the cameras.
 
         Args:
             scaling_factor: Scaling factor to apply to the output resolution.
         """
-        if isinstance(scaling_factor, float):
-            scaling_factor = torch.tensor([scaling_factor]).to(self.device).broadcast_to((self.size))
+        if isinstance(scaling_factor, (float, int)):
+            scaling_factor = torch.tensor([scaling_factor]).to(self.device).broadcast_to((self.cx.shape))
+        elif isinstance(scaling_factor, torch.Tensor) and scaling_factor.shape == self.shape:
+            scaling_factor = scaling_factor.unsqueeze(-1)
+        elif isinstance(scaling_factor, torch.Tensor) and scaling_factor.shape == (*self.shape, 1):
+            pass
+        else:
+            raise ValueError(
+                f"Scaling factor must be a float, int, or a tensor of shape {self.shape} or {(*self.shape, 1)}."
+            )
 
         self.fx = self.fx * scaling_factor
         self.fy = self.fy * scaling_factor
         self.cx = self.cx * scaling_factor
         self.cy = self.cy * scaling_factor
-        self._image_heights = (self._image_heights * scaling_factor).to(torch.int64)
-        self._image_widths = (self._image_widths * scaling_factor).to(torch.int64)
-
-    def __getitem__(self, indices):
-        if isinstance(indices, torch.Tensor):
-            return Cameras(
-                self.camera_to_worlds[indices],
-                self.fx[indices],
-                self.fy[indices],
-                self.cx[indices],
-                self.cy[indices],
-                height=self._image_heights[indices],
-                width=self._image_widths[indices],
-                distortion_params=self.distortion_params[indices] if self.distortion_params is not None else None,
-                camera_type=self.camera_type[indices],
-                times=self.times[indices] if self.times is not None else None,
-            )
-        if isinstance(indices, (int, slice)):
-            indices = (indices,)
-        return Cameras(
-            self.camera_to_worlds[indices + (slice(None),)],
-            self.fx[indices + (slice(None),)],
-            self.fy[indices + (slice(None),)],
-            self.cx[indices + (slice(None),)],
-            self.cy[indices + (slice(None),)],
-            height=self._image_heights[indices + (slice(None),)],
-            width=self._image_widths[indices + (slice(None),)],
-            distortion_params=self.distortion_params[indices + (slice(None),)]
-            if self.distortion_params is not None
-            else None,
-            camera_type=self.camera_type[indices + (slice(None),)],
-            times=self.times[indices + (slice(None),)] if self.times is not None else None,
-        )
+        self.height = (self.height * scaling_factor).to(torch.int64)
+        self.width = (self.width * scaling_factor).to(torch.int64)

nerfstudio/configs/method_configs.py

@@ -31,10 +31,16 @@ from nerfstudio.configs.base_config import (
 )
 from nerfstudio.data.datamanagers.base_datamanager import VanillaDataManagerConfig
 from nerfstudio.data.datamanagers.semantic_datamanager import SemanticDataManagerConfig
+from nerfstudio.data.datamanagers.variable_res_datamanager import (
+    VariableResDataManagerConfig,
+)
 from nerfstudio.data.dataparsers.blender_dataparser import BlenderDataParserConfig
 from nerfstudio.data.dataparsers.dnerf_dataparser import DNeRFDataParserConfig
 from nerfstudio.data.dataparsers.friends_dataparser import FriendsDataParserConfig
 from nerfstudio.data.dataparsers.nerfstudio_dataparser import NerfstudioDataParserConfig
+from nerfstudio.data.dataparsers.phototourism_dataparser import (
+    PhototourismDataParserConfig,
+)
 from nerfstudio.engine.optimizers import AdamOptimizerConfig, RAdamOptimizerConfig
 from nerfstudio.field_components.temporal_distortions import TemporalDistortionKind
 from nerfstudio.models.instant_ngp import InstantNGPModelConfig
@@ -55,6 +61,7 @@ descriptions = {
     "vanilla-nerf": "Original NeRF model. (slow)",
     "tensorf": "tensorf",
     "dnerf": "Dynamic-NeRF model. (slow)",
+    "phototourism": "Uses the Phototourism data.",
 }
 
 method_configs["nerfacto"] = Config(
@@ -212,6 +219,36 @@ method_configs["dnerf"] = Config(
     },
 )
 
+method_configs["phototourism"] = Config(
+    method_name="phototourism",
+    trainer=TrainerConfig(
+        steps_per_eval_batch=500, steps_per_save=2000, max_num_iterations=30000, mixed_precision=True
+    ),
+    pipeline=VanillaPipelineConfig(
+        datamanager=VariableResDataManagerConfig(  # NOTE: one of the only differences with nerfacto
+            dataparser=PhototourismDataParserConfig(),  # NOTE: one of the only differences with nerfacto
+            train_num_rays_per_batch=4096,
+            eval_num_rays_per_batch=4096,
+            camera_optimizer=CameraOptimizerConfig(
+                mode="SO3xR3", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)
+            ),
+        ),
+        model=NerfactoModelConfig(eval_num_rays_per_chunk=1 << 15),
+    ),
+    optimizers={
+        "proposal_networks": {
+            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
+            "scheduler": None,
+        },
+        "fields": {
+            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
+            "scheduler": None,
+        },
+    },
+    viewer=ViewerConfig(num_rays_per_chunk=1 << 15),
+    vis="viewer",
+)
+
 AnnotatedBaseConfigUnion = tyro.conf.SuppressFixed[  # Don't show unparseable (fixed) arguments in helptext.
     tyro.conf.FlagConversionOff[
         tyro.extras.subcommand_type_from_defaults(defaults=method_configs, descriptions=descriptions)

nerfstudio/data/datamanagers/base_datamanager.py

@@ -42,6 +42,9 @@ from nerfstudio.data.dataparsers.instant_ngp_dataparser import (
 )
 from nerfstudio.data.dataparsers.nerfstudio_dataparser import NerfstudioDataParserConfig
 from nerfstudio.data.dataparsers.nuscenes_dataparser import NuScenesDataParserConfig
+from nerfstudio.data.dataparsers.phototourism_dataparser import (
+    PhototourismDataParserConfig,
+)
 from nerfstudio.data.dataparsers.record3d_dataparser import Record3DDataParserConfig
 from nerfstudio.data.datasets.base_dataset import InputDataset
 from nerfstudio.data.pixel_samplers import PixelSampler
@@ -50,6 +53,7 @@ from nerfstudio.data.utils.dataloaders import (
     FixedIndicesEvalDataloader,
     RandIndicesEvalDataloader,
 )
+from nerfstudio.data.utils.nerfstudio_collate import nerfstudio_collate
 from nerfstudio.engine.callbacks import TrainingCallback, TrainingCallbackAttributes
 from nerfstudio.model_components.ray_generators import RayGenerator
 from nerfstudio.utils.misc import IterableWrapper
@@ -66,6 +70,7 @@ AnnotatedDataParserUnion = tyro.conf.OmitSubcommandPrefixes[  # Omit prefixes of
             "nuscenes-data": NuScenesDataParserConfig(),
             "record3d-data": Record3DDataParserConfig(),
             "dnerf-data": DNeRFDataParserConfig(),
+            "phototourism-data": PhototourismDataParserConfig(),
         },
         prefix_names=False,  # Omit prefixes in subcommands themselves.
     )
@@ -266,6 +271,8 @@ class VanillaDataManagerConfig(InstantiateConfig):
     camera_optimizer: CameraOptimizerConfig = CameraOptimizerConfig()
     """Specifies the camera pose optimizer used during training. Helpful if poses are noisy, such as for data from
     Record3D."""
+    collate_fn = staticmethod(nerfstudio_collate)
+    """Specifies the collate function to use for the train and eval dataloaders."""
     camera_res_scale_factor: float = 1.0
     """The scale factor for scaling spatial data such as images, mask, semantics
     along with relevant information about camera intrinsics
@@ -336,6 +343,7 @@ class VanillaDataManager(DataManager):  # pylint: disable=abstract-method
             device=self.device,
             num_workers=self.world_size * 4,
             pin_memory=True,
+            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)
@@ -358,6 +366,7 @@ class VanillaDataManager(DataManager):  # pylint: disable=abstract-method
             device=self.device,
             num_workers=self.world_size * 4,
             pin_memory=True,
+            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)

nerfstudio/data/datamanagers/variable_res_datamanager.py

@@ -0,0 +1,66 @@
+# Copyright 2022 The Nerfstudio Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Data loader for variable resolution datasets, where batching raw image tensors isn't possible.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Dict, List
+
+from nerfstudio.data.datamanagers.base_datamanager import VanillaDataManagerConfig
+from nerfstudio.data.utils.nerfstudio_collate import nerfstudio_collate
+
+
+def variable_res_collate(batch: List[Dict]) -> Dict:
+    """Default collate function for the cached dataloader.
+    Args:
+        batch: Batch of samples from the dataset.
+    Returns:
+        Collated batch.
+    """
+    images = []
+    masks = []
+    for data in batch:
+        image = data.pop("image")
+        mask = data.pop("mask", None)
+        images.append(image)
+        if mask:
+            masks.append(mask)
+
+    new_batch: dict = nerfstudio_collate(batch)
+    new_batch["image"] = images
+    if masks:
+        new_batch["mask"] = masks
+
+    return new_batch
+
+
+@dataclass
+class VariableResDataManagerConfig(VanillaDataManagerConfig):
+    """A datamanager for variable resolution datasets, with presets to optimize
+    for the fact that we are now dealing with lists of images and masks.
+    """
+
+    train_num_images_to_sample_from: int = 40
+    """Number of images to sample during training iteration."""
+    train_num_times_to_repeat_images: int = 100
+    """When not training on all images, number of iterations before picking new
+    images. If -1, never pick new images."""
+    eval_num_images_to_sample_from: int = 40
+    """Number of images to sample during eval iteration."""
+    eval_num_times_to_repeat_images: int = 100
+    collate_fn = staticmethod(variable_res_collate)

nerfstudio/data/dataparsers/phototourism_dataparser.py

@@ -0,0 +1,180 @@
+# Copyright 2022 The Nerfstudio Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Phototourism dataset parser. Datasets and documentation here: http://phototour.cs.washington.edu/datasets/"""
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Type
+
+import numpy as np
+import torch
+from rich.progress import Console
+from typing_extensions import Literal
+
+from nerfstudio.cameras import camera_utils
+from nerfstudio.cameras.cameras import Cameras, CameraType
+from nerfstudio.data.dataparsers.base_dataparser import (
+    DataParser,
+    DataParserConfig,
+    DataparserOutputs,
+)
+from nerfstudio.data.scene_box import SceneBox
+from nerfstudio.data.utils.colmap_utils import read_cameras_binary, read_images_binary
+
+CONSOLE = Console(width=120)
+
+
+@dataclass
+class PhototourismDataParserConfig(DataParserConfig):
+    """Phototourism dataset parser config"""
+
+    _target: Type = field(default_factory=lambda: Phototourism)
+    """target class to instantiate"""
+    data: Path = Path("data/phototourism/brandenburg-gate")
+    """Directory specifying location of data."""
+    scale_factor: float = 3.0
+    """How much to scale the camera origins by."""
+    alpha_color: str = "white"
+    """alpha color of background"""
+    train_split_percentage: float = 0.9
+    """The percent of images to use for training. The remaining images are for eval."""
+    scene_scale: float = 1.0
+    """How much to scale the region of interest by."""
+    orientation_method: Literal["pca", "up", "none"] = "up"
+    """The method to use for orientation."""
+    auto_scale_poses: bool = True
+    """Whether to automatically scale the poses to fit in +/- 1 bounding box."""
+    center_poses: bool = True
+    """Whether to center the poses."""
+
+
+@dataclass
+class Phototourism(DataParser):
+    """Phototourism dataset. This is based on https://github.com/kwea123/nerf_pl/blob/nerfw/datasets/phototourism.py
+    and uses colmap's utils file to read the poses.
+    """
+
+    config: PhototourismDataParserConfig
+
+    def __init__(self, config: PhototourismDataParserConfig):
+        super().__init__(config=config)
+        self.data: Path = config.data
+
+    # pylint: disable=too-many-statements
+    def _generate_dataparser_outputs(self, split="train"):
+
+        image_filenames = []
+        poses = []
+
+        with CONSOLE.status(f"[bold green]Reading phototourism images and poses for {split} split...") as _:
+            cams = read_cameras_binary(self.data / "dense/sparse/cameras.bin")
+            imgs = read_images_binary(self.data / "dense/sparse/images.bin")
+
+        poses = []
+        fxs = []
+        fys = []
+        cxs = []
+        cys = []
+        image_filenames = []
+
+        flip = torch.eye(3)
+        flip[0, 0] = -1.0
+        flip = flip.double()
+
+        for _id, cam in cams.items():
+            img = imgs[_id]
+
+            assert cam.model == "PINHOLE", "Only pinhole (perspective) camera model is supported at the moment"
+
+            pose = torch.cat([torch.tensor(img.qvec2rotmat()), torch.tensor(img.tvec.reshape(3, 1))], dim=1)
+            pose = torch.cat([pose, torch.tensor([[0.0, 0.0, 0.0, 1.0]])], dim=0)
+            poses.append(torch.linalg.inv(pose))
+            fxs.append(torch.tensor(cam.params[0]))
+            fys.append(torch.tensor(cam.params[1]))
+            cxs.append(torch.tensor(cam.params[2]))
+            cys.append(torch.tensor(cam.params[3]))
+
+            image_filenames.append(self.data / "dense/images" / img.name)
+
+        poses = torch.stack(poses).float()
+        poses[..., 1:3] *= -1
+        fxs = torch.stack(fxs).float()
+        fys = torch.stack(fys).float()
+        cxs = torch.stack(cxs).float()
+        cys = torch.stack(cys).float()
+
+        # filter image_filenames and poses based on train/eval split percentage
+        num_images = len(image_filenames)
+        num_train_images = math.ceil(num_images * self.config.train_split_percentage)
+        num_eval_images = num_images - num_train_images
+        i_all = np.arange(num_images)
+        i_train = np.linspace(
+            0, num_images - 1, num_train_images, dtype=int
+        )  # equally spaced training images starting and ending at 0 and num_images-1
+        i_eval = np.setdiff1d(i_all, i_train)  # eval images are the remaining images
+        i_all = torch.tensor(i_all)
+        i_train = torch.tensor(i_train)
+        i_eval = torch.tensor(i_eval)
+        assert len(i_eval) == num_eval_images
+        if split == "train":
+            indices = i_train
+        elif split in ["val", "test"]:
+            indices = i_eval
+        else:
+            raise ValueError(f"Unknown dataparser split {split}")
+
+        poses = camera_utils.auto_orient_and_center_poses(
+            poses, method=self.config.orientation_method, center_poses=self.config.center_poses
+        )
+
+        # Scale poses
+        scale_factor = 1.0
+        if self.config.auto_scale_poses:
+            scale_factor /= torch.max(torch.abs(poses[:, :3, 3]))
+
+        poses[:, :3, 3] *= scale_factor * self.config.scale_factor
+
+        # in x,y,z order
+        # assumes that the scene is centered at the origin
+        aabb_scale = self.config.scene_scale
+        scene_box = SceneBox(
+            aabb=torch.tensor(
+                [[-aabb_scale, -aabb_scale, -aabb_scale], [aabb_scale, aabb_scale, aabb_scale]], dtype=torch.float32
+            )
+        )
+
+        cameras = Cameras(
+            camera_to_worlds=poses[:, :3, :4],
+            fx=fxs,
+            fy=fys,
+            cx=cxs,
+            cy=cys,
+            camera_type=CameraType.PERSPECTIVE,
+        )
+
+        cameras = cameras[indices]
+        image_filenames = [image_filenames[i] for i in indices]
+
+        assert len(cameras) == len(image_filenames)
+
+        dataparser_outputs = DataparserOutputs(
+            image_filenames=image_filenames,
+            cameras=cameras,
+            scene_box=scene_box,
+        )
+
+        return dataparser_outputs

nerfstudio/data/pixel_samplers.py

@@ -62,6 +62,77 @@ def collate_image_dataset_batch(batch: Dict, num_rays_per_batch: int, keep_full_
     return collated_batch
 
 
+def collate_image_dataset_batch_list(batch: Dict, num_rays_per_batch: int, keep_full_image: bool = False):
+    """
+    Does the same as collate_image_dataset_batch, except it will operate over a list of images / masks inside
+    a list.
+
+    We will use this with the intent of DEPRECIATING it as soon as we find a viable alternative.
+    The intention will be to replace this with a more efficient implementation that doesn't require a for loop, but
+    since pytorch's ragged tensors are still in beta (this would allow for some vectorization), this will do
+
+    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"][0].device
+    num_images = len(batch["image"])
+
+    # only sample within the mask, if the mask is in the batch
+    all_indices = []
+    all_images = []
+
+    if "mask" in batch:
+        num_rays_in_batch = num_rays_per_batch // num_images
+        for i in range(num_images):
+            if i == num_images - 1:
+                num_rays_in_batch = num_rays_per_batch - (num_images - 1) * num_rays_in_batch
+            nonzero_indices = torch.nonzero(batch["mask"][i][..., 0], as_tuple=False)
+            chosen_indices = random.sample(range(len(nonzero_indices)), k=num_rays_in_batch)
+            indices = nonzero_indices[chosen_indices]
+            indices = torch.cat([torch.full((num_rays_in_batch, 1), i, device=device), indices], dim=-1)
+            all_indices.append(indices)
+            all_images.append(batch["image"][i][indices[:, 1], indices[:, 2]])
+
+    else:
+        num_rays_in_batch = num_rays_per_batch // num_images
+        for i in range(num_images):
+            image_height, image_width, _ = batch["image"][i].shape
+            if i == num_images - 1:
+                num_rays_in_batch = num_rays_per_batch - (num_images - 1) * num_rays_in_batch
+            indices = torch.floor(
+                torch.rand((num_rays_in_batch, 3), device=device)
+                * torch.tensor([1, image_height, image_width], device=device)
+            ).long()
+            indices[:, 0] = i
+            all_indices.append(indices)
+            all_images.append(batch["image"][i][indices[:, 1], indices[:, 2]])
+
+    indices = torch.cat(all_indices, dim=0)
+
+    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 key != "image" and key != "mask" and value is not None
+    }
+
+    collated_batch["image"] = torch.cat(all_images, dim=0)
+
+    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 PixelSampler:  # pylint: disable=too-few-public-methods
     """Samples 'pixel_batch's from 'image_batch's.
 
@@ -88,7 +159,15 @@ class PixelSampler:  # pylint: disable=too-few-public-methods
         Args:
             image_batch: batch of images to sample from
         """
-        pixel_batch = collate_image_dataset_batch(
-            image_batch, self.num_rays_per_batch, keep_full_image=self.keep_full_image
-        )
+        if isinstance(image_batch["image"], list):
+            image_batch = dict(image_batch.items())  # copy the dictioary so we don't modify the original
+            pixel_batch = collate_image_dataset_batch_list(
+                image_batch, self.num_rays_per_batch, keep_full_image=self.keep_full_image
+            )
+        elif isinstance(image_batch["image"], torch.Tensor):
+            pixel_batch = collate_image_dataset_batch(
+                image_batch, self.num_rays_per_batch, keep_full_image=self.keep_full_image
+            )
+        else:
+            raise ValueError("image_batch['image'] must be a list or torch.Tensor")
         return pixel_batch

nerfstudio/data/utils/colmap_utils.py

@@ -0,0 +1,303 @@
+# Copyright (c) 2018, ETH Zurich and UNC Chapel Hill.
+# All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+#     * Redistributions of source code must retain the above copyright
+#       notice, this list of conditions and the following disclaimer.
+#
+#     * Redistributions in binary form must reproduce the above copyright
+#       notice, this list of conditions and the following disclaimer in the
+#       documentation and/or other materials provided with the distribution.
+#
+#     * Neither the name of ETH Zurich and UNC Chapel Hill nor the names of
+#       its contributors may be used to endorse or promote products derived
+#       from this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS BE
+# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+# POSSIBILITY OF SUCH DAMAGE.
+#
+# Author: Johannes L. Schoenberger (jsch at inf.ethz.ch)
+
+# pylint: disable-all
+# Disabling all pylinting since this file is copy-pasted over from COLMAP
+
+import collections
+import os
+import struct
+
+import numpy as np
+
+CameraModel = collections.namedtuple("CameraModel", ["model_id", "model_name", "num_params"])
+Camera = collections.namedtuple("Camera", ["id", "model", "width", "height", "params"])
+BaseImage = collections.namedtuple("Image", ["id", "qvec", "tvec", "camera_id", "name", "xys", "point3D_ids"])
+Point3D = collections.namedtuple("Point3D", ["id", "xyz", "rgb", "error", "image_ids", "point2D_idxs"])
+
+
+class Image(BaseImage):
+    def qvec2rotmat(self):
+        return qvec2rotmat(self.qvec)
+
+
+CAMERA_MODELS = {
+    CameraModel(model_id=0, model_name="SIMPLE_PINHOLE", num_params=3),
+    CameraModel(model_id=1, model_name="PINHOLE", num_params=4),
+    CameraModel(model_id=2, model_name="SIMPLE_RADIAL", num_params=4),
+    CameraModel(model_id=3, model_name="RADIAL", num_params=5),
+    CameraModel(model_id=4, model_name="OPENCV", num_params=8),
+    CameraModel(model_id=5, model_name="OPENCV_FISHEYE", num_params=8),
+    CameraModel(model_id=6, model_name="FULL_OPENCV", num_params=12),
+    CameraModel(model_id=7, model_name="FOV", num_params=5),
+    CameraModel(model_id=8, model_name="SIMPLE_RADIAL_FISHEYE", num_params=4),
+    CameraModel(model_id=9, model_name="RADIAL_FISHEYE", num_params=5),
+    CameraModel(model_id=10, model_name="THIN_PRISM_FISHEYE", num_params=12),
+}
+CAMERA_MODEL_IDS = dict([(camera_model.model_id, camera_model) for camera_model in CAMERA_MODELS])
+
+
+def read_next_bytes(fid, num_bytes, format_char_sequence, endian_character="<"):
+    """Read and unpack the next bytes from a binary file.
+    :param fid:
+    :param num_bytes: Sum of combination of {2, 4, 8}, e.g. 2, 6, 16, 30, etc.
+    :param format_char_sequence: List of {c, e, f, d, h, H, i, I, l, L, q, Q}.
+    :param endian_character: Any of {@, =, <, >, !}
+    :return: Tuple of read and unpacked values.
+    """
+    data = fid.read(num_bytes)
+    return struct.unpack(endian_character + format_char_sequence, data)
+
+
+def read_cameras_text(path):
+    """
+    see: src/base/reconstruction.cc
+        void Reconstruction::WriteCamerasText(const std::string& path)
+        void Reconstruction::ReadCamerasText(const std::string& path)
+    """
+    cameras = {}
+    with open(path, "r") as fid:
+        while True:
+            line = fid.readline()
+            if not line:
+                break
+            line = line.strip()
+            if len(line) > 0 and line[0] != "#":
+                elems = line.split()
+                camera_id = int(elems[0])
+                model = elems[1]
+                width = int(elems[2])
+                height = int(elems[3])
+                params = np.array(tuple(map(float, elems[4:])))
+                cameras[camera_id] = Camera(id=camera_id, model=model, width=width, height=height, params=params)
+    return cameras
+
+
+def read_cameras_binary(path_to_model_file):
+    """
+    see: src/base/reconstruction.cc
+        void Reconstruction::WriteCamerasBinary(const std::string& path)
+        void Reconstruction::ReadCamerasBinary(const std::string& path)
+    """
+    cameras = {}
+    with open(path_to_model_file, "rb") as fid:
+        num_cameras = read_next_bytes(fid, 8, "Q")[0]
+        for camera_line_index in range(num_cameras):
+            camera_properties = read_next_bytes(fid, num_bytes=24, format_char_sequence="iiQQ")
+            camera_id = camera_properties[0]
+            model_id = camera_properties[1]
+            model_name = CAMERA_MODEL_IDS[camera_properties[1]].model_name
+            width = camera_properties[2]
+            height = camera_properties[3]
+            num_params = CAMERA_MODEL_IDS[model_id].num_params
+            params = read_next_bytes(fid, num_bytes=8 * num_params, format_char_sequence="d" * num_params)
+            cameras[camera_id] = Camera(
+                id=camera_id, model=model_name, width=width, height=height, params=np.array(params)
+            )
+        assert len(cameras) == num_cameras
+    return cameras
+
+
+def read_images_text(path):
+    """
+    see: src/base/reconstruction.cc
+        void Reconstruction::ReadImagesText(const std::string& path)
+        void Reconstruction::WriteImagesText(const std::string& path)
+    """
+    images = {}
+    with open(path, "r") as fid:
+        while True:
+            line = fid.readline()
+            if not line:
+                break
+            line = line.strip()
+            if len(line) > 0 and line[0] != "#":
+                elems = line.split()
+                image_id = int(elems[0])
+                qvec = np.array(tuple(map(float, elems[1:5])))
+                tvec = np.array(tuple(map(float, elems[5:8])))
+                camera_id = int(elems[8])
+                image_name = elems[9]
+                elems = fid.readline().split()
+                xys = np.column_stack([tuple(map(float, elems[0::3])), tuple(map(float, elems[1::3]))])
+                point3D_ids = np.array(tuple(map(int, elems[2::3])))
+                images[image_id] = Image(
+                    id=image_id,
+                    qvec=qvec,
+                    tvec=tvec,
+                    camera_id=camera_id,
+                    name=image_name,
+                    xys=xys,
+                    point3D_ids=point3D_ids,
+                )
+    return images
+
+
+def read_images_binary(path_to_model_file):
+    """
+    see: src/base/reconstruction.cc
+        void Reconstruction::ReadImagesBinary(const std::string& path)
+        void Reconstruction::WriteImagesBinary(const std::string& path)
+    """
+    images = {}
+    with open(path_to_model_file, "rb") as fid:
+        num_reg_images = read_next_bytes(fid, 8, "Q")[0]
+        for image_index in range(num_reg_images):
+            binary_image_properties = read_next_bytes(fid, num_bytes=64, format_char_sequence="idddddddi")
+            image_id = binary_image_properties[0]
+            qvec = np.array(binary_image_properties[1:5])
+            tvec = np.array(binary_image_properties[5:8])
+            camera_id = binary_image_properties[8]
+            image_name = ""
+            current_char = read_next_bytes(fid, 1, "c")[0]
+            while current_char != b"\x00":  # look for the ASCII 0 entry
+                image_name += current_char.decode("utf-8")
+                current_char = read_next_bytes(fid, 1, "c")[0]
+            num_points2D = read_next_bytes(fid, num_bytes=8, format_char_sequence="Q")[0]
+            x_y_id_s = read_next_bytes(fid, num_bytes=24 * num_points2D, format_char_sequence="ddq" * num_points2D)
+            xys = np.column_stack([tuple(map(float, x_y_id_s[0::3])), tuple(map(float, x_y_id_s[1::3]))])
+            point3D_ids = np.array(tuple(map(int, x_y_id_s[2::3])))
+            images[image_id] = Image(
+                id=image_id,
+                qvec=qvec,
+                tvec=tvec,
+                camera_id=camera_id,
+                name=image_name,
+                xys=xys,
+                point3D_ids=point3D_ids,
+            )
+    return images
+
+
+def read_points3D_text(path):
+    """
+    see: src/base/reconstruction.cc
+        void Reconstruction::ReadPoints3DText(const std::string& path)
+        void Reconstruction::WritePoints3DText(const std::string& path)
+    """
+    points3D = {}
+    with open(path, "r") as fid:
+        while True:
+            line = fid.readline()
+            if not line:
+                break
+            line = line.strip()
+            if len(line) > 0 and line[0] != "#":
+                elems = line.split()
+                point3D_id = int(elems[0])
+                xyz = np.array(tuple(map(float, elems[1:4])))
+                rgb = np.array(tuple(map(int, elems[4:7])))
+                error = float(elems[7])
+                image_ids = np.array(tuple(map(int, elems[8::2])))
+                point2D_idxs = np.array(tuple(map(int, elems[9::2])))
+                points3D[point3D_id] = Point3D(
+                    id=point3D_id, xyz=xyz, rgb=rgb, error=error, image_ids=image_ids, point2D_idxs=point2D_idxs
+                )
+    return points3D
+
+
+def read_points3d_binary(path_to_model_file):
+    """
+    see: src/base/reconstruction.cc
+        void Reconstruction::ReadPoints3DBinary(const std::string& path)
+        void Reconstruction::WritePoints3DBinary(const std::string& path)
+    """
+    points3D = {}
+    with open(path_to_model_file, "rb") as fid:
+        num_points = read_next_bytes(fid, 8, "Q")[0]
+        for point_line_index in range(num_points):
+            binary_point_line_properties = read_next_bytes(fid, num_bytes=43, format_char_sequence="QdddBBBd")
+            point3D_id = binary_point_line_properties[0]
+            xyz = np.array(binary_point_line_properties[1:4])
+            rgb = np.array(binary_point_line_properties[4:7])
+            error = np.array(binary_point_line_properties[7])
+            track_length = read_next_bytes(fid, num_bytes=8, format_char_sequence="Q")[0]
+            track_elems = read_next_bytes(fid, num_bytes=8 * track_length, format_char_sequence="ii" * track_length)
+            image_ids = np.array(tuple(map(int, track_elems[0::2])))
+            point2D_idxs = np.array(tuple(map(int, track_elems[1::2])))
+            points3D[point3D_id] = Point3D(
+                id=point3D_id, xyz=xyz, rgb=rgb, error=error, image_ids=image_ids, point2D_idxs=point2D_idxs
+            )
+    return points3D
+
+
+def read_model(path, ext):
+    if ext == ".txt":
+        cameras = read_cameras_text(os.path.join(path, "cameras" + ext))
+        images = read_images_text(os.path.join(path, "images" + ext))
+        points3D = read_points3D_text(os.path.join(path, "points3D") + ext)
+    else:
+        cameras = read_cameras_binary(os.path.join(path, "cameras" + ext))
+        images = read_images_binary(os.path.join(path, "images" + ext))
+        points3D = read_points3d_binary(os.path.join(path, "points3D") + ext)
+    return cameras, images, points3D
+
+
+def qvec2rotmat(qvec):
+    return np.array(
+        [
+            [
+                1 - 2 * qvec[2] ** 2 - 2 * qvec[3] ** 2,
+                2 * qvec[1] * qvec[2] - 2 * qvec[0] * qvec[3],
+                2 * qvec[3] * qvec[1] + 2 * qvec[0] * qvec[2],
+            ],
+            [
+                2 * qvec[1] * qvec[2] + 2 * qvec[0] * qvec[3],
+                1 - 2 * qvec[1] ** 2 - 2 * qvec[3] ** 2,
+                2 * qvec[2] * qvec[3] - 2 * qvec[0] * qvec[1],
+            ],
+            [
+                2 * qvec[3] * qvec[1] - 2 * qvec[0] * qvec[2],
+                2 * qvec[2] * qvec[3] + 2 * qvec[0] * qvec[1],
+                1 - 2 * qvec[1] ** 2 - 2 * qvec[2] ** 2,
+            ],
+        ]
+    )
+
+
+def rotmat2qvec(R):
+    Rxx, Ryx, Rzx, Rxy, Ryy, Rzy, Rxz, Ryz, Rzz = R.flat
+    K = (
+        np.array(
+            [
+                [Rxx - Ryy - Rzz, 0, 0, 0],
+                [Ryx + Rxy, Ryy - Rxx - Rzz, 0, 0],
+                [Rzx + Rxz, Rzy + Ryz, Rzz - Rxx - Ryy, 0],
+                [Ryz - Rzy, Rzx - Rxz, Rxy - Ryx, Rxx + Ryy + Rzz],
+            ]
+        )
+        / 3.0
+    )
+    eigvals, eigvecs = np.linalg.eigh(K)
+    qvec = eigvecs[[3, 0, 1, 2], np.argmax(eigvals)]
+    if qvec[0] < 0:
+        qvec *= -1
+    return qvec

nerfstudio/data/utils/dataloaders.py

@@ -25,12 +25,13 @@ from typing import Dict, Optional, Tuple, Union
 
 import torch
 from rich.progress import Console, track
-from torch.utils.data import Dataset, default_collate
+from torch.utils.data import Dataset
 from torch.utils.data.dataloader import DataLoader
 
 from nerfstudio.cameras.cameras import Cameras
 from nerfstudio.cameras.rays import RayBundle
 from nerfstudio.data.datasets.base_dataset import InputDataset
+from nerfstudio.data.utils.nerfstudio_collate import nerfstudio_collate
 from nerfstudio.utils.misc import get_dict_to_torch
 
 CONSOLE = Console(width=120)
@@ -45,6 +46,7 @@ class CacheDataloader(DataLoader):
         num_samples_to_collate: How many images to sample rays for each batch. -1 for all images.
         num_times_to_repeat_images: How often to collate new images. -1 to never pick new images.
         device: Device to perform computation.
+        collate_fn: The function we will use to collate our training data
     """
 
     def __init__(
@@ -53,13 +55,16 @@ class CacheDataloader(DataLoader):
         num_images_to_sample_from: int = -1,
         num_times_to_repeat_images: int = -1,
         device: Union[torch.device, str] = "cpu",
+        collate_fn=nerfstudio_collate,
         **kwargs,
     ):
         self.dataset = dataset
+        super().__init__(dataset=dataset, **kwargs)  # This will set self.dataset
         self.num_times_to_repeat_images = num_times_to_repeat_images
         self.cache_all_images = (num_images_to_sample_from == -1) or (num_images_to_sample_from >= len(self.dataset))
         self.num_images_to_sample_from = len(self.dataset) if self.cache_all_images else num_images_to_sample_from
         self.device = device
+        self.collate_fn = collate_fn
         self.num_workers = kwargs.get("num_workers", 0)
 
         self.num_repeated = self.num_times_to_repeat_images  # starting value
@@ -82,7 +87,6 @@ class CacheDataloader(DataLoader):
                 f"Caching {self.num_images_to_sample_from} out of {len(self.dataset)} images, "
                 f"resampling every {self.num_times_to_repeat_images} iters."
             )
-        super().__init__(dataset=dataset, **kwargs)
 
     def __getitem__(self, idx):
         return self.dataset.__getitem__(idx)
@@ -103,7 +107,7 @@ class CacheDataloader(DataLoader):
                 res = executor.submit(self.dataset.__getitem__, idx)
                 results.append(res)
 
-            for res in track(results, description="Loading data batch"):
+            for res in track(results, description="Loading data batch", transient=True):
                 batch_list.append(res.result())
 
         return batch_list
@@ -111,7 +115,7 @@ class CacheDataloader(DataLoader):
     def _get_collated_batch(self):
         """Returns a collated batch."""
         batch_list = self._get_batch_list()
-        collated_batch = default_collate(batch_list)
+        collated_batch = self.collate_fn(batch_list)
         collated_batch = get_dict_to_torch(collated_batch, device=self.device, exclude=["image"])
         return collated_batch
 
@@ -169,23 +173,7 @@ class EvalDataloader(DataLoader):
         Args:
             image_idx: Camera image index
         """
-        distortion_params = None
-        if self.cameras.distortion_params is not None:
-            distortion_params = self.cameras.distortion_params[image_idx]
-
-        camera = Cameras(
-            fx=self.cameras.fx[image_idx : image_idx + 1],
-            fy=self.cameras.fy[image_idx : image_idx + 1],
-            cx=self.cameras.cx[image_idx : image_idx + 1],
-            cy=self.cameras.cy[image_idx : image_idx + 1],
-            height=self.cameras.image_height[image_idx : image_idx + 1],
-            width=self.cameras.image_width[image_idx : image_idx + 1],
-            camera_to_worlds=self.cameras.camera_to_worlds[image_idx : image_idx + 1],
-            distortion_params=distortion_params,
-            camera_type=self.cameras.camera_type[image_idx : image_idx + 1],
-            times=self.cameras.times[image_idx : image_idx + 1] if self.cameras.times else None,
-        )
-        return camera
+        return self.cameras[image_idx]
 
     def get_data_from_image_idx(self, image_idx: int) -> Tuple[RayBundle, Dict]:
         """Returns the data for a specific image index.
@@ -193,7 +181,7 @@ class EvalDataloader(DataLoader):
         Args:
             image_idx: Camera image index
         """
-        ray_bundle = self.cameras.generate_rays(camera_indices=image_idx)
+        ray_bundle = self.cameras.generate_rays(camera_indices=image_idx, keep_shape=True)
         batch = self.input_dataset[image_idx]
         batch = get_dict_to_torch(batch, device=self.device, exclude=["image"])
         return ray_bundle, batch

nerfstudio/data/utils/nerfstudio_collate.py

@@ -0,0 +1,198 @@
+# Copyright 2022 The Plenoptix Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Custom collate function that includes cases for nerfstudio types.
+"""
+
+import collections
+import collections.abc
+import re
+from typing import Callable, Dict, Union
+
+import torch
+import torch.utils.data
+from torch._six import string_classes
+
+from nerfstudio.cameras.cameras import Cameras
+
+# pylint: disable=implicit-str-concat
+NERFSTUDIO_COLLATE_ERR_MSG_FORMAT = (
+    "default_collate: batch must contain tensors, numpy arrays, numbers, " "dicts, lists or anything in {}; found {}"
+)
+np_str_obj_array_pattern = re.compile(r"[SaUO]")
+
+
+def nerfstudio_collate(
+    batch, extra_mappings: Union[Dict[type, Callable], None] = None
+):  # pylint: disable=too-many-return-statements
+    r"""
+    This is the default pytorch collate function, but with support for nerfstudio types. All documentation
+    below is copied straight over from pytorch's default_collate function, python version 3.8.13,
+    pytorch version '1.12.1+cu113'. Custom nerfstudio types are accounted for at the end, and extra
+    mappings can be passed in to handle custom types. These mappings are from types: callable (types
+    being like int or float or the return value of type(3.), etc). The only code before we parse for custom types that
+    was changed from default pytorch was the addition of the extra_mappings argument, a find and replace operation
+    from default_collate to nerfstudio_collate, and the addition of the nerfstudio_collate_err_msg_format variable.
+
+
+    Function that takes in a batch of data and puts the elements within the batch
+    into a tensor with an additional outer dimension - batch size. The exact output type can be
+    a :class:`torch.Tensor`, a `Sequence` of :class:`torch.Tensor`, a
+    Collection of :class:`torch.Tensor`, or left unchanged, depending on the input type.
+    This is used as the default function for collation when
+    `batch_size` or `batch_sampler` is defined in :class:`~torch.utils.data.DataLoader`.
+
+    Here is the general input type (based on the type of the element within the batch) to output type mapping:
+
+        * :class:`torch.Tensor` -> :class:`torch.Tensor` (with an added outer dimension batch size)
+        * NumPy Arrays -> :class:`torch.Tensor`
+        * `float` -> :class:`torch.Tensor`
+        * `int` -> :class:`torch.Tensor`
+        * `str` -> `str` (unchanged)
+        * `bytes` -> `bytes` (unchanged)
+        * `Mapping[K, V_i]` -> `Mapping[K, nerfstudio_collate([V_1, V_2, ...])]`
+        * `NamedTuple[V1_i, V2_i, ...]` -> `NamedTuple[nerfstudio_collate([V1_1, V1_2, ...]),
+          nerfstudio_collate([V2_1, V2_2, ...]), ...]`
+        * `Sequence[V1_i, V2_i, ...]` -> `Sequence[nerfstudio_collate([V1_1, V1_2, ...]),
+          nerfstudio_collate([V2_1, V2_2, ...]), ...]`
+
+    Args:
+        batch: a single batch to be collated
+
+    Examples:
+        >>> # Example with a batch of `int`s:
+        >>> nerfstudio_collate([0, 1, 2, 3])
+        tensor([0, 1, 2, 3])
+        >>> # Example with a batch of `str`s:
+        >>> nerfstudio_collate(['a', 'b', 'c'])
+        ['a', 'b', 'c']
+        >>> # Example with `Map` inside the batch:
+        >>> nerfstudio_collate([{'A': 0, 'B': 1}, {'A': 100, 'B': 100}])
+        {'A': tensor([  0, 100]), 'B': tensor([  1, 100])}
+        >>> # Example with `NamedTuple` inside the batch:
+        >>> Point = namedtuple('Point', ['x', 'y'])
+        >>> nerfstudio_collate([Point(0, 0), Point(1, 1)])
+        Point(x=tensor([0, 1]), y=tensor([0, 1]))
+        >>> # Example with `Tuple` inside the batch:
+        >>> nerfstudio_collate([(0, 1), (2, 3)])
+        [tensor([0, 2]), tensor([1, 3])]
+        >>> # Example with `List` inside the batch:
+        >>> nerfstudio_collate([[0, 1], [2, 3]])
+        [tensor([0, 2]), tensor([1, 3])]
+    """
+    if extra_mappings is None:
+        extra_mappings = {}
+    elem = batch[0]
+    elem_type = type(elem)
+    if isinstance(elem, torch.Tensor):  # pylint: disable=no-else-return
+        out = None
+        if torch.utils.data.get_worker_info() is not None:
+            # If we're in a background process, concatenate directly into a
+            # shared memory tensor to avoid an extra copy
+            numel = sum(x.numel() for x in batch)
+            storage = elem.storage()._new_shared(numel, device=elem.device)  # pylint: disable=protected-access
+            out = elem.new(storage).resize_(len(batch), *list(elem.size()))
+        return torch.stack(batch, 0, out=out)
+    elif elem_type.__module__ == "numpy" and elem_type.__name__ != "str_" and elem_type.__name__ != "string_":
+        # pylint: disable=no-else-return, consider-using-in
+        if elem_type.__name__ == "ndarray" or elem_type.__name__ == "memmap":
+            # array of string classes and object
+            if np_str_obj_array_pattern.search(elem.dtype.str) is not None:
+                raise TypeError(NERFSTUDIO_COLLATE_ERR_MSG_FORMAT.format(elem.dtype))
+
+            return nerfstudio_collate([torch.as_tensor(b) for b in batch], extra_mappings=extra_mappings)
+        elif elem.shape == ():  # scalars
+            return torch.as_tensor(batch)
+    elif isinstance(elem, float):
+        return torch.tensor(batch, dtype=torch.float64)
+    elif isinstance(elem, int):
+        return torch.tensor(batch)
+    elif isinstance(elem, string_classes):
+        return batch
+    elif isinstance(elem, collections.abc.Mapping):
+        try:
+            return elem_type(
+                {key: nerfstudio_collate([d[key] for d in batch], extra_mappings=extra_mappings) for key in elem}
+            )
+        except TypeError:
+            # The mapping type may not support `__init__(iterable)`.
+            return {key: nerfstudio_collate([d[key] for d in batch], extra_mappings=extra_mappings) for key in elem}
+    elif isinstance(elem, tuple) and hasattr(elem, "_fields"):  # namedtuple
+        return elem_type(*(nerfstudio_collate(samples, extra_mappings=extra_mappings) for samples in zip(*batch)))
+    elif isinstance(elem, collections.abc.Sequence):
+        # check to make sure that the elements in batch have consistent size
+        it = iter(batch)
+        elem_size = len(next(it))
+        if not all(len(elem) == elem_size for elem in it):
+            raise RuntimeError("each element in list of batch should be of equal size")
+        transposed = list(zip(*batch))  # It may be accessed twice, so we use a list.
+
+        if isinstance(elem, tuple):
+            return [
+                nerfstudio_collate(samples, extra_mappings=extra_mappings) for samples in transposed
+            ]  # Backwards compatibility.
+        else:
+            try:
+                return elem_type([nerfstudio_collate(samples, extra_mappings=extra_mappings) for samples in transposed])
+            except TypeError:
+                # The sequence type may not support `__init__(iterable)` (e.g., `range`).
+                return [nerfstudio_collate(samples, extra_mappings=extra_mappings) for samples in transposed]
+
+    # NerfStudio types supported below
+
+    elif isinstance(elem, Cameras):
+        # If a camera, just concatenate along the batch dimension. In the future, this may change to stacking
+        assert all((isinstance(cam, Cameras) for cam in batch))
+        assert all((cam.distortion_params is None for cam in batch)) or all(
+            (cam.distortion_params is not None for cam in batch)
+        ), "All cameras must have distortion parameters or none of them should have distortion parameters.\
+            Generalized batching will be supported in the future."
+
+        # If no batch dimension exists, then we need to stack everything and create a batch dimension on 0th dim
+        if elem.shape == ():
+            op = torch.stack
+        # If batch dimension exists, then we need to concatenate along the 0th dimension
+        else:
+            op = torch.cat
+
+        return Cameras(
+            op([cameras.camera_to_worlds for cameras in batch], dim=0),
+            op([cameras.fx for cameras in batch], dim=0),
+            op([cameras.fy for cameras in batch], dim=0),
+            op([cameras.cx for cameras in batch], dim=0),
+            op([cameras.cy for cameras in batch], dim=0),
+            height=op([cameras.height for cameras in batch], dim=0),
+            width=op([cameras.width for cameras in batch], dim=0),
+            distortion_params=op(
+                [
+                    cameras.distortion_params
+                    if cameras.distortion_params is not None
+                    else torch.zeros_like(cameras.distortion_params)
+                    for cameras in batch
+                ],
+                dim=0,
+            ),
+            camera_type=op([cameras.camera_type for cameras in batch], dim=0),
+            times=torch.stack(
+                [cameras.times if cameras.times is not None else -torch.ones_like(cameras.times) for cameras in batch],
+                dim=0,
+            ),
+        )
+
+    for type_key in extra_mappings:
+        if isinstance(elem, type_key):
+            return extra_mappings[type_key](batch)
+
+    raise TypeError(NERFSTUDIO_COLLATE_ERR_MSG_FORMAT.format(elem_type))

nerfstudio/model_components/ray_generators.py

@@ -52,7 +52,7 @@ class RayGenerator(nn.Module):
         camera_opt_to_camera = self.pose_optimizer(c)
 
         ray_bundle = self.cameras.generate_rays(
-            camera_indices=c,
+            camera_indices=c.unsqueeze(-1),
             coords=coords,
             camera_opt_to_camera=camera_opt_to_camera,
         )

nerfstudio/pipelines/base_pipeline.py

@@ -338,7 +338,7 @@ class VanillaPipeline(Pipeline):
                 metrics_dict, _ = self.model.get_image_metrics_and_images(outputs, batch)
                 assert "num_rays_per_sec" not in metrics_dict
                 metrics_dict["num_rays_per_sec"] = num_rays / (time() - inner_start)
-                fps_str = f"fps_at_{height}x{width}"
+                fps_str = "fps"
                 assert fps_str not in metrics_dict
                 metrics_dict[fps_str] = metrics_dict["num_rays_per_sec"] / (height * width)
                 metrics_dict_list.append(metrics_dict)

nerfstudio/utils/tensor_dataclass.py

@@ -29,6 +29,8 @@ class TensorDataclass:
     Fields that are not Tensors will not be batched unless they are also a TensorDataclass.
     Any fields that are dictionaries will have their Tensors or TensorDataclasses batched, and
     dictionaries will have their tensors or TensorDataclasses considered in the initial broadcast.
+    Tensor fields must have at least 1 dimension, meaning that you must convert a field like torch.Tensor(1)
+    to torch.Tensor([1])
 
     Example:
 
@@ -56,12 +58,24 @@ class TensorDataclass:
 
     _shape: tuple
 
+    # A mapping from field-name (str): n (int)
+    # Any field OR any key in a dictionary field with this name (field-name) and a corresponding
+    # torch.Tensor will be assumed to have n dimensions after the batch dims. These n final dimensions
+    # will remain the same shape when doing reshapes, broadcasting, etc on the tensordataclass
+    _field_custom_dimensions: Dict[str, int] = {}
+
     def __post_init__(self) -> None:
         """Finishes setting up the TensorDataclass
 
         This will 1) find the broadcasted shape and 2) broadcast all fields to this shape 3)
         set _shape to be the broadcasted shape.
         """
+        if self._field_custom_dimensions is not None:
+            for k, v in self._field_custom_dimensions.items():
+                assert (
+                    isinstance(v, int) and v > 1
+                ), f"Custom dimensions must be an integer greater than 1, since 1 is the default, received {k}: {v}"
+
         if not dataclasses.is_dataclass(self):
             raise TypeError("TensorDataclass must be a dataclass")
 
@@ -90,9 +104,13 @@ class TensorDataclass:
             The batch shapes of all tensors in the dictionary.
         """
         batch_shapes = []
-        for v in dict_.values():
+        for k, v in dict_.items():
             if isinstance(v, torch.Tensor):
-                batch_shapes.append(v.shape[:-1])
+                if isinstance(self._field_custom_dimensions, dict) and k in self._field_custom_dimensions:
+                    # pylint: disable=unsubscriptable-object
+                    batch_shapes.append(v.shape[: -self._field_custom_dimensions[k]])
+                else:
+                    batch_shapes.append(v.shape[:-1])
             elif isinstance(v, TensorDataclass):
                 batch_shapes.append(v.shape)
             elif isinstance(v, Dict):
@@ -111,7 +129,17 @@ class TensorDataclass:
         new_dict = {}
         for k, v in dict_.items():
             if isinstance(v, torch.Tensor):
-                new_dict[k] = v.broadcast_to((*batch_shape, v.shape[-1]))
+                # If custom dimension key, then we need to
+                if isinstance(self._field_custom_dimensions, dict) and k in self._field_custom_dimensions:
+                    # pylint: disable=unsubscriptable-object
+                    new_dict[k] = v.broadcast_to(
+                        (
+                            *batch_shape,
+                            *v.shape[-self._field_custom_dimensions[k] :],
+                        )
+                    )
+                else:
+                    new_dict[k] = v.broadcast_to((*batch_shape, v.shape[-1]))
             elif isinstance(v, TensorDataclass):
                 new_dict[k] = v.broadcast_to(batch_shape)
             elif isinstance(v, Dict):
@@ -119,18 +147,26 @@ class TensorDataclass:
         return new_dict
 
     def __getitem__(self: TensorDataclassT, indices) -> TensorDataclassT:
-        if isinstance(indices, torch.Tensor):
+        if isinstance(indices, (torch.Tensor)):
             return self._apply_fn_to_fields(lambda x: x[indices])
-        if isinstance(indices, (int, slice)):
+        if isinstance(indices, (int, slice, type(Ellipsis))):
             indices = (indices,)
+        assert isinstance(indices, tuple)
         tensor_fn = lambda x: x[indices + (slice(None),)]
         dataclass_fn = lambda x: x[indices]
-        return self._apply_fn_to_fields(tensor_fn, dataclass_fn)
+
+        def custom_tensor_dims_fn(k, v):
+            custom_dims = self._field_custom_dimensions[k]  # pylint: disable=unsubscriptable-object
+            return v[indices + ((slice(None),) * custom_dims)]
+
+        return self._apply_fn_to_fields(tensor_fn, dataclass_fn, custom_tensor_dims_fn=custom_tensor_dims_fn)
 
     def __setitem__(self, indices, value) -> NoReturn:
         raise RuntimeError("Index assignment is not supported for TensorDataclass")
 
     def __len__(self) -> int:
+        if len(self._shape) == 0:
+            raise TypeError("len() of a 0-d tensor")
         return self.shape[0]
 
     def __bool__(self) -> bool:
@@ -173,7 +209,12 @@ class TensorDataclass:
             shape = (shape,)
         tensor_fn = lambda x: x.reshape((*shape, x.shape[-1]))
         dataclass_fn = lambda x: x.reshape(shape)
-        return self._apply_fn_to_fields(tensor_fn, dataclass_fn)
+
+        def custom_tensor_dims_fn(k, v):
+            custom_dims = self._field_custom_dimensions[k]  # pylint: disable=unsubscriptable-object
+            return v.reshape((*shape, *v.shape[-custom_dims:]))
+
+        return self._apply_fn_to_fields(tensor_fn, dataclass_fn, custom_tensor_dims_fn=custom_tensor_dims_fn)
 
     def flatten(self: TensorDataclassT) -> TensorDataclassT:
         """Returns a new TensorDataclass with flattened batch dimensions
@@ -195,7 +236,14 @@ class TensorDataclass:
         Returns:
             A new TensorDataclass with the same data but with a new shape.
         """
-        return self._apply_fn_to_fields(lambda x: x.broadcast_to((*shape, x.shape[-1])))
+
+        def custom_tensor_dims_fn(k, v):
+            custom_dims = self._field_custom_dimensions[k]  # pylint: disable=unsubscriptable-object
+            return v.broadcast_to((*shape, *v.shape[-custom_dims:]))
+
+        return self._apply_fn_to_fields(
+            lambda x: x.broadcast_to((*shape, x.shape[-1])), custom_tensor_dims_fn=custom_tensor_dims_fn
+        )
 
     def to(self: TensorDataclassT, device) -> TensorDataclassT:
         """Returns a new TensorDataclass with the same data but on the specified device.
@@ -209,7 +257,10 @@ class TensorDataclass:
         return self._apply_fn_to_fields(lambda x: x.to(device))
 
     def _apply_fn_to_fields(
-        self: TensorDataclassT, fn: Callable, dataclass_fn: Optional[Callable] = None
+        self: TensorDataclassT,
+        fn: Callable,
+        dataclass_fn: Optional[Callable] = None,
+        custom_tensor_dims_fn: Optional[Callable] = None,
     ) -> TensorDataclassT:
         """Applies a function to all fields of the tensor dataclass.
 
@@ -231,12 +282,21 @@ class TensorDataclass:
         """
 
         new_fields = self._apply_fn_to_dict(
-            {f.name: self.__getattribute__(f.name) for f in dataclasses.fields(self)}, fn, dataclass_fn
+            {f.name: self.__getattribute__(f.name) for f in dataclasses.fields(self)},
+            fn,
+            dataclass_fn,
+            custom_tensor_dims_fn,
         )
 
         return dataclasses.replace(self, **new_fields)
 
-    def _apply_fn_to_dict(self, dict_: Dict, fn: Callable, dataclass_fn: Optional[Callable] = None) -> Dict:
+    def _apply_fn_to_dict(
+        self,
+        dict_: Dict,
+        fn: Callable,
+        dataclass_fn: Optional[Callable] = None,
+        custom_tensor_dims_fn: Optional[Callable] = None,
+    ) -> Dict:
         """A helper function for _apply_fn_to_fields, applying a function to all fields of dict_
 
         Args:
@@ -254,6 +314,14 @@ class TensorDataclass:
             if v is not None:
                 if isinstance(v, TensorDataclass) and dataclass_fn is not None:
                     new_dict[f] = dataclass_fn(v)
+                # This is the case when we have a custom dimensions tensor
+                elif (
+                    isinstance(v, torch.Tensor)
+                    and isinstance(self._field_custom_dimensions, dict)
+                    and f in self._field_custom_dimensions
+                    and custom_tensor_dims_fn is not None
+                ):
+                    new_dict[f] = custom_tensor_dims_fn(f, v)
                 elif isinstance(v, (torch.Tensor, TensorDataclass)):
                     new_dict[f] = fn(v)
                 elif isinstance(v, Dict):

scripts/downloads/download_data.py

@@ -250,9 +250,7 @@ def main(
     dataset: DatasetDownload,
 ):
     """Script to download existing datasets.
-
     We currently support the following datasets:
-
     - nerfstudio: Growing collection of real-world scenes. Use the `capture_name` argument to specify
         which capture to download.
     - blender: Blender synthetic scenes realeased with NeRF.

scripts/render.py

@@ -137,7 +137,7 @@ class RenderTrajectory:
 
         # TODO(ethan): use camera information from parsing args
         if self.traj == "spiral":
-            camera_start = pipeline.datamanager.eval_dataloader.get_camera(image_idx=0)
+            camera_start = pipeline.datamanager.eval_dataloader.get_camera(image_idx=0).flatten()
             # TODO(ethan): pass in the up direction of the camera
             camera_path = get_spiral_path(camera_start, steps=30, radius=0.1)
         elif self.traj == "filename":

tests/cameras/test_cameras.py

@@ -1,12 +1,111 @@
 """
 Test the camera classes.
 """
+import dataclasses
+from itertools import product
 
 import torch
 
 from nerfstudio.cameras.cameras import Cameras
 from nerfstudio.cameras.rays import RayBundle
 
+BATCH_SIZE = 2
+H_W = 800
+FX_Y = 10.0
+CX_Y = H_W / 2.0
+# pylint: disable=unnecessary-comprehension
+C2W_FLAT = torch.eye(4)[:3, :]
+CAMERA_TO_WORLDS = [
+    C2W_FLAT,
+    torch.stack([C2W_FLAT] * BATCH_SIZE),
+    torch.stack([torch.stack([C2W_FLAT] * BATCH_SIZE)] * BATCH_SIZE),
+]
+FX_YS = [
+    FX_Y,
+    torch.tensor([1]).float() * FX_Y,
+    torch.ones(BATCH_SIZE, 1) * FX_Y,
+    torch.ones((BATCH_SIZE, BATCH_SIZE, 1)) * FX_Y,
+]
+H_WS = [
+    None,
+    H_W,
+    torch.tensor([1]) * H_W,
+    torch.ones(BATCH_SIZE, 1).int() * H_W,
+    torch.ones((BATCH_SIZE, BATCH_SIZE, 1)).int() * H_W,
+]
+CX_YS = [
+    CX_Y,
+    torch.tensor([1]).float() * CX_Y,
+    torch.ones(BATCH_SIZE, 1) * CX_Y,
+    torch.ones((BATCH_SIZE, BATCH_SIZE, 1)) * CX_Y,
+]
+DISTORTION_PARAMS = [None, torch.zeros(6), torch.zeros((BATCH_SIZE, 6)), torch.zeros((BATCH_SIZE, BATCH_SIZE, 6))]
+camera_types = [
+    1,
+    torch.tensor([1]),
+    torch.ones(BATCH_SIZE, 1).int(),
+    torch.ones((BATCH_SIZE, BATCH_SIZE, 1)).int(),
+]
+C = Cameras(
+    CAMERA_TO_WORLDS[1],
+    FX_YS[0],
+    FX_YS[0],
+    CX_YS[0],
+    CX_YS[0],
+    H_WS[0],
+    H_WS[0],
+    DISTORTION_PARAMS[0],
+    camera_types[0],
+)
+
+C0 = Cameras(
+    CAMERA_TO_WORLDS[0],
+    FX_YS[0],
+    FX_YS[0],
+    CX_YS[0],
+    CX_YS[0],
+    H_WS[0],
+    H_WS[0],
+    DISTORTION_PARAMS[0],
+    camera_types[0],
+)
+
+C1 = Cameras(
+    CAMERA_TO_WORLDS[1],
+    FX_YS[0],
+    FX_YS[0],
+    CX_YS[0],
+    CX_YS[0],
+    H_WS[0],
+    H_WS[0],
+    DISTORTION_PARAMS[0],
+    camera_types[0],
+)
+
+C2 = Cameras(
+    CAMERA_TO_WORLDS[2],
+    FX_YS[0],
+    FX_YS[0],
+    CX_YS[0],
+    CX_YS[0],
+    H_WS[0],
+    H_WS[0],
+    DISTORTION_PARAMS[0],
+    camera_types[0],
+)
+
+C2_DIST = Cameras(
+    CAMERA_TO_WORLDS[2],
+    FX_YS[0],
+    FX_YS[0],
+    CX_YS[0],
+    CX_YS[0],
+    H_WS[0],
+    H_WS[0],
+    DISTORTION_PARAMS[1],
+    camera_types[0],
+)
+
 
 def test_pinhole_camera():
     """Test that the pinhole camera model works."""
@@ -22,5 +121,148 @@ def test_pinhole_camera():
     pinhole_camera.generate_rays(camera_indices=0, coords=coords)
 
 
+def test_camera_as_tensordataclass():
+    """Test that the camera class move to Tensordataclass works."""
+    _ = C2_DIST[torch.tensor([0]), torch.tensor([0])]
+
+    assert C0.shape == ()
+    assert C1[...].shape == torch.Size([BATCH_SIZE])
+
+    assert _check_cam_shapes(C0, ())
+    assert _check_cam_shapes(C1, (2,))
+    assert _check_cam_shapes(C2_DIST, (2, 2))
+
+    assert C0.generate_rays(0).shape == (H_W, H_W)
+    assert C0.generate_rays(0, coords=torch.ones(10, 2)).shape == (10,)
+    C1.generate_rays(0)
+    C1.generate_rays(torch.tensor([0, 1]).unsqueeze(-1))
+
+    # Make sure rays generated are same when distortion params are identity (all zeros) and None
+    assert C2.shape == C2_DIST.shape
+
+    c2_rays = C2.generate_rays(torch.tensor([0, 0]))
+    c_dist_rays = C2_DIST.generate_rays(torch.tensor([0, 0]))
+    assert _check_dataclass_allclose(c2_rays, c_dist_rays)
+    assert c2_rays.shape == (H_W, H_W)
+    assert c_dist_rays.shape == (H_W, H_W)
+
+    camera_indices = torch.tensor([[0, 0]])
+    assert C2.generate_rays(camera_indices).shape == (H_W, H_W, 1)
+
+    for args in product(
+        CAMERA_TO_WORLDS,
+        FX_YS,
+        FX_YS[:1],
+        CX_YS,
+        CX_YS[:1],
+        H_WS,
+        H_WS[:1],
+        DISTORTION_PARAMS[:-1],
+        camera_types[:-1],
+    ):
+        _ = Cameras(*args)
+
+
+def check_generate_rays_shape():
+    """Checking the output shapes from Cameras.generate_rays"""
+    coord = torch.tensor([1, 1])
+    combos = [
+        (0, None, torch.Size((800, 800))),  # First camera, all pixels
+        (
+            0,
+            coord,
+            torch.Size(()),
+        ),  # First camera, selected pixels in coords, output zero dimensional since no extra batch dim in coord
+        (
+            torch.zeros(1, 1),
+            coord.broadcast_to(1, 2),
+            torch.Size((1,)),
+        ),  # [0]th camera and selected coords, output is [1] dimensional
+        (
+            0,
+            coord.broadcast_to(1, 2),
+            torch.Size((1,)),
+        ),  # First camera and selected coords, output is [1] dimensional since one extra batch dim in coords
+        (
+            torch.zeros(1),
+            None,
+            torch.Size((800, 800)),
+        ),  # [0]th camera, all pixels, output is [HxW] dimensional since no extra output batch dim
+        (
+            torch.zeros(5, 1),
+            coord.broadcast_to(5, 2),
+            torch.Size((5,)),
+        ),  # [0]th camera and selected coords, output is [5] dimensional since extra batch dim is provided in inputs
+        (0, coord.broadcast_to(5, 2), torch.Size((5,))),  # First camera and selected coords, output is [5] dimensional
+        (
+            torch.zeros(5, 1),
+            None,
+            torch.Size((800, 800, 5)),
+        ),  # [0]th camera and all pixels, HxWx5 since coords is none and our extra batch dim in our inputs is [5]
+        (
+            torch.zeros(11, 5, 1),
+            coord.broadcast_to(11, 5, 2),
+            torch.Size((11, 5)),
+        ),  # [0]th camera and selected coords since inputs have 2 extra batch dims of [11,5]
+        (
+            0,
+            coord.broadcast_to(11, 5, 2),
+            torch.Size((11, 5)),
+        ),  # [0]th camera and selected coords since coords have 2 extra batch dims of [11, 5]
+        (
+            torch.zeros(11, 5, 1),
+            None,
+            torch.Size((800, 800, 11, 5)),
+        ),  # [0]th camera and all pixels since coords is none but we still have 2 extra batch dims of [11, 5] in coords
+    ]
+    for camera_indices, coords, output_size in combos:
+        shape = C0.generate_rays(camera_indices=camera_indices, coords=coords).shape
+        assert shape == output_size
+
+    assert C1.shape == (2,)
+    for camera_indices, coords, output_size in combos:
+        shape = C1.generate_rays(camera_indices=camera_indices, coords=coords).shape
+        assert shape == output_size
+
+    # camera_indices can't be an int anymore since our cameras object have 2 batch dimensions
+    # camera_indices last dim needs to be (..., 2) since len(cameras.shape) == 2
+    combos = [
+        (torch.zeros(2), coord, ()),
+        (torch.zeros(1, 2), coord.broadcast_to(1, 2), (1,)),
+        (torch.zeros(2), None, (800, 800)),
+        (torch.zeros(5, 2), coord.broadcast_to(5, 2), (5,)),
+        (torch.zeros(5, 2), None, (800, 800, 5)),
+        (torch.zeros(11, 5, 2), coord.broadcast_to(11, 5, 2), (11, 5)),
+        (torch.zeros(11, 5, 2), None, (800, 800, 11, 5)),
+    ]
+    for camera_indices, coords, output_size in combos:
+        shape = C2.generate_rays(camera_indices=camera_indices, coords=coords).shape
+        assert shape == output_size
+
+
+def _check_dataclass_allclose(ipt, other):
+    for field in dataclasses.fields(ipt):
+        if getattr(ipt, field.name) is not None:
+            assert torch.allclose(getattr(ipt, field.name), getattr(other, field.name))
+    return True
+
+
+def _check_cam_shapes(cam: Cameras, _batch_size):
+    if _batch_size:
+        assert len(cam) == _batch_size[0]
+    assert cam.shape == (*_batch_size,)
+    assert cam.camera_to_worlds.shape == (*_batch_size, 3, 4)
+    assert cam.fx.shape == (*_batch_size, 1)
+    assert cam.fy.shape == (*_batch_size, 1)
+    assert cam.cx.shape == (*_batch_size, 1)
+    assert cam.cy.shape == (*_batch_size, 1)
+    assert cam.height.shape == (*_batch_size, 1)
+    assert cam.width.shape == (*_batch_size, 1)
+    assert cam.distortion_params is None or cam.distortion_params.shape == (*_batch_size, 6)
+    assert cam.camera_type.shape == (*_batch_size, 1)
+    return True
+
+
 if __name__ == "__main__":
     test_pinhole_camera()
+    test_camera_as_tensordataclass()

tests/test_train.py

@@ -14,7 +14,7 @@ from nerfstudio.configs.method_configs import method_configs
 from nerfstudio.data.dataparsers.blender_dataparser import BlenderDataParserConfig
 from scripts.train import train_loop
 
-BLACKLIST = ["base", "semantic-nerfw", "instant-ngp", "nerfacto"]
+BLACKLIST = ["base", "semantic-nerfw", "instant-ngp", "nerfacto", "phototourism"]
 
 
 def set_reduced_config(config: Config):