iPhone (Record3D) Dataset (#235)

* init branch commit.

* intrinsics done.

* created dataset for record3d.

* clean up and documentation.

* cleanup.

* fixed orientation.

nerfactory/dataloaders/datasets.py

@@ -24,6 +24,7 @@ from typing import Optional, Union
 import imageio
 import numpy as np
 import torch
+from scipy.spatial.transform import Rotation
 
 from nerfactory.dataloaders.colmap_utils import (
     read_cameras_binary,
@@ -277,7 +278,7 @@ class Mipnerf360(Dataset):
     aabb_scale = 4
 
     @classmethod
-    def _normalize_orientation(cls, poses: np.ndarray):
+    def normalize_orientation(cls, poses: np.ndarray):
         """Set the _up_ direction to be in the positive Y direction.
 
         Args:
@@ -341,7 +342,7 @@ class Mipnerf360(Dataset):
         poses = np.concatenate([poses[:, :, 1:2], -poses[:, :, 0:1], poses[:, :, 2:]], axis=-1)
 
         # Center poses and rotate. (Compute up from average of all poses)
-        poses = self._normalize_orientation(poses)
+        poses = self.normalize_orientation(poses)
 
         # Scale factor used in mipnerf
         if self.auto_scale:
@@ -376,6 +377,110 @@ class Mipnerf360(Dataset):
         return dataset_inputs
 
 
+@dataclass
+class Record3D(Dataset):
+    """Record3D Dataset
+
+    Args:
+        data_directory: Location of data
+        downscale_factor: How much to downscale images. Defaults to 1.
+        val_skip: 1/val_skip images to use for validation. Defaults to 8.
+        aabb_scale: Scene scale, Defaults to 4.0.
+    """
+
+    data_directory: str
+    downscale_factor: int = 1
+    val_skip: int = 8
+    aabb_scale = 4.0
+
+    def _generate_dataset_inputs(self, split: str = "train") -> DatasetInputs:
+        abs_dir = get_absolute_path(self.data_directory)
+
+        image_dir = os.path.join(abs_dir, "rgb")
+
+        if not os.path.exists(image_dir):
+            raise ValueError(f"Image directory {image_dir} doesn't exist")
+
+        ext = ".jpg"
+        image_filenames = []
+        for f in os.listdir(image_dir):
+            image_filenames.append(os.path.join(image_dir, f))
+        image_filenames = sorted(image_filenames, key=lambda fn: int(os.path.basename(fn)[: -len(ext)]))
+        num_images = len(image_filenames)
+
+        metadata_path = os.path.join(abs_dir, "metadata.json")
+        metadata_dict = load_from_json(metadata_path)
+
+        poses_data = np.array(metadata_dict["poses"])
+        # (N, 3, 4)
+        poses = np.concatenate(
+            [Rotation.from_quat(poses_data[:, :4]).as_matrix(), poses_data[:, 4:, None]],
+            axis=-1,
+        ).astype(np.float32)
+
+        # Normalization similar to Mipnerf360
+        poses = Mipnerf360.normalize_orientation(poses)
+
+        bottom = np.reshape([0, 0, 0, 1.0], [1, 4])
+        bottom = np.tile(np.reshape(bottom, [1, 1, 4]), [poses.shape[0], 1, 1])
+        poses = np.concatenate([poses[:, :3, :4], bottom], -2).astype(np.float32)
+
+        rotation_matrix = np.array(
+            [
+                [1.0, 0.0, 0.0, 0.0],
+                [0.0, 0.0, -1.0, 0.0],
+                [0.0, 1.0, 0.0, 0.0],
+                [0.0, 0.0, 0.0, 1.0],
+            ],
+            dtype=np.float32,
+        )
+        poses = rotation_matrix @ poses
+
+        idx_test = np.arange(num_images)[:: self.val_skip]
+        idx_train = np.array([i for i in np.arange(num_images) if i not in idx_test])
+        idx = idx_train if split == "train" else idx_test
+        if num_images != poses.shape[0]:
+            raise RuntimeError(f"Different number of images ({num_images}), and poses ({poses.shape[0]})")
+
+        image_filenames = np.array(image_filenames)[idx]
+        poses = poses[idx]
+
+        # Centering poses
+        poses[:, :3, 3] = poses[:, :3, 3] - np.mean(poses[:, :3, 3], axis=0)
+
+        camera_to_world = torch.from_numpy(poses[:, :3, :4])  # camera to world transform
+
+        # Camera intrinsics
+        K = np.array(metadata_dict["K"]).reshape((3, 3)).T
+        focal_length = K[0, 0]
+
+        H = metadata_dict["h"]
+        W = metadata_dict["w"]
+
+        # TODO(akristoffersen): The metadata dict comes with principle points,
+        # but caused errors in image coord indexing. Should update once that is fixed.
+        cx, cy = W / 2, H / 2
+
+        num_cameras = len(image_filenames)
+        num_intrinsics_params = 3
+        intrinsics = torch.ones((num_cameras, num_intrinsics_params), dtype=torch.float32)
+        intrinsics *= torch.tensor([cx, cy, focal_length])
+
+        # scene_bounds = SceneBounds.from_camera_poses(camera_to_world, self.aabb_scale)
+        aabb = torch.tensor([[-1, -1, -1], [1, 1, 1]], dtype=torch.float32) * self.aabb_scale
+        scene_bounds = SceneBounds(aabb=aabb)
+
+        dataset_inputs = DatasetInputs(
+            image_filenames=image_filenames,
+            downscale_factor=1,
+            intrinsics=intrinsics,
+            camera_to_world=camera_to_world,
+            scene_bounds=scene_bounds,
+        )
+
+        return dataset_inputs
+
+
 @dataclass
 class Friends(Dataset):
     """Friends Dataset

nerfactory/dataloaders/structs.py

@@ -99,6 +99,13 @@ class SceneBounds:
         aabb = torch.tensor([json_[0], json_[1]])
         return SceneBounds(aabb=aabb)
 
+    @staticmethod
+    def from_camera_poses(poses: TensorType[..., 3, 4], scale_factor: float) -> "SceneBounds":
+        """Returns the instance of SceneBounds that fully envelopes a set of poses"""
+        xyzs = poses[..., :3, -1]
+        aabb = torch.stack([torch.min(xyzs, dim=0)[0], torch.max(xyzs, dim=0)[0]])
+        return SceneBounds(aabb=aabb * scale_factor)
+
 
 @dataclass
 class DatasetInputs:
@@ -106,12 +113,14 @@ class DatasetInputs:
 
     Args:
         image_filenames: Filenames for the images.
+        intrinsics: Tensor of per-image camera intrisics.
+        camera_to_world: Tensor of per-image c2w matrices, in [R | t] format.
         ...
     """
 
     image_filenames: List[str]
-    intrinsics: torch.Tensor
-    camera_to_world: torch.Tensor
+    intrinsics: TensorType["num_cameras", "num_intrinsics_params"]
+    camera_to_world: TensorType["num_cameras", 3, 4]
     downscale_factor: int = 1
     mask_filenames: Optional[List[str]] = None
     depth_filenames: Optional[List[str]] = None

scripts/run_eval.py

@@ -149,7 +149,7 @@ def main(
     Args:
         config_name: Name of the config file to use.
         checkpoint_dir: Directory to load the checkpoint from.
-        rendered_output_name: Name of the renderer output to use.
+        rendered_output_name: Name of the renderer output to use (rgb, depth, etc.).
         method: Method to use for evaluation. PSNR computes metrics, TRAJ renders a trajectory.
         traj: Trajectory to render.
         output_filename: Name of the output file.