vendored model

sharp/utils/camera.py

@@ -0,0 +1,386 @@
+"""Contains utility functionality to render different modalities.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+from typing import Literal, NamedTuple
+
+import numpy as np
+import torch
+
+from .gaussians import Gaussians3D
+from .linalg import eyes
+
+TrajetoryType = Literal["swipe", "shake", "rotate", "rotate_forward"]
+LookAtMode = Literal["point", "ahead"]
+
+
+@dataclasses.dataclass
+class CameraInfo:
+    """Camera info for a pinhole camera."""
+
+    intrinsics: torch.Tensor
+    extrinsics: torch.Tensor
+    width: int
+    height: int
+
+
+class FocusRange(NamedTuple):
+    """Parametrizes a range of depth / disparity values."""
+
+    min: float
+    focus: float
+    max: float
+
+
+@dataclasses.dataclass
+class TrajectoryParams:
+    """Parameters for trajectory."""
+
+    type: TrajetoryType = "rotate_forward"
+    lookat_mode: LookAtMode = "point"
+    max_disparity: float = 0.08
+    max_zoom: float = 0.15
+    distance_m: float = 0.0
+    num_steps: int = 60
+    num_repeats: int = 1
+
+
+def compute_max_offset(
+    scene: Gaussians3D,
+    params: TrajectoryParams,
+    resolution_px: tuple[int, int],
+    f_px: float,
+) -> np.ndarray:
+    """Compute the maximum offset for camera along X/Y/Z axis."""
+    scene_points = scene.mean_vectors
+    extrinsics = torch.eye(4).to(scene_points.device)
+    min_depth, _, _ = _compute_depth_quantiles(scene_points, extrinsics)
+
+    r_px = resolution_px
+    diagonal = np.sqrt((r_px[0] / f_px) ** 2 + (r_px[1] / f_px) ** 2)
+    max_lateral_offset_m = params.max_disparity * diagonal * min_depth
+
+    max_medial_offset_m = params.max_zoom * min_depth
+    max_offset_xyz_m = np.array([max_lateral_offset_m, max_lateral_offset_m, max_medial_offset_m])
+
+    return max_offset_xyz_m
+
+
+def create_eye_trajectory(
+    scene: Gaussians3D,
+    params: TrajectoryParams,
+    resolution_px: tuple[int, int],
+    f_px: float,
+) -> list[torch.Tensor]:
+    """Create eye trajectory for trajectory type."""
+    max_offset_xyz_m = compute_max_offset(
+        scene,
+        params,
+        resolution_px,
+        f_px,
+    )
+    # We place the eye trajectory at z=distance plane (default=0),
+    # assuming portal plane is placed at z=natural_distance.
+    if params.type == "swipe":
+        return create_eye_trajectory_swipe(
+            max_offset_xyz_m, params.distance_m, params.num_steps, params.num_repeats
+        )
+    elif params.type == "shake":
+        return create_eye_trajectory_shake(
+            max_offset_xyz_m, params.distance_m, params.num_steps, params.num_repeats
+        )
+    elif params.type == "rotate":
+        return create_eye_trajectory_rotate(
+            max_offset_xyz_m, params.distance_m, params.num_steps, params.num_repeats
+        )
+    elif params.type == "rotate_forward":
+        return create_eye_trajectory_rotate_forward(
+            max_offset_xyz_m, params.distance_m, params.num_steps, params.num_repeats
+        )
+    else:
+        raise ValueError(f"Invalid trajectory type {params.type}.")
+
+
+def create_eye_trajectory_swipe(
+    offset_xyz_m: np.ndarray,
+    distance_m: float,
+    num_steps: int,
+    num_repeats: int,
+) -> list[torch.Tensor]:
+    """Create a left to right swipe trajectory."""
+    offset_x_m, _, _ = offset_xyz_m
+    eye_positions = [
+        torch.tensor([x, 0, distance_m], dtype=torch.float32)
+        for x in np.linspace(-offset_x_m, offset_x_m, num_steps)
+    ]
+    return eye_positions * num_repeats
+
+
+def create_eye_trajectory_shake(
+    offset_xyz_m: np.ndarray,
+    distance_m: float,
+    num_steps: int,
+    num_repeats: int,
+) -> list[torch.Tensor]:
+    """Create a left right shake followed by an up down shake trajectory."""
+    num_steps_total = num_steps * num_repeats
+    num_steps_horizontal = num_steps_total // 2
+    num_steps_vertical = num_steps_total - num_steps_horizontal
+
+    offset_x_m, offset_y_m, _ = offset_xyz_m
+    eye_positions: list[torch.Tensor] = []
+    eye_positions.extend(
+        torch.tensor(
+            [offset_x_m * np.sin(2 * np.pi * t), 0.0, distance_m],
+            dtype=torch.float32,
+        )
+        for t in np.linspace(0, num_repeats, num_steps_horizontal)
+    )
+    eye_positions.extend(
+        torch.tensor(
+            [0.0, offset_y_m * np.sin(2 * np.pi * t), distance_m],
+            dtype=torch.float32,
+        )
+        for t in np.linspace(0, num_repeats, num_steps_vertical)
+    )
+
+    return eye_positions
+
+
+def create_eye_trajectory_rotate(
+    offset_xyz_m: np.ndarray,
+    distance_m: float,
+    num_steps: int,
+    num_repeats: int,
+) -> list[torch.Tensor]:
+    """Create a rotating trajectory."""
+    num_steps_total = num_steps * num_repeats
+    offset_x_m, offset_y_m, _ = offset_xyz_m
+    eye_positions = [
+        torch.tensor(
+            [
+                offset_x_m * np.sin(2 * np.pi * t),
+                offset_y_m * np.cos(2 * np.pi * t),
+                distance_m,
+            ],
+            dtype=torch.float32,
+        )
+        for t in np.linspace(0, num_repeats, num_steps_total)
+    ]
+
+    return eye_positions
+
+
+def create_eye_trajectory_rotate_forward(
+    offset_xyz_m: np.ndarray,
+    distance_m: float,
+    num_steps: int,
+    num_repeats: int,
+) -> list[torch.Tensor]:
+    """Create a rotating trajectory."""
+    num_steps_total = num_steps * num_repeats
+    offset_x_m, _, offset_z_m = offset_xyz_m
+    eye_positions = [
+        torch.tensor(
+            [
+                offset_x_m * np.sin(2 * np.pi * t),
+                0.0,
+                distance_m + offset_z_m * (1.0 - np.cos(2 * np.pi * t)) / 2,
+            ],
+            dtype=torch.float32,
+        )
+        for t in np.linspace(0, num_repeats, num_steps_total)
+    ]
+
+    return eye_positions
+
+
+def create_camera_model(
+    scene: Gaussians3D,
+    intrinsics: torch.Tensor,
+    resolution_px: tuple[int, int],
+    lookat_mode: LookAtMode = "point",
+) -> PinholeCameraModel:
+    """Create camera model to simulate general pinhole camera."""
+    screen_extrinsics = torch.eye(4)
+    screen_intrinsics = intrinsics.clone()
+
+    image_width, image_height = resolution_px
+    screen_resolution_px = get_screen_resolution_px_from_input(
+        width=image_width, height=image_height
+    )
+
+    screen_intrinsics[0] *= screen_resolution_px[0] / image_width
+    screen_intrinsics[1] *= screen_resolution_px[1] / image_height
+
+    camera_model = PinholeCameraModel(
+        scene,
+        screen_extrinsics=screen_extrinsics,
+        screen_intrinsics=screen_intrinsics,
+        screen_resolution_px=screen_resolution_px,
+        focus_depth_quantile=0.1,
+        min_depth_focus=2.0,
+        lookat_mode=lookat_mode,
+    )
+    return camera_model
+
+
+def create_camera_matrix(
+    position: torch.Tensor,
+    look_at_position: torch.Tensor | None = None,
+    world_up: torch.Tensor | None = None,
+    inverse: bool = False,
+) -> torch.Tensor:
+    """Create camera matrix from vectors."""
+    device = position.device
+
+    if look_at_position is None:
+        look_at_position = torch.zeros(3, device=device)
+    if world_up is None:
+        world_up = torch.tensor([0.0, 0.0, 1.0], device=device)
+
+    position, look_at_position, world_up = torch.broadcast_tensors(
+        position, look_at_position, world_up
+    )
+
+    camera_front = look_at_position - position
+    camera_front = camera_front / camera_front.norm(dim=-1, keepdim=True)
+
+    camera_right = torch.cross(camera_front, world_up, dim=-1)
+    camera_right = camera_right / camera_right.norm(dim=-1, keepdim=True)
+
+    camera_down = torch.cross(camera_front, camera_right, dim=-1)
+    rotation_matrix = torch.stack([camera_right, camera_down, camera_front], dim=-1)
+
+    matrix = eyes(dim=4, shape=position.shape[:-1], device=device)
+    if inverse:
+        matrix[..., :3, :3] = rotation_matrix.transpose(-1, -2)
+        matrix[..., :3, 3:4] = -rotation_matrix.transpose(-1, -2) @ position[..., None]
+    else:
+        matrix[..., :3, :3] = rotation_matrix
+        matrix[..., :3, 3] = position
+
+    return matrix
+
+
+class PinholeCameraModel:
+    """Camera model that focuses on point."""
+
+    def __init__(
+        self,
+        scene: Gaussians3D,
+        screen_extrinsics: torch.Tensor,
+        screen_intrinsics: torch.Tensor,
+        screen_resolution_px: tuple[int, int],
+        focus_depth_quantile: float = 0.1,
+        min_depth_focus: float = 2.0,
+        lookat_point: tuple[float, float, float] | None = None,
+        lookat_mode: LookAtMode = "point",
+    ) -> None:
+        """Initialize GeneralPinholeCameraModel.
+
+        Args:
+            scene: The scene to display.
+            screen_extrinsics: Extrinsics of the default position.
+            screen_intrinsics: Intrinsics to use for rendering.
+            screen_resolution_px: Width and height to render.
+            focus_depth_quantile: Where inside the depth range to focus on.
+            min_depth_focus: Depth to focus at.
+            lookat_point: a point that the camera's Z axis directs towards.
+            lookat_mode: "point" to look at a fixed point,
+                "ahead" to look straight ahead.
+        """
+        self.scene = scene
+        self.screen_extrinsics = screen_extrinsics
+        self.screen_intrinsics = screen_intrinsics
+        self.screen_resolution_px = screen_resolution_px
+
+        self.focus_depth_quantile = focus_depth_quantile
+        self.min_depth_focus = min_depth_focus
+        self.lookat_point = lookat_point
+        self.lookat_mode = lookat_mode
+
+        scene_points = scene.mean_vectors
+        if scene_points.ndim == 3:
+            scene_points = scene_points[0]
+        elif scene_points.ndim != 2:
+            raise ValueError("Unsupported dimensionality of scene points.")
+        self._scene_points = scene_points.cpu()
+
+        self.depth_quantiles = _compute_depth_quantiles(
+            self._scene_points,
+            self.screen_extrinsics,
+            q_focus=self.focus_depth_quantile,
+        )
+
+    def compute(self, eye_pos: torch.Tensor) -> CameraInfo:
+        """Compute camera for eye position."""
+        extrinsics = self.screen_extrinsics.clone()
+
+        origin = eye_pos if self.lookat_mode == "ahead" else torch.zeros(3)
+
+        if self.lookat_point is None:
+            depth_focus = max(self.min_depth_focus, self.depth_quantiles.focus)
+            look_at_position = origin + torch.tensor([0.0, 0.0, depth_focus])
+        else:
+            look_at_position = origin + torch.tensor([*self.lookat_point])
+
+        world_up = torch.tensor([0.0, -1.0, 0.0])
+        extrinsics_modifier = create_camera_matrix(
+            eye_pos, look_at_position, world_up, inverse=True
+        )
+        extrinsics = extrinsics_modifier @ self.screen_extrinsics
+
+        camera_info = CameraInfo(
+            intrinsics=self.screen_intrinsics,
+            extrinsics=extrinsics,
+            width=self.screen_resolution_px[0],
+            height=self.screen_resolution_px[1],
+        )
+        return camera_info
+
+    def set_screen_extrinsics(self, new_value: torch.Tensor) -> None:
+        """Modify the default extrinsics."""
+        self.screen_extrinsics = new_value
+        self.depth_quantiles = _compute_depth_quantiles(self._scene_points, self.screen_extrinsics)
+
+
+def get_screen_resolution_px_from_input(width: int, height: int) -> tuple[int, int]:
+    """Get resolution for metadata dictionary."""
+    resolution_px = (width, height)
+    # halve the dimensions for super large image
+    if resolution_px[1] > 3000:
+        resolution_px = (resolution_px[0] // 2, resolution_px[1] // 2)
+    # for mp4 compatibility, enforce dimensions to even number,
+    # otherwise could not be played in browser
+    if resolution_px[0] % 2 != 0:
+        resolution_px = (resolution_px[0] + 1, resolution_px[1])
+    if resolution_px[1] % 2 != 0:
+        resolution_px = (resolution_px[0], resolution_px[1] + 1)
+    return resolution_px
+
+
+def _compute_depth_quantiles(
+    points: torch.Tensor,
+    extrinsics: torch.Tensor,
+    q_near: float = 0.001,
+    q_focus: float = 0.1,
+    q_far: float = 0.999,
+) -> FocusRange:
+    """Compute disparity quantiles for scene and extrinsics id."""
+    points_local = points @ extrinsics[:3, :3].T + extrinsics[:3, 3]
+    depth_values = points_local[..., 2].flatten()
+    depth_values = depth_values[depth_values > 0]
+    q_values = torch.tensor([q_near, q_focus, q_far])
+    depth_quantiles_pt = torch.quantile(depth_values.cpu(), q_values)
+    depth_quantiles = FocusRange(
+        min=float(depth_quantiles_pt[0]),
+        focus=float(depth_quantiles_pt[1]),
+        max=float(depth_quantiles_pt[2]),
+    )
+    return depth_quantiles