From 7da588c5d2bfa7879bdf14485312b71ba849ba14 Mon Sep 17 00:00:00 2001 From: Georgy Ponimatkin <38923619+ponimatkin@users.noreply.github.com> Date: Sun, 26 Feb 2023 18:30:19 +0100 Subject: [PATCH] Extending ns-render to support loop rendering using interpolated cameras (#1508) * Extending ns-render to support loop rendering using interpolated cameras. * Adding docstring for render.py. Applying formating. * Expanding docstrings for class. Renaming trajectory type. * Removing mentions of loop from render.py. --- nerfstudio/cameras/camera_paths.py | 2 +- nerfstudio/cameras/camera_utils.py | 12 ++++++++---- scripts/render.py | 28 ++++++++++++++++++++++------ 3 files changed, 31 insertions(+), 11 deletions(-) diff --git a/nerfstudio/cameras/camera_paths.py b/nerfstudio/cameras/camera_paths.py index a43504694..a313cad4d 100644 --- a/nerfstudio/cameras/camera_paths.py +++ b/nerfstudio/cameras/camera_paths.py @@ -38,7 +38,7 @@ def get_interpolated_camera_path(cameras: Cameras, steps: int) -> Cameras: A new set of cameras along a path. """ Ks = cameras.get_intrinsics_matrices().cpu().numpy() - poses = cameras.camera_to_worlds().cpu().numpy() + poses = cameras.camera_to_worlds.cpu().numpy() poses, Ks = get_interpolated_poses_many(poses, Ks, steps_per_transition=steps) cameras = Cameras(fx=Ks[:, 0, 0], fy=Ks[:, 1, 1], cx=Ks[0, 0, 2], cy=Ks[0, 1, 2], camera_to_worlds=poses) diff --git a/nerfstudio/cameras/camera_utils.py b/nerfstudio/cameras/camera_utils.py index bb6430717..5e7de6179 100644 --- a/nerfstudio/cameras/camera_utils.py +++ b/nerfstudio/cameras/camera_utils.py @@ -178,7 +178,7 @@ def get_interpolated_poses(pose_a, pose_b, steps: int = 10) -> List[float]: pose = np.identity(4) pose[:3, :3] = quaternion_matrix(quat)[:3, :3] pose[:3, 3] = tran - poses_ab.append(pose) + poses_ab.append(pose[:3]) return poses_ab @@ -215,14 +215,18 @@ def get_interpolated_poses_many( tuple of new poses and intrinsics """ traj = [] - Ks = [] + k_interp = [] for idx in range(poses.shape[0] - 1): pose_a = poses[idx] pose_b = poses[idx + 1] poses_ab = get_interpolated_poses(pose_a, pose_b, steps=steps_per_transition) traj += poses_ab - Ks += get_interpolated_k(Ks[idx], Ks[idx + 1], steps_per_transition) - return torch.stack(traj, dim=0), torch.stack(Ks, dim=0) + k_interp += get_interpolated_k(Ks[idx], Ks[idx + 1], steps=steps_per_transition) + + traj = np.stack(traj, axis=0) + k_interp = np.stack(k_interp, axis=0) + + return torch.tensor(traj), torch.tensor(k_interp) def normalize(x) -> TensorType[...]: diff --git a/scripts/render.py b/scripts/render.py index 4466c5385..ce2f79dc2 100644 --- a/scripts/render.py +++ b/scripts/render.py @@ -28,7 +28,11 @@ from rich.progress import ( from torchtyping import TensorType from typing_extensions import Literal, assert_never -from nerfstudio.cameras.camera_paths import get_path_from_json, get_spiral_path +from nerfstudio.cameras.camera_paths import ( + get_interpolated_camera_path, + get_path_from_json, + get_spiral_path, +) from nerfstudio.cameras.cameras import Cameras, CameraType from nerfstudio.data.scene_box import SceneBox from nerfstudio.model_components import renderers @@ -94,7 +98,6 @@ def _render_trajectory_video( with progress: for camera_idx in progress.track(range(cameras.size), description=""): - aabb_box = None if crop_data is not None: bounding_box_min = crop_data.center - crop_data.scale / 2.0 @@ -249,14 +252,20 @@ def get_crop_from_json(camera_json: Dict[str, Any]) -> Optional[CropData]: @dataclass class RenderTrajectory: - """Load a checkpoint, render a trajectory, and save to a video file.""" + """Load a checkpoint, render a trajectory, and save to a video file. + The following trajectory options are available, + filename: Load from trajectory created using viewer or blender vfx plugin. + interpolate: Create trajectory by interpolating between eval dataset images. + spiral: Create a spiral trajectory (can be hit or miss). + """ load_config: Path """Path to config YAML file.""" rendered_output_names: List[str] = field(default_factory=lambda: ["rgb"]) """Name of the renderer outputs to use. rgb, depth, etc. concatenates them along y axis""" - traj: Literal["spiral", "filename"] = "spiral" - """Trajectory to render.""" + traj: Literal["spiral", "filename", "interpolate"] = "spiral" + """Trajectory type to render. Select between spiral-shaped trajectory, trajectory loaded from + a viewer-generated file and interpolated camera paths from the eval dataset.""" downscale_factor: int = 1 """Scaling factor to apply to the camera image resolution.""" camera_path_filename: Path = Path("camera_path.json") @@ -267,6 +276,8 @@ class RenderTrajectory: """How long the video should be.""" output_format: Literal["images", "video"] = "video" """How to save output data.""" + interpolation_steps: int = 10 + """Number of interpolation steps between eval dataset cameras.""" eval_num_rays_per_chunk: Optional[int] = None """Specifies number of rays per chunk during eval.""" @@ -275,7 +286,7 @@ class RenderTrajectory: _, pipeline, _ = eval_setup( self.load_config, eval_num_rays_per_chunk=self.eval_num_rays_per_chunk, - test_mode="test" if self.traj == "spiral" else "inference", + test_mode="test" if self.traj in ["spiral", "interpolate"] else "inference", ) install_checks.check_ffmpeg_installed() @@ -303,6 +314,11 @@ class RenderTrajectory: camera_type = CameraType.PERSPECTIVE crop_data = get_crop_from_json(camera_path) camera_path = get_path_from_json(camera_path) + elif self.traj == "interpolate": + camera_type = CameraType.PERSPECTIVE + camera_path = get_interpolated_camera_path( + cameras=pipeline.datamanager.eval_dataloader.cameras, steps=self.interpolation_steps + ) else: assert_never(self.traj)