Render specific frames, panoramic views from specific poses, or uniformly sampled from camera path

nerfstudio/scripts/camera_path_generator.py

@@ -0,0 +1,282 @@
+import json
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Literal, Optional
+
+import numpy as np
+import torch
+import tyro
+
+from nerfstudio.cameras.cameras import Cameras
+from nerfstudio.utils.eval_utils import eval_setup
+
+
+def calculate_camera_to_world_matrix(eye: torch.Tensor, target: torch.Tensor, up: torch.Tensor) -> torch.Tensor:
+ forward = (eye - target) / torch.norm(target - eye)
+ right = torch.cross((up - eye), forward)
+ right /= torch.norm(right)
+ up = torch.cross(forward, right)
+ # Create rotation matrices
+ rotation_matrices = torch.stack([right, up, forward], dim=-1)
+
+ # Combine rotation and translation matrices
+ camera_to_world_matrices = torch.eye(4) # Initialize with identity matrices
+ camera_to_world_matrices[:3, :3] = rotation_matrices
+ camera_to_world_matrices[:3, 3] = eye
+ return camera_to_world_matrices[:3, :].unsqueeze(0)
+
+
+# Function to create a rotation matrix around the y-axis
+def rotation_matrix_y(angles: torch.Tensor) -> torch.Tensor:
+ cos_angles = torch.cos(angles)
+ sin_angles = torch.sin(angles)
+ rotation_matrices = torch.zeros((angles.shape[0], 3, 3), device=angles.device)
+ rotation_matrices[:, 0, 0] = cos_angles
+ rotation_matrices[:, 0, 2] = sin_angles
+ rotation_matrices[:, 1, 1] = 1
+ rotation_matrices[:, 2, 0] = -sin_angles
+ rotation_matrices[:, 2, 2] = cos_angles
+ return rotation_matrices
+
+
+def rotate_camera_to_world_matrices(
+ camera_to_worlds: torch.Tensor, start_angle: float, end_angle: float, intermediate_steps: int, camera: Cameras
+) -> Cameras:
+ device = camera.device
+ # Generate views by rotating around the y-axis in the world coordinate system and transform these rotated targets into the NeRF coordinate system afterwards
+ start_angle = np.deg2rad(start_angle)
+ end_angle = np.deg2rad(end_angle)
+ angles = torch.linspace(start_angle, end_angle, intermediate_steps, device=device)
+ # Create rotation matrices for all angles in parallel
+ rotation_matrices = rotation_matrix_y(angles)
+ # Transpose the rotation matrices to get the clockwise rotation
+ rotation_matrices = rotation_matrices.transpose(1, 2)
+ # Number of cameras and rotations
+ num_cameras = camera_to_worlds.shape[0]
+ num_rotations = rotation_matrices.shape[0]
+
+ # Repeat the camera-to-world matrices for each rotation
+ # Shape: (num_cameras, num_rotations, 3, 4)
+ expanded_camera_matrices = camera_to_worlds.unsqueeze(1).repeat(1, num_rotations, 1, 1)
+ # Repeat the rotation matrices for each camera
+ # Shape: (num_cameras, num_rotations, 3, 3)
+ expanded_rotation_matrices = rotation_matrices.unsqueeze(0).repeat(num_cameras, 1, 1, 1)
+
+ # Apply the rotation matrices to the rotation part of the camera-to-world matrices
+ # Shape: (num_cameras, num_rotations, 3, 3)
+ rotated_matrices = torch.matmul(expanded_camera_matrices[..., :3], expanded_rotation_matrices)
+ # Combine the rotated parts with the original translation parts
+ # Shape: (num_cameras, num_rotations, 3, 4)
+ result_matrices = torch.cat((rotated_matrices, expanded_camera_matrices[..., 3:]), dim=-1)
+
+ # Reshape to the desired shape (num_cameras*num_rotations, 3, 4)
+ result_matrices = result_matrices.view(-1, 3, 4)
+ cameras = Cameras(
+ camera_to_worlds=result_matrices,
+ fx=camera.fx,
+ fy=camera.fy,
+ cx=camera.cx,
+ cy=camera.cy,
+ width=camera.width,
+ height=camera.height,
+ distortion_params=camera.distortion_params,
+ camera_type=camera.camera_type,
+ times=camera.times,
+ metadata=camera.metadata,
+ )
+ return cameras
+
+
+def get_pose_view_camera_matrix(
+ transformation_matrix: list, scale_factor: float, eye: list, target: list, device: torch.device
+) -> torch.Tensor:
+ transformation_matrix = torch.tensor(transformation_matrix, device=device)
+ transformation_matrix = torch.cat(
+ (transformation_matrix, torch.tensor([[0, 0, 0, 1]], device=device)), dim=0
+ ) # Homogeneous coordinate
+ # Apply scale to transformation matrix
+ scaled_transformation_matrix = transformation_matrix * scale_factor
+
+ eye = torch.tensor(eye, device=device)
+ target = torch.tensor(target, device=device)
+
+ # Transform the eye into NeRF coordinate system
+ eye_nerf = torch.cat((eye, torch.tensor([1.0], device=device))) # Homogeneous coordinate
+ eye_nerf = scaled_transformation_matrix @ eye_nerf
+ eye_nerf = eye_nerf[:-1] # Remove homogeneous component
+
+ target_nerf = torch.cat((target, torch.tensor([1.0], device=device))) # Homogeneous coordinate
+ target_nerf = scaled_transformation_matrix @ target_nerf
+ target_nerf = target_nerf[:-1] # Remove homogeneous component
+
+ up = torch.tensor(
+ [eye[0], -1, eye[2]], device=device
+ ) # [0, -1, 0] in world coordinate system becomes the view towards the ceiling in the camera coordinate system
+ up = torch.cat((up, torch.tensor([1.0], device=device)))
+ up_nerf = scaled_transformation_matrix @ up
+ up_nerf = up_nerf[:-1] # Remove homogeneous component
+ camera_to_world_matrix = calculate_camera_to_world_matrix(eye_nerf, target_nerf, up_nerf)
+ return camera_to_world_matrix
+
+
+@dataclass
+class GeneratePoseView:
+ """Generate a camera path by rotating around a pose and output it to a JSON file."""
+
+ load_config: Path
+ """Path to config YAML file."""
+ output_path: Path = Path("camera_path.json")
+ """Path to output camera path JSON file."""
+ pose_source: Literal["eval", "train"] = "eval"
+ """Pose source to render."""
+ viewpoints: Optional[int] = None
+ """Number of viewpoints to sample from the camera path"""
+ eye: Optional[str] = None
+ """Eye of the pose to render of shape "x y z" in the original data coordinate system."""
+ target: Optional[str] = None
+ """Target of the pose to render of shape "x y z" in the original data coordinate system. If None, they eye targets the origin in parallel to the floor."""
+ start_angle: float = 0.0
+ """Start angle of the perspective relative to the origin."""
+ end_angle: float = 345.0
+ """End angle of the perspective relative to the origin."""
+ intermediate_steps: int = 20
+ """Number of intermediate steps of the arc that will be rendered."""
+ load_dataparser_transforms: Optional[Path] = None
+ """Path to dataparser_transforms JSON file."""
+
+ def main(self) -> None:
+ """Main function."""
+ self._setup_pipeline()
+
+ if self.eye:
+ assert (
+ self.viewpoints is None
+ ), "Either use eye and target point or viewpoints sampled uniformly from the camera path to render views."
+ pose_view_cameras = self._render_eye_view()
+ elif self.viewpoints:
+ assert (
+ self.eye is None
+ ), "Either use eye and target point or viewpoints sampled uniformly from the camera path to render views."
+ assert (
+ self.target is None
+ ), "Either use eye and target point or viewpoints sampled uniformly from the camera path to render views."
+ assert (
+ self.load_dataparser_transforms is None
+ ), "Either use eye and target point or viewpoints sampled uniformly from the camera path to render views."
+ pose_view_cameras = self._render_viewpoints()
+ else:
+ raise ValueError("Either eye and target or viewpoints must be specified.")
+ assert self.start_angle <= self.end_angle, "The start angle must be less than or equal to the end angle."
+ assert self.start_angle > -360, "The start angle must not exceed -360 degrees."
+ assert self.end_angle < 360, "The end angle must not exceed 360 degrees."
+
+ camera_path = self._generate_camera_path(pose_view_cameras)
+ self._save_camera_path(camera_path)
+
+ def _setup_pipeline(self) -> None:
+ """Sets up the pipeline and necessary components."""
+ _, self.pipeline, _, _ = eval_setup(
+ self.load_config,
+ eval_num_rays_per_chunk=None,
+ test_mode="test",
+ )
+
+ if self.pose_source == "eval":
+ assert self.pipeline.datamanager.eval_dataset is not None
+ self.cameras = self.pipeline.datamanager.eval_dataset.cameras
+ else:
+ assert self.pipeline.datamanager.train_dataset is not None
+ self.cameras = self.pipeline.datamanager.train_dataset.cameras
+
+ self.camera = self.cameras[0] # Use first camera object to get intrinsics and distortion params
+
+ def _render_eye_view(self) -> None:
+ """Renders views by rotating around a specified eye and target point."""
+ assert self.load_dataparser_transforms, "dataparser_transforms.json path must be provided."
+ assert (
+ self.load_dataparser_transforms.is_file()
+ ), f"dataparser_transforms.json could not be found in {self.load_dataparser_transforms}."
+
+ with open(self.load_dataparser_transforms) as f:
+ transforms = json.load(f)
+ assert (
+ "transform" in transforms
+ ), f"Transformation matrix could not be found in {self.load_dataparser_transforms}."
+ assert "scale" in transforms, f"Scale factor could not be found in {self.load_dataparser_transforms}."
+ transform = transforms["transform"]
+ scale = transforms["scale"]
+ assert len(transform) == 3 and len(transform[0]) == 4, "Transformation matrix must be of shape [3, 4]."
+ assert isinstance(scale, float), "Scale factor must be a scalar."
+
+ eye = [float(x) for x in self.eye.split()]
+ target = [float(x) for x in self.target.split()] if self.target else [0, eye[1], 0]
+ self._validate_eye_target(eye, target)
+ pose_view_camera_matrix = get_pose_view_camera_matrix(transform, scale, eye, target, self.camera.device)
+ pose_view_cameras = rotate_camera_to_world_matrices(
+ pose_view_camera_matrix, self.start_angle, self.end_angle, self.intermediate_steps, self.camera
+ )
+ return pose_view_cameras
+
+ def _render_viewpoints(self) -> None:
+ """Renders views by rotating around sampled viewpoints."""
+ step_size = int(np.ceil(len(self.cameras) / self.viewpoints))
+ camera_indices = torch.arange(0, len(self.cameras), step_size, device=self.cameras.device)
+ pose_view_cameras = self.cameras[camera_indices]
+ pose_view_cameras = rotate_camera_to_world_matrices(
+ pose_view_cameras.camera_to_worlds, self.start_angle, self.end_angle, self.intermediate_steps, self.camera
+ )
+ return pose_view_cameras
+
+ def _validate_eye_target(self, eye, target) -> None:
+ """Validates the eye and target points."""
+ assert len(eye) == 3, "Eye must be of shape 'x y z' in the original data coordinate system."
+ assert len(target) == 3, "Target must be of shape 'x y z' in the original data coordinate system."
+
+ def _generate_camera_path(self, cameras: Cameras) -> dict:
+ """Generates a camera path from the given cameras."""
+ camera_path = {
+ "camera_type": cameras.camera_type[0].value,
+ "cameras": [],
+ "seconds": self.intermediate_steps / 24.0, # Assuming 24 FPS
+ }
+
+ for idx, camera_to_world in enumerate(cameras.camera_to_worlds):
+ camera_path["cameras"].append(
+ {
+ "camera_to_worlds": camera_to_world.cpu().numpy().tolist(),
+ "fx": cameras.fx[idx].item(),
+ "fy": cameras.fy[idx].item(),
+ "cx": cameras.cx[idx].item(),
+ "cy": cameras.cy[idx].item(),
+ "width": cameras.width[idx].item(),
+ "height": cameras.height[idx].item(),
+ "distortion_params": cameras.distortion_params[idx].cpu().numpy().tolist(),
+ "metadata": cameras.metadata[idx],
+ }
+ )
+
+ return camera_path
+
+ def _save_camera_path(self, camera_path: dict) -> None:
+ """Saves the generated camera path to a JSON file."""
+ with open(self.output_path, "w") as f:
+ json.dump(camera_path, f, indent=4)
+ print(f"Camera path saved to {self.output_path}")
+
+
+Commands = tyro.conf.FlagConversionOff[Literal["pose-view"]]
+
+
+def entrypoint():
+ """Entrypoint for use with pyproject scripts."""
+ tyro.extras.set_accent_color("bright_yellow")
+ tyro.cli(GeneratePoseView).main()
+
+
+if __name__ == "__main__":
+ entrypoint()
+
+
+def get_parser_fn():
+ """Get the parser function for the sphinx docs."""
+ return tyro.extras.get_parser(Commands) # noqa

nerfstudio/scripts/render.py

@@ -16,6 +16,7 @@
 """
 render.py
 """
+
 from __future__ import annotations
 
 import gzip

pyproject.toml

@@ -132,6 +132,7 @@ ns-train = "nerfstudio.scripts.train:entrypoint"
 ns-viewer = "nerfstudio.scripts.viewer.run_viewer:entrypoint"
 ns-eval = "nerfstudio.scripts.eval:entrypoint"
 ns-render = "nerfstudio.scripts.render:entrypoint"
+ns-generate-camera-path = "nerfstudio.scripts.camera_path_generator:entrypoint"
 ns-export = "nerfstudio.scripts.exporter:entrypoint"
 ns-dev-test = "nerfstudio.scripts.github.run_actions:entrypoint"
 ns-dev-sync-viser-message-defs = "nerfstudio.scripts.viewer.sync_viser_message_defs:entrypoint"