From b471736a61ca74fc274fdb9dd4e750f694b1685e Mon Sep 17 00:00:00 2001 From: Alexander Kristoffersen Date: Tue, 16 Aug 2022 23:24:49 -0700 Subject: [PATCH] iPhone (Record3D) Dataset (#235) * init branch commit. * intrinsics done. * created dataset for record3d. * clean up and documentation. * cleanup. * fixed orientation. --- nerfactory/dataloaders/datasets.py | 109 ++++++++++++++++++++++++++++- nerfactory/dataloaders/structs.py | 13 +++- scripts/run_eval.py | 2 +- 3 files changed, 119 insertions(+), 5 deletions(-) diff --git a/nerfactory/dataloaders/datasets.py b/nerfactory/dataloaders/datasets.py index 711c35faa..b1762e086 100644 --- a/nerfactory/dataloaders/datasets.py +++ b/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 diff --git a/nerfactory/dataloaders/structs.py b/nerfactory/dataloaders/structs.py index efdcec2a1..af17a605a 100644 --- a/nerfactory/dataloaders/structs.py +++ b/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 diff --git a/scripts/run_eval.py b/scripts/run_eval.py index 12e52b6cf..689a081c3 100644 --- a/scripts/run_eval.py +++ b/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.