vendored model

sharp/utils/__init__.py

@@ -0,0 +1,5 @@
+"""Contains utils packages.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""

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

sharp/utils/color_space.py

@@ -0,0 +1,88 @@
+"""Contains color space utility functions.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Literal
+
+import torch
+
+from sharp.utils.robust import robust_where
+
+LOGGER = logging.getLogger(__name__)
+
+ColorSpace = Literal["sRGB", "linearRGB"]
+
+
+def encode_color_space(color_space: ColorSpace) -> int:
+    """Encode color space to integer."""
+    return 0 if color_space == "sRGB" else 1
+
+
+def decode_color_space(color_space_index: int) -> ColorSpace:
+    """Decode color space index to color space."""
+    return "sRGB" if color_space_index == 0 else "linearRGB"
+
+
+def sRGB2linearRGB(sRGB: torch.Tensor) -> torch.Tensor:
+    """SRGB to linearRGB conversion function.
+
+    Reference:
+    https://developer.apple.com/metal/Metal-Shading-Language-Specification.pdf
+    Section 7.7.7
+
+    Args:
+        sRGB: Input image tensor in sRGB space.
+    """
+    # We need to use robust_where to clamp the second branch.
+    # Otherwise, torch.where will lead to NaN in the backward pass, see
+    # https://github.com/pytorch/pytorch/issues/68425
+    THRESHOLD = 0.04045
+
+    def branch_true_func(x):
+        return x / 12.92
+
+    def branch_false_func(x):
+        return ((x + 0.055) / 1.055) ** 2.4
+
+    return robust_where(
+        sRGB <= THRESHOLD,
+        sRGB,
+        branch_true_func,
+        branch_false_func,
+        branch_false_safe_value=THRESHOLD,
+    )
+
+
+def linearRGB2sRGB(linearRGB: torch.Tensor) -> torch.Tensor:
+    """LinearRGB to sRGB conversion function.
+
+    Reference:
+    https://developer.apple.com/metal/Metal-Shading-Language-Specification.pdf
+    Section 7.7.7
+
+    Args:
+        linearRGB: Input image tensor in linearRGB space.
+    """
+    # We need to use robust_where to clamp the second branch.
+    # Otherwise, torch.where will lead to NaN in the backward pass, see
+    # https://github.com/pytorch/pytorch/issues/68425
+    THRESHOLD = 0.0031308
+
+    def branch_true_func(x):
+        return x * 12.92
+
+    def branch_false_func(x):
+        return 1.055 * (x ** (1 / 2.4)) - 0.055
+
+    return robust_where(
+        linearRGB <= THRESHOLD,
+        linearRGB,
+        branch_true_func,
+        branch_false_func,
+        branch_false_safe_value=THRESHOLD,
+    )

sharp/utils/gaussians.py

@@ -0,0 +1,483 @@
+"""Contains basic data structures and functionality for 3D Gaussians.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+from typing import Any, Literal, NamedTuple
+
+import numpy as np
+import torch
+from plyfile import PlyData, PlyElement
+
+from sharp.utils import color_space as cs_utils
+from sharp.utils import linalg
+
+LOGGER = logging.getLogger(__name__)
+
+
+BackgroundColor = Literal["black", "white", "random_color", "random_pixel"]
+
+
+class Gaussians3D(NamedTuple):
+    """Represents a collection of 3D Gaussians."""
+
+    mean_vectors: torch.Tensor
+    singular_values: torch.Tensor
+    quaternions: torch.Tensor
+    colors: torch.Tensor
+    opacities: torch.Tensor
+
+    def to(self, device: torch.device) -> Gaussians3D:
+        """Move Gaussians to device."""
+        return Gaussians3D(
+            mean_vectors=self.mean_vectors.to(device),
+            singular_values=self.singular_values.to(device),
+            quaternions=self.quaternions.to(device),
+            colors=self.colors.to(device),
+            opacities=self.opacities.to(device),
+        )
+
+
+class SceneMetaData(NamedTuple):
+    """Meta data about Gaussian scene."""
+
+    focal_length_px: float
+    resolution_px: tuple[int, int]
+    color_space: cs_utils.ColorSpace
+
+
+def get_unprojection_matrix(
+    extrinsics: torch.Tensor,
+    intrinsics: torch.Tensor,
+    image_shape: tuple[int, int],
+) -> torch.Tensor:
+    """Compute unprojection matrix to transform Gaussians to Euclidean space.
+
+    Args:
+        extrinsics: The 4x4 extrinsics matrix of the camera view.
+        intrinsics: The 4x4 intrinsics matrix of the camera view.
+        image_shape: The (width, height) of the input image.
+
+    Returns:
+        A 4x4 matrix to transform Gaussians from NDC space to Euclidean space.
+    """
+    device = intrinsics.device
+    image_width, image_height = image_shape
+    # This matrix converts OpenCV pixel coordinates to NDC coordinates where
+    # (-1, 1) denotes the top left and (1, 1) the bottom right of the image.
+    #
+    # Note that premultiplying the intrinsics with ndc_matrix typically yields a matrix
+    # that simply scales the x-axis by 2 * focal_length / image_width and the y-axis by
+    # 2 * focal_length / image_height.
+    ndc_matrix = torch.tensor(
+        [
+            [2.0 / image_width, 0.0, -1.0, 0.0],
+            [0.0, 2.0 / image_height, -1.0, 0.0],
+            [0.0, 0.0, 1.0, 0.0],
+            [0.0, 0.0, 0.0, 1.0],
+        ],
+        device=device,
+    )
+    return torch.linalg.inv(ndc_matrix @ intrinsics @ extrinsics)
+
+
+def unproject_gaussians(
+    gaussians_ndc: Gaussians3D,
+    extrinsics: torch.Tensor,
+    intrinsics: torch.Tensor,
+    image_shape: tuple[int, int],
+) -> Gaussians3D:
+    """Unproject Gaussians from NDC space to world coordinates."""
+    unprojection_matrix = get_unprojection_matrix(extrinsics, intrinsics, image_shape)
+    gaussians = apply_transform(gaussians_ndc, unprojection_matrix[:3])
+    return gaussians
+
+
+def apply_transform(gaussians: Gaussians3D, transform: torch.Tensor) -> Gaussians3D:
+    """Apply an affine transformation to 3D Gaussians.
+
+    Args:
+        gaussians: The Gaussians to transform.
+        transform: An affine transform with shape 3x4.
+
+    Returns:
+        The transformed Gaussians.
+
+    Note: This operation is not differentiable.
+    """
+    transform_linear = transform[..., :3, :3]
+    transform_offset = transform[..., :3, 3]
+
+    mean_vectors = gaussians.mean_vectors @ transform_linear.T + transform_offset
+    covariance_matrices = compose_covariance_matrices(
+        gaussians.quaternions, gaussians.singular_values
+    )
+    covariance_matrices = (
+        transform_linear @ covariance_matrices @ transform_linear.transpose(-1, -2)
+    )
+    quaternions, singular_values = decompose_covariance_matrices(covariance_matrices)
+
+    return Gaussians3D(
+        mean_vectors=mean_vectors,
+        singular_values=singular_values,
+        quaternions=quaternions,
+        colors=gaussians.colors,
+        opacities=gaussians.opacities,
+    )
+
+
+def decompose_covariance_matrices(
+    covariance_matrices: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Decompose 3D covariance matrices into quaternions and singular values.
+
+    Args:
+        covariance_matrices: The covariance matrices to decompose.
+
+    Returns:
+        Quaternion and singular values corresponding to the orientation and scales of
+        the diagonalized matrix.
+
+    Note: This operation is not differentiable.
+    """
+    device = covariance_matrices.device
+    dtype = covariance_matrices.dtype
+
+    # We convert to fp64 to avoid numerical errors.
+    covariance_matrices = covariance_matrices.detach().cpu().to(torch.float64)
+    rotations, singular_values_2, _ = torch.linalg.svd(covariance_matrices)
+
+    # NOTE: in SVD, it is possible that U and VT are both reflections.
+    # We need to correct them.
+    batch_idx, gaussian_idx = torch.where(torch.linalg.det(rotations) < 0)
+    num_reflections = len(gaussian_idx)
+    if num_reflections > 0:
+        LOGGER.warning(
+            "Received %d reflection matrices from SVD. Flipping them to rotations.",
+            num_reflections,
+        )
+        # Flip the last column of reflection and make it a rotation.
+        rotations[batch_idx, gaussian_idx, :, -1] *= -1
+    quaternions = linalg.quaternions_from_rotation_matrices(rotations)
+    quaternions = quaternions.to(dtype=dtype, device=device)
+    singular_values = singular_values_2.sqrt().to(dtype=dtype, device=device)
+    return quaternions, singular_values
+
+
+def compose_covariance_matrices(
+    quaternions: torch.Tensor, singular_values: torch.Tensor
+) -> torch.Tensor:
+    """Compose 3D covariance matrices into quaternions and singular values.
+
+    Args:
+        quaternions: The quaternions describing the principal basis.
+        singular_values: The scales of the diagonalized matrix.
+
+    Returns:
+        The 3x3 covariances matrices.
+    """
+    device = quaternions.device
+    rotations = linalg.rotation_matrices_from_quaternions(quaternions)
+    diagonal_matrix = torch.eye(3, device=device) * singular_values[..., :, None]
+    return rotations @ diagonal_matrix.square() @ rotations.transpose(-1, -2)
+
+
+def convert_spherical_harmonics_to_rgb(sh0: torch.Tensor) -> torch.Tensor:
+    """Convert degree-0 spherical harmonics to RGB.
+
+    Reference:
+        https://en.wikipedia.org/wiki/Table_of_spherical_harmonics
+    """
+    coeff_degree0 = np.sqrt(1.0 / (4.0 * np.pi))
+    return sh0 * coeff_degree0 + 0.5
+
+
+def convert_rgb_to_spherical_harmonics(rgb: torch.Tensor) -> torch.Tensor:
+    """Convert RGB to degree-0 spherical harmonics.
+
+    Reference:
+        https://en.wikipedia.org/wiki/Table_of_spherical_harmonics
+    """
+    coeff_degree0 = np.sqrt(1.0 / (4.0 * np.pi))
+    return (rgb - 0.5) / coeff_degree0
+
+
+def load_ply(path: Path) -> tuple[Gaussians3D, SceneMetaData]:
+    """Loads a ply from a file."""
+    plydata = PlyData.read(path)
+
+    vertices = next(filter(lambda x: x.name == "vertex", plydata.elements))
+
+    properties = ["x", "y", "z"]
+    properties.extend([f"f_dc_{i}" for i in range(3)])
+    properties.extend([f"scale_{i}" for i in range(3)])
+    properties.extend([f"rot_{i}" for i in range(3)])
+
+    for prop in properties:
+        if prop not in vertices:
+            raise KeyError(f"Incompatible ply file: property {prop} not found in ply elements.")
+    mean_vectors = np.stack(
+        (
+            np.asarray(vertices["x"]),
+            np.asarray(vertices["y"]),
+            np.asarray(vertices["z"]),
+        ),
+        axis=1,
+    )
+
+    scale_logits = np.stack(
+        (
+            np.asarray(vertices["scale_0"]),
+            np.asarray(vertices["scale_1"]),
+            np.asarray(vertices["scale_2"]),
+        ),
+        axis=1,
+    )
+
+    quaternions = np.stack(
+        (
+            np.asarray(vertices["rot_0"]),
+            np.asarray(vertices["rot_1"]),
+            np.asarray(vertices["rot_2"]),
+            np.asarray(vertices["rot_3"]),
+        ),
+        axis=1,
+    )
+
+    spherical_harmonics_deg0 = np.stack(
+        (
+            np.asarray(vertices["f_dc_0"]),
+            np.asarray(vertices["f_dc_1"]),
+            np.asarray(vertices["f_dc_2"]),
+        ),
+        axis=1,
+    )
+
+    colors = convert_spherical_harmonics_to_rgb(spherical_harmonics_deg0)
+
+    opacity_logits = np.asarray(vertices["opacity"])[..., None]
+
+    supplement_elements = [element for element in plydata.elements if element.name != "vertex"]
+    supplement_data: dict[str, Any] = {}
+    supplement_keys = ["extrinsic", "intrinsic", "color_space", "image_size"]
+
+    for element in supplement_elements:
+        for key in supplement_keys:
+            if key not in supplement_data and key in element:
+                supplement_data[key] = np.asarray(element[key])
+
+    # Parse intrinsics and image_size.
+    if "intrinsic" in supplement_data:
+        intrinsics_data = supplement_data["intrinsic"]
+
+        # Legacy: image_size is contained in intrinsic element.
+        if "image_size" not in supplement_data:
+            if len(intrinsics_data) != 4:
+                raise ValueError(
+                    "Expect legacy intrinsics with len=4 containing image size, "
+                    f"but received len={len(intrinsics_data)}"
+                )
+            focal_length_px = (intrinsics_data[0], intrinsics_data[1])
+            width = int(intrinsics_data[2])
+            height = int(intrinsics_data[3])
+
+        else:
+            if len(intrinsics_data) != 9:
+                raise ValueError(
+                    "Expect 9 elements in intrinsics, " f"but received {len(intrinsics_data)}."
+                )
+            intrinsics_matrix = intrinsics_data.reshape((3, 3))
+            focal_length_px = (intrinsics_matrix[0, 0], intrinsics_matrix[1, 1])
+
+            image_size_data = supplement_data["image_size"]
+            width = image_size_data[0]
+            height = image_size_data[1]
+
+    # Default to VGA resolution: focal length = 512, image size = (640, 480).
+    else:
+        focal_length_px = (512, 512)
+        width = 640
+        height = 480
+
+    # Parse extrinsics.
+    extrinsics_data = supplement_data.get("extrinsic", np.eye(4).flatten())
+    extrinsics_matrix = np.eye(4)
+
+    # Legacy: extrinsics store 12 elements.
+    if len(extrinsics_data) == 12:
+        extrinsics_matrix[:3] = extrinsics_data.reshape((3, 4))
+        extrinsics_matrix[:3, :3] = extrinsics_matrix[:3, :3].copy().T
+    elif len(extrinsics_data) == 16:
+        extrinsics_matrix[:] = extrinsics_data.reshape((4, 4))
+    else:
+        raise ValueError(f"Unrecognized extrinsics matrix shape {len(extrinsics_data)}")
+
+    # Parse color space.
+    color_space_index = supplement_data.get("color_space", 1)
+    color_space = cs_utils.decode_color_space(color_space_index)
+    colors = torch.from_numpy(colors).view(1, -1, 3).float()
+
+    if color_space == "sRGB":
+        # Convert to linearRGB for proper alpha blending.
+        colors = cs_utils.sRGB2linearRGB(colors.flatten(0, 1)).view(1, -1, 3)
+        color_space = "linearRGB"
+
+    mean_vectors = torch.from_numpy(mean_vectors).view(1, -1, 3).float()
+    quaternions = torch.from_numpy(quaternions).view(1, -1, 4).float()
+    singular_values = torch.exp(torch.from_numpy(scale_logits).view(1, -1, 3)).float()
+    opacities = torch.sigmoid(torch.from_numpy(opacity_logits).view(1, -1)).float()
+
+    gaussians = Gaussians3D(
+        mean_vectors=mean_vectors,
+        quaternions=quaternions,
+        singular_values=singular_values,
+        opacities=opacities,
+        colors=colors,
+    )
+    metadata = SceneMetaData(focal_length_px[0], (width, height), color_space)
+    return gaussians, metadata
+
+
+@torch.no_grad()
+def save_ply(
+    gaussians: Gaussians3D, f_px: float, image_shape: tuple[int, int], path: Path
+) -> PlyData:
+    """Save a predicted Gaussian3D to a ply file."""
+
+    def _inverse_sigmoid(tensor: torch.Tensor) -> torch.Tensor:
+        return torch.log(tensor / (1.0 - tensor))
+
+    xyz = gaussians.mean_vectors.flatten(0, 1)
+    scale_logits = torch.log(gaussians.singular_values).flatten(0, 1)
+    quaternions = gaussians.quaternions.flatten(0, 1)
+
+    # SHARP takes an image, convert it to sRGB color space as input,
+    # and predicts linearRGB Gaussians as output.
+    # The SHARP renderer would blend linearRGB Gaussians and convert rendered images and videos
+    # back to sRGB for the best display quality.
+    #
+    # However, public renderers do not have such linear2sRGB conversions after rendering.
+    # If they render linearRGB Gaussians as-is, the output would be dark without Gamma correction.
+    #
+    # To make it compatible to public renderers, we force convert linearRGB to sRGB during export.
+    # - The SHARP renderer will still handle conversions properly.
+    # - Public renderers will be mostly working fine when regarding sRGB images as linearRGB images,
+    #   although for the best performance, it is recommended to apply the conversions.
+    colors = convert_rgb_to_spherical_harmonics(
+        cs_utils.linearRGB2sRGB(gaussians.colors.flatten(0, 1))
+    )
+    color_space_index = cs_utils.encode_color_space("sRGB")
+
+    # Store opacity logits.
+    opacity_logits = _inverse_sigmoid(gaussians.opacities).flatten(0, 1).unsqueeze(-1)
+
+    attributes = torch.cat(
+        (
+            xyz,
+            colors,
+            opacity_logits,
+            scale_logits,
+            quaternions,
+        ),
+        dim=1,
+    )
+
+    dtype_full = [
+        (attribute, "f4")
+        for attribute in ["x", "y", "z"]
+        + [f"f_dc_{i}" for i in range(3)]
+        + ["opacity"]
+        + [f"scale_{i}" for i in range(3)]
+        + [f"rot_{i}" for i in range(4)]
+    ]
+
+    num_gaussians = len(xyz)
+    elements = np.empty(num_gaussians, dtype=dtype_full)
+    elements[:] = list(map(tuple, attributes.detach().cpu().numpy()))
+    vertex_elements = PlyElement.describe(elements, "vertex")
+
+    # Load image-wise metadata.
+    image_height, image_width = image_shape
+
+    # Export image size.
+    dtype_image_size = [("image_size", "u4")]
+    image_size_array = np.empty(2, dtype=dtype_image_size)
+    image_size_array[:] = np.array([image_width, image_height])
+    image_size_element = PlyElement.describe(image_size_array, "image_size")
+
+    # Export intrinsics.
+    dtype_intrinsic = [("intrinsic", "f4")]
+    intrinsic_array = np.empty(9, dtype=dtype_intrinsic)
+    intrinsic = np.array(
+        [
+            f_px,
+            0,
+            image_width * 0.5,
+            0,
+            f_px,
+            image_height * 0.5,
+            0,
+            0,
+            1,
+        ]
+    )
+    intrinsic_array[:] = intrinsic.flatten()
+    intrinsic_element = PlyElement.describe(intrinsic_array, "intrinsic")
+
+    # Export dummy extrinsics.
+    dtype_extrinsic = [("extrinsic", "f4")]
+    extrinsic_array = np.empty(16, dtype=dtype_extrinsic)
+    extrinsic_array[:] = np.eye(4).flatten()
+    extrinsic_element = PlyElement.describe(extrinsic_array, "extrinsic")
+
+    # Export number of frames and particles per frame.
+    dtype_frames = [("frame", "i4")]
+    frame_array = np.empty(2, dtype=dtype_frames)
+    frame_array[:] = np.array([1, num_gaussians], dtype=np.int32)
+    frame_element = PlyElement.describe(frame_array, "frame")
+
+    # Export disparity ranges for transform.
+    dtype_disparity = [("disparity", "f4")]
+    disparity_array = np.empty(2, dtype=dtype_disparity)
+
+    disparity = 1.0 / gaussians.mean_vectors[0, ..., -1]
+    quantiles = (
+        torch.quantile(disparity, q=torch.tensor([0.1, 0.9], device=disparity.device))
+        .float()
+        .cpu()
+        .numpy()
+    )
+    disparity_array[:] = quantiles
+    disparity_element = PlyElement.describe(disparity_array, "disparity")
+
+    # Export colorspace.
+    dtype_color_space = [("color_space", "u1")]
+    color_space_array = np.empty(1, dtype=dtype_color_space)
+    color_space_array[:] = np.array([color_space_index]).flatten()
+    color_space_element = PlyElement.describe(color_space_array, "color_space")
+
+    dtype_version = [("version", "u1")]
+    version_array = np.empty(3, dtype=dtype_version)
+    version_array[:] = np.array([1, 5, 0], dtype=np.uint8).flatten()
+    version_element = PlyElement.describe(version_array, "version")
+
+    plydata = PlyData(
+        [
+            vertex_elements,
+            extrinsic_element,
+            intrinsic_element,
+            image_size_element,
+            frame_element,
+            disparity_element,
+            color_space_element,
+            version_element,
+        ]
+    )
+
+    plydata.write(path)
+    return plydata

sharp/utils/gsplat.py

@@ -0,0 +1,191 @@
+"""Contains utility code for gsplat renderer.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import NamedTuple
+
+import gsplat
+import torch
+from torch import nn
+
+from sharp.utils import color_space as cs_utils
+from sharp.utils import io, vis
+from sharp.utils.gaussians import BackgroundColor, Gaussians3D
+
+
+class RenderingOutputs(NamedTuple):
+    """Outputs of 3D Gaussians renderer."""
+
+    color: torch.Tensor
+    depth: torch.Tensor
+    alpha: torch.Tensor
+
+
+def write_renderings(rendering: RenderingOutputs, output_folder: Path, filename: str):
+    """Write rendered color/depth/alpha to files."""
+    batch_size = len(rendering.color)
+    if batch_size != 1:
+        raise RuntimeError("We only support saving rendering of batch size = 1")
+
+    def _save_image_tensor(tensor: torch.Tensor, suffix: str):
+        np_array = tensor.permute(1, 2, 0).numpy()
+        io.save_image(np_array, (output_folder / filename).with_suffix(suffix))
+
+    color = (rendering.color[0].cpu() * 255.0).to(dtype=torch.uint8)
+    colorized_depth = vis.colorize_depth(rendering.depth[0], val_max=100.0)
+    colorized_alpha = vis.colorize_alpha(rendering.alpha[0])
+
+    _save_image_tensor(color, ".color.png")
+    _save_image_tensor(colorized_depth, ".depth.png")
+    _save_image_tensor(colorized_alpha, ".alpha.png")
+
+
+class GSplatRenderer(nn.Module):
+    """Module to render 3D Gaussians to images using gsplat."""
+
+    color_space: cs_utils.ColorSpace
+    background_color: BackgroundColor
+
+    def __init__(
+        self,
+        color_space: cs_utils.ColorSpace = "sRGB",
+        background_color: BackgroundColor = "black",
+        low_pass_filter_eps: float = 0.0,
+    ) -> None:
+        """Initialize gsplat renderer.
+
+        Args:
+            color_space: The color space to use for rendering.
+            background_color: The background color to use for rendering.
+            low_pass_filter_eps: The epsilon value for the low pass filter.
+        """
+        super().__init__()
+        self.color_space = color_space
+        self.background_color = background_color
+        self.low_pass_filter_eps = low_pass_filter_eps
+
+    def forward(
+        self,
+        gaussians: Gaussians3D,
+        extrinsics: torch.Tensor,
+        intrinsics: torch.Tensor,
+        image_width: int,
+        image_height: int,
+    ) -> RenderingOutputs:
+        """Predict images from gaussians.
+
+        Args:
+            gaussians: The Gaussians to render.
+            extrinsics: The extrinsics of the camera to render to in OpenCV format.
+            intrinsics: The intriniscs of the camera to render to in OpenCV format.
+            image_width: The desired output image width.
+            image_height: The desired output image height.
+        """
+        batch_size = len(gaussians.mean_vectors)
+        outputs_list: list[RenderingOutputs] = []
+
+        for ib in range(batch_size):
+            colors, alphas, meta = gsplat.rendering.rasterization(
+                means=gaussians.mean_vectors[ib],
+                quats=gaussians.quaternions[ib],
+                scales=gaussians.singular_values[ib],
+                opacities=gaussians.opacities[ib],
+                colors=gaussians.colors[ib],
+                viewmats=extrinsics[ib : ib + 1],
+                Ks=intrinsics[ib : ib + 1, :3, :3],
+                width=image_width,
+                height=image_height,
+                render_mode="RGB+D",
+                rasterize_mode="classic",
+                absgrad=False,
+                packed=False,
+                eps2d=self.low_pass_filter_eps,
+            )
+
+            rendered_color = colors[..., 0:3].permute([0, 3, 1, 2])
+            rendered_depth_unnormalized = colors[..., 3:4].permute([0, 3, 1, 2])
+            rendered_alpha = alphas.permute([0, 3, 1, 2])
+
+            # Compose with background color.
+            rendered_color = self.compose_with_background(
+                rendered_color, rendered_alpha, self.background_color
+            )
+
+            # Colorspace conversion.
+            if self.color_space == "sRGB":
+                pass
+            elif self.color_space == "linearRGB":
+                rendered_color = cs_utils.linearRGB2sRGB(rendered_color)
+            else:
+                ValueError("Unsupported ColorSpace type.")
+
+            # splats: (B, N, 10)
+            cov2d = self._conics_to_covars2d(meta["conics"])
+            # Set the cov2d of invisible splats to 1 to avoid nan in condition number calculation..
+            splats_visible_mask = meta["depths"] > 1e-2
+            cov2d[~splats_visible_mask][..., 0, 0] = 1
+            cov2d[~splats_visible_mask][..., 1, 1] = 1
+            cov2d[~splats_visible_mask][..., 0, 1] = 0
+
+            # Normalize the depth by alpha.
+            rendered_depth = rendered_depth_unnormalized / torch.clip(rendered_alpha, min=1e-8)
+
+            outputs = RenderingOutputs(
+                color=rendered_color,
+                depth=rendered_depth,
+                alpha=rendered_alpha,
+            )
+            outputs_list.append(outputs)
+
+        return RenderingOutputs(
+            color=torch.cat([item.color for item in outputs_list], dim=0).contiguous(),
+            depth=torch.cat([item.depth for item in outputs_list], dim=0).contiguous(),
+            alpha=torch.cat([item.alpha for item in outputs_list], dim=0).contiguous(),
+        )
+
+    @staticmethod
+    def compose_with_background(
+        rendered_rgb: torch.Tensor,
+        rendered_alpha: torch.Tensor,
+        background_color: BackgroundColor,
+    ) -> torch.Tensor:
+        """Compose rendered RGB with background color."""
+        if background_color == "black":
+            return rendered_rgb
+        elif background_color == "white":
+            return rendered_rgb + (1.0 - rendered_alpha)
+        elif background_color == "random_color":
+            return (
+                rendered_rgb
+                + (1.0 - rendered_alpha)
+                * torch.rand(3, dtype=rendered_rgb.dtype, device=rendered_rgb.device)[
+                    None, :, None, None
+                ]
+            )
+        elif background_color == "random_pixel":
+            return rendered_rgb + (1.0 - rendered_alpha) * torch.rand_like(rendered_rgb)
+        else:
+            raise ValueError("Unsupported BackgroundColor type.")
+
+    @staticmethod
+    def _conics_to_covars2d(conics: torch.Tensor, eps=1e-8) -> torch.Tensor:
+        """Convert conics to covariance matrices."""
+        a = conics[..., 0]
+        b = conics[..., 1]
+        c = conics[..., 2]
+        # Reconstruct determinant.
+        det = 1 / (a * c - b**2 + eps)
+        det = det.clamp(min=eps)
+        # Reconstruct covars2d.
+        covars2d = torch.zeros(*conics.shape[:-1], 2, 2, device=conics.device)
+        covars2d[..., 1, 1] = a * det
+        covars2d[..., 0, 0] = c * det
+        covars2d[..., 0, 1] = -b * det
+        covars2d[..., 1, 0] = -b * det
+        covars2d = torch.nan_to_num(covars2d, nan=0.0, posinf=0.0, neginf=0.0)
+        return covars2d

sharp/utils/io.py

@@ -0,0 +1,213 @@
+"""Contains image IO.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+from __future__ import annotations
+
+import io
+import logging
+from pathlib import Path
+from typing import IO, Any, Protocol
+
+import imageio.v2 as iio
+import numpy as np
+import pillow_heif
+import torch
+from PIL import ExifTags, Image, TiffTags
+
+from .vis import METRIC_DEPTH_MAX_CLAMP_METER, colorize_depth
+
+LOGGER = logging.getLogger(__name__)
+
+
+# NOTE: unused, kept for reference.
+Image.MAX_IMAGE_PIXELS = 200000000
+
+
+def load_rgb(
+    path: Path, auto_rotate: bool = True, remove_alpha: bool = True
+) -> tuple[np.ndarray, list[bytes] | None, float]:
+    """Load an RGB image."""
+    LOGGER.debug(f"Loading image {path} ...")
+
+    if path.suffix.lower() in [".heic"]:
+        heif_file = pillow_heif.open_heif(path, convert_hdr_to_8bit=True)
+        img_pil = heif_file.to_pillow()
+    else:
+        img_pil = Image.open(path)
+
+    img_exif = extract_exif(img_pil)
+    icc_profile = img_pil.info.get("icc_profile", None)
+
+    # Rotate the image.
+    if auto_rotate:
+        exif_orientation = img_exif.get("Orientation", 1)
+        if exif_orientation == 3:
+            img_pil = img_pil.transpose(Image.ROTATE_180)
+        elif exif_orientation == 6:
+            img_pil = img_pil.transpose(Image.ROTATE_270)
+        elif exif_orientation == 8:
+            img_pil = img_pil.transpose(Image.ROTATE_90)
+        elif exif_orientation != 1:
+            LOGGER.warning(f"Ignoring image orientation {exif_orientation}.")
+
+    # Extract the focal length.
+    f_35mm = img_exif.get("FocalLengthIn35mmFilm", img_exif.get("FocalLenIn35mmFilm", None))
+    if f_35mm is None or f_35mm < 1:
+        f_35mm = img_exif.get("FocalLength", None)
+        if f_35mm is None:
+            LOGGER.warn(f"Did not find focallength in exif data of {path} - Setting to 30mm.")
+            f_35mm = 30.0
+        if f_35mm < 10.0:
+            LOGGER.info("Found focal length below 10mm, assuming it's not for 35mm.")
+            # This is a very crude approximation.
+            f_35mm *= 8.4
+
+    img = np.asarray(img_pil)
+    # Convert to RGB if single channel.
+    if img.ndim < 3 or img.shape[2] == 1:
+        img = np.dstack((img, img, img))
+
+    if remove_alpha:
+        img = img[:, :, :3]
+
+    LOGGER.debug(f"\tHxW: {img.shape[0]}x{img.shape[1]}")
+    LOGGER.debug(f"\tfocal length @ 35mm film: {f_35mm}mm")
+    f_px = convert_focallength(img.shape[1], img.shape[0], f_35mm)
+    LOGGER.debug(f"\tfocal length: {f_px:.2f}px")
+
+    return img, icc_profile, f_px
+
+
+def extract_exif(img_pil: Image.Image) -> dict[str, Any]:
+    """Return exif information as a dictionary."""
+    # Get full exif description from get_ifd(0x8769):
+    # cf https://pillow.readthedocs.io/en/stable/releasenotes/8.2.0.html#image-getexif-exif-and-gps-ifd # noqa
+    img_exif = img_pil.getexif().get_ifd(0x8769)
+    exif_dict = {ExifTags.TAGS[k]: v for k, v in img_exif.items() if k in ExifTags.TAGS}
+
+    # https://pillow.readthedocs.io/en/stable/_modules/PIL/TiffTags.html# # noqa
+    tiff_tags = img_pil.getexif()
+    tiff_dict = {TiffTags.TAGS_V2[k].name: v for k, v in tiff_tags.items() if k in TiffTags.TAGS_V2}
+    return {**exif_dict, **tiff_dict}
+
+
+def convert_focallength(width: float, height: float, f_mm: float = 30) -> float:
+    """Converts a focal length given in mm to pixels."""
+    return f_mm * np.sqrt(width**2.0 + height**2.0) / np.sqrt(36**2 + 24**2)
+
+
+def save_image(
+    image: np.ndarray,
+    output_path: Path,
+    icc_profile: list[bytes] | None = None,
+    jpeg_quality: int = 92,
+) -> None:
+    """Save image to given path."""
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+
+    extensions_to_format = Image.registered_extensions()
+    try:
+        format = extensions_to_format[output_path.suffix.lower()]
+    except KeyError:
+        raise ValueError(f"Unsupported output format {output_path.suffix}.")
+
+    with output_path.open("wb") as file_handle:
+        write_image(
+            image,
+            file_handle,
+            format,
+            icc_profile=icc_profile,
+            jpeg_quality=jpeg_quality,
+        )
+
+
+def write_image(
+    image: np.ndarray,
+    output_io: IO[bytes],
+    format="jpg",
+    icc_profile: list[bytes] | None = None,
+    jpeg_quality: int = 92,
+):
+    """Write image to binary stream."""
+    pil_config = {}
+    if format == "JPEG":
+        pil_config["quality"] = jpeg_quality
+
+    image_pil = Image.fromarray(image)
+
+    # Workaround to error [io.UnsupportedOperation: seek].
+    if format == "TIFF":
+        bytes_io = io.BytesIO()
+        image_pil.save(bytes_io, format="TIFF")
+        bytes_io.seek(0)
+        output_io.write(bytes_io.read())
+        return
+
+    image_pil.save(output_io, format, icc_profile=icc_profile, **pil_config)
+
+
+def get_supported_image_extensions(with_heic: bool = True) -> list[str]:
+    """Return supported image extensions."""
+    exts = Image.registered_extensions()
+    supported_extensions = {ex for ex, f in exts.items() if f in Image.OPEN}
+    if with_heic:
+        supported_extensions.add(".heic")
+
+    supported_extensions_upper = {ex.upper() for ex in supported_extensions}
+    return list(supported_extensions | supported_extensions_upper)
+
+
+def get_supported_video_extensions():
+    """Return supported video extensions."""
+    supported_extensions = {".mp4", ".mov"}
+    supported_extensions_upper = {ext.upper() for ext in supported_extensions}
+    return list(supported_extensions | supported_extensions_upper)
+
+
+class OutputWriter(Protocol):
+    """Protocol for writing output to disk."""
+
+    def add_frame(self, image: torch.Tensor, depth: torch.Tensor) -> None:
+        """Add a single frame to output."""
+        ...
+
+    def close(self) -> None:
+        """Finish writing."""
+        ...
+
+
+class VideoWriter(OutputWriter):
+    """Output writer for video output."""
+
+    def __init__(self, output_path: Path, fps: float = 30.0, render_depth: bool = True) -> None:
+        """Initialize VideoWriter."""
+        output_path.parent.mkdir(exist_ok=True, parents=True)
+        self.output_path = output_path
+        self.image_writer = iio.get_writer(output_path, fps=fps)
+
+        self.max_depth_estimate = None
+        if render_depth:
+            self.depth_writer = iio.get_writer(output_path.with_suffix(".depth.mp4"), fps=fps)
+
+    def add_frame(self, image: torch.Tensor, depth: torch.Tensor) -> None:
+        """Add a single frame to output."""
+        image_np = image.detach().cpu().numpy()
+        self.image_writer.append_data(image_np)
+
+        if self.depth_writer is not None:
+            if self.max_depth_estimate is None:
+                self.max_depth_estimate = depth.max().item()
+
+            colored_depth_pt = colorize_depth(
+                depth,
+                min(self.max_depth_estimate, METRIC_DEPTH_MAX_CLAMP_METER),  # type: ignore[call-overload]
+            )
+            colored_depth_np = colored_depth_pt.squeeze(0).permute(1, 2, 0).cpu().numpy()
+            self.depth_writer.append_data(colored_depth_np)
+
+    def close(self):
+        """Finish writing."""
+        self.image_writer.close()

sharp/utils/linalg.py

@@ -0,0 +1,104 @@
+"""Contains linear algebra related utility functions.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+from scipy.spatial.transform import Rotation
+
+
+def rotation_matrices_from_quaternions(quaternions: torch.Tensor) -> torch.Tensor:
+    """Convert batch of quaternions into rotations matrices.
+
+    Args:
+        quaternions: The quaternions convert to matrices.
+
+    Returns:
+        The rotations matrices corresponding to the (normalized) quaternions.
+    """
+    device = quaternions.device
+    shape = quaternions.shape[:-1]
+
+    quaternions = quaternions / torch.linalg.norm(quaternions, dim=-1, keepdim=True)
+    real_part = quaternions[..., 0]
+    vector_part = quaternions[..., 1:]
+
+    vector_cross = get_cross_product_matrix(vector_part)
+    real_part = real_part[..., None, None]
+
+    matrix_outer = vector_part[..., :, None] * vector_part[..., None, :]
+    matrix_diag = real_part.square() * eyes(3, shape=shape, device=device)
+    matrix_cross_1 = 2 * real_part * vector_cross
+    matrix_cross_2 = vector_cross @ vector_cross
+
+    return matrix_outer + matrix_diag + matrix_cross_1 + matrix_cross_2
+
+
+def quaternions_from_rotation_matrices(matrices: torch.Tensor) -> torch.Tensor:
+    """Convert batch of rotation matrices to quaternions.
+
+    Args:
+        matrices: The matrices to convert to quaternions.
+
+    Returns:
+        The quaternions corresponding to the rotation matrices.
+
+    Note: this operation is not differentiable and will be performed on the CPU.
+    """
+    if not matrices.shape[-2:] == (3, 3):
+        raise ValueError(f"matrices have invalid shape {matrices.shape}")
+    matrices_np = matrices.detach().cpu().numpy()
+    quaternions_np = Rotation.from_matrix(matrices_np.reshape(-1, 3, 3)).as_quat()
+    # We use a convention where the w component is at the start of the quaternion.
+    quaternions_np = quaternions_np[:, [3, 0, 1, 2]]
+    quaternions_np = quaternions_np.reshape(matrices_np.shape[:-2] + (4,))
+    return torch.as_tensor(quaternions_np, device=matrices.device, dtype=matrices.dtype)
+
+
+def get_cross_product_matrix(vectors: torch.Tensor) -> torch.Tensor:
+    """Generate cross product matrix for vector exterior product."""
+    if not vectors.shape[-1] == 3:
+        raise ValueError("Only 3-dimensional vectors are supported")
+    device = vectors.device
+    shape = vectors.shape[:-1]
+    unit_basis = eyes(3, shape=shape, device=device)
+    # We compute the matrix by multiplying each column of unit_basis with the
+    # corresponding vector.
+    return torch.cross(vectors[..., :, None], unit_basis, dim=-2)
+
+
+def eyes(
+    dim: int, shape: tuple[int, ...], device: torch.device | str | None = None
+) -> torch.Tensor:
+    """Create batch of identity matrices."""
+    return torch.eye(dim, device=device).broadcast_to(shape + (dim, dim)).clone()
+
+
+def quaternion_product(q1, q2):
+    """Compute dot product between two quaternions."""
+    real_1 = q1[..., :1]
+    real_2 = q2[..., :1]
+    vector_1 = q1[..., 1:]
+    vector_2 = q2[..., 1:]
+
+    real_out = real_1 * real_2 - (vector_1 * vector_2).sum(dim=-1, keepdim=True)
+    vector_out = real_1 * vector_2 + real_2 * vector_1 + torch.cross(vector_1, vector_2)
+    return torch.concatenate([real_out, vector_out], dim=-1)
+
+
+def quaternion_conj(q):
+    """Get conjugate of a quaternion."""
+    real = q[..., :1]
+    vector = q[..., 1:]
+    return torch.concatenate([real, -vector], dim=-1)
+
+
+def project(u: torch.Tensor, basis: torch.Tensor) -> torch.Tensor:
+    """Project tensor u to unit basis a."""
+    unit_u = F.normalize(u, dim=-1)
+    inner_prod = (unit_u * basis).sum(dim=-1, keepdim=True)
+    return inner_prod * u

sharp/utils/logging.py

@@ -0,0 +1,45 @@
+"""Contains logging related utility functions.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+from __future__ import annotations
+
+import logging
+import sys
+from pathlib import Path
+
+
+def configure(log_level: int, log_path: Path | None = None, prefix: str | None = None) -> None:
+    """Configure logger globally.
+
+    Args:
+        log_level: The desired verbosity level.
+        log_path: The path to write logs to.
+        prefix: The prefix of the logger.
+    """
+    logger = logging.getLogger(prefix)
+
+    # Reset logger to initial state (e.g. to avoid side effects from imports).
+    for handler in logger.handlers:
+        logger.removeHandler(handler)
+
+    for filter in logger.filters:
+        logger.removeFilter(filter)
+
+    # Set level.
+    logger.setLevel(log_level)
+
+    formatter = logging.Formatter("%(asctime)s | %(levelname)s | %(message)s")
+
+    # Set up console handler.
+    stdout_handler = logging.StreamHandler(sys.stdout)
+    stdout_handler.setFormatter(formatter)
+    logger.addHandler(stdout_handler)
+
+    # Set up file handler.
+    if log_path is not None:
+        file_handler = logging.FileHandler(log_path, mode="w")
+        file_handler.setFormatter(formatter)
+        logger.addHandler(file_handler)

sharp/utils/math.py

@@ -0,0 +1,183 @@
+"""Contains utility math functions.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Callable, Literal, NamedTuple, Tuple, Union
+
+import torch
+from torch import autograd
+
+ActivationType = Literal[
+    "linear",
+    "exp",
+    "sigmoid",
+    "softplus",
+    "relu_with_pushback",
+    "hard_sigmoid_with_pushback",
+]
+ActivationFunction = Callable[[torch.Tensor], torch.Tensor]
+
+
+class ActivationPair(NamedTuple):
+    """A pair of forward and inverse activation functions."""
+
+    forward: ActivationFunction
+    inverse: ActivationFunction
+
+
+def create_activation_pair(activation_type: ActivationType) -> ActivationPair:
+    """Create activation function and corresponding inverse function.
+
+    Args:
+        activation_type: The activation type to create.
+
+    Returns:
+        The corresponding activation functions and the corresponding inverse function.
+    """
+    if activation_type == "linear":
+        return ActivationPair(lambda x: x, lambda x: x)
+    elif activation_type == "exp":
+        return ActivationPair(torch.exp, torch.log)
+    elif activation_type == "sigmoid":
+        return ActivationPair(torch.sigmoid, inverse_sigmoid)
+    elif activation_type == "softplus":
+        return ActivationPair(torch.nn.functional.softplus, inverse_softplus)
+    elif activation_type == "relu_with_pushback":
+        return ActivationPair(relu_with_pushback, lambda x: x)
+    elif activation_type == "hard_sigmoid_with_pushback":
+        return ActivationPair(hard_sigmoid_with_pushback, lambda x: 6.0 * x - 3.0)
+    else:
+        raise ValueError(f"Unsupported activation function: {activation_type}.")
+
+
+def inverse_sigmoid(tensor: torch.Tensor) -> torch.Tensor:
+    """Compute inverse sigmoid."""
+    return torch.log(tensor / (1.0 - tensor))
+
+
+def inverse_softplus(tensor: torch.Tensor, eps: float = 1e-06) -> torch.Tensor:
+    """Compute inverse softplus."""
+    tensor = tensor.clamp_min(eps)
+    sigmoid = torch.sigmoid(-tensor)
+    exp = sigmoid / (1.0 - sigmoid)
+    return tensor + torch.log(-exp + 1.0)
+
+
+# The first value describes the threshold from where clamping will be applied, while
+# the second value describes the value to clamp with.
+SoftClampRange = Tuple[Union[torch.Tensor, float], Union[torch.Tensor, float]]
+
+
+def softclamp(
+    tensor: torch.Tensor,
+    min: SoftClampRange | None = None,
+    max: SoftClampRange | None = None,
+) -> torch.Tensor:
+    """Clamp tensor to min/max in differentiable way.
+
+    Args:
+        tensor: The tensor to clamp.
+        min: Pair of threshold to start clamping and value to clamp to.
+            The first value should be larger than the second.
+        max: Pair of threshold to start clamping and value to clamp to.
+            The first value should be smaller than the second.
+
+    Returns:
+        The clamped tensor.
+    """
+
+    def normalize(clamp_range: SoftClampRange) -> torch.Tensor:
+        value0, value1 = clamp_range
+        return value0 + (value1 - value0) * torch.tanh((tensor - value0) / (value1 - value0))
+
+    tensor_clamped = tensor
+    if min is not None:
+        tensor_clamped = torch.maximum(tensor_clamped, normalize(min))
+    if max is not None:
+        tensor_clamped = torch.minimum(tensor_clamped, normalize(max))
+
+    return tensor_clamped
+
+
+class ClampWithPushback(autograd.Function):
+    """Implementation of clamp_with_pushback function."""
+
+    @staticmethod
+    def forward(
+        ctx: Any,
+        tensor: torch.Tensor,
+        min: float | None,
+        max: float | None,
+        pushback: float,
+    ) -> torch.Tensor:
+        """Apply clamp."""
+        if min is not None and max is not None and min >= max:
+            raise ValueError("Only min < max is supported.")
+
+        ctx.save_for_backward(tensor)
+        ctx.min = min
+        ctx.max = max
+        ctx.pushback = pushback
+        return torch.clamp(tensor, min=min, max=max)
+
+    @staticmethod
+    def backward(  # type: ignore[override] # Deal with buggy torch annotations.
+        ctx: Any, grad_in: torch.Tensor
+    ) -> tuple[torch.Tensor, None, None, None]:
+        """Compute gradient of clamp with pushback."""
+        grad_out = grad_in.clone()
+        (tensor,) = ctx.saved_tensors
+
+        if ctx.min is not None:
+            mask_min = tensor < ctx.min
+            grad_out[mask_min] = -ctx.pushback
+
+        if ctx.max is not None:
+            mask_max = tensor > ctx.max
+            grad_out[mask_max] = ctx.pushback
+
+        return grad_out, None, None, None
+
+
+def clamp_with_pushback(
+    tensor: torch.Tensor,
+    min: float | None = None,
+    max: float | None = None,
+    pushback: float = 1e-2,
+) -> torch.Tensor:
+    """Variant of clamp function which avoid the vanishing gradient problem.
+
+    This function is equivalent to adding a regularizer of the form
+
+        pushback * sum_i (
+            relu(min - preactivation_i) + relu(preactivation_i - max)
+        )
+
+    to the full loss function, which pushes clamped values back.
+
+    When used in minimization problems, pushback should be greater than
+    zero. In maximization problems, pushback should be smaller than zero.
+    """
+    output = ClampWithPushback.apply(tensor, min, max, pushback)
+    assert isinstance(output, torch.Tensor)
+    return output
+
+
+def hard_sigmoid_with_pushback(x: torch.Tensor, slope: float = 1.0 / 6.0) -> torch.Tensor:
+    """Apply hard sigmoid with pushback.
+
+    For compatibility reasons, we follow the default PyTorch implementation with a
+    default slope of 1/6:
+
+        https://pytorch.org/docs/stable/generated/torch.nn.Hardsigmoid.html
+    """
+    return clamp_with_pushback(slope * x + 0.5, min=0.0, max=1.0)
+
+
+def relu_with_pushback(x: torch.Tensor) -> torch.Tensor:
+    """Compute relu with pushback."""
+    return clamp_with_pushback(x, min=0.0)

sharp/utils/module_surgery.py

@@ -0,0 +1,32 @@
+"""Contains utility functionality to modify torch modules.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from torch import nn
+
+NORM_LAYER_TYPES = tuple(module_type for name, module_type in nn.__dict__.items() if "Norm" in name)
+BATCH_NORM_LAYER_TYPES = tuple(
+    module_type for name, module_type in nn.__dict__.items() if "BatchNorm" in name
+)
+
+
+def freeze_norm_layer(module: nn.Module) -> nn.Module:
+    """Freeze all normalization layers."""
+
+    def set_module_eval_mode(module: nn.Module, _: Any) -> None:
+        module.eval()
+
+    for submodule in module.modules():
+        if isinstance(submodule, NORM_LAYER_TYPES):
+            submodule.requires_grad_(False)
+            # This is to ensure that batch norm layers are always called
+            # with the precomputed running statistics.
+            submodule.register_forward_pre_hook(set_module_eval_mode)
+
+    return module

sharp/utils/robust.py

@@ -0,0 +1,45 @@
+"""Utility function for loss implementations.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+from __future__ import annotations
+
+from typing import Callable
+
+import torch
+
+
+def robust_where(
+    condition: torch.Tensor,
+    input: torch.Tensor,
+    branch_true_func: Callable[[torch.Tensor], torch.Tensor],
+    branch_false_func: Callable[[torch.Tensor], torch.Tensor],
+    branch_true_safe_value: float | None = None,
+    branch_false_safe_value: float | None = None,
+) -> torch.Tensor:
+    """Robust torch.where function to avoid NaN in backward pass.
+
+    See https://github.com/pytorch/pytorch/issues/68425
+
+    Args:
+        condition: When True (nonzero), yield branch_true_func(input),
+            otherwise yield branch_false_func(input)
+        input: The input tensor for torch.where
+        branch_true_func: Callable for values at indices where condition is True.
+        branch_false_func: Callable for values at indices where condition is False.
+        branch_true_safe_value: Safe value to replace the true branch.
+        branch_false_safe_value: Safe value to replace the false branch.
+    """
+    input_1 = input
+    input_2 = input
+    if branch_true_safe_value is not None:
+        input_1 = torch.where(condition, input_1, branch_true_safe_value)
+    if branch_false_safe_value is not None:
+        input_2 = torch.where(~condition, input_2, branch_false_safe_value)
+    return torch.where(
+        condition,
+        branch_true_func(input_1),
+        branch_false_func(input_2),
+    )

sharp/utils/training.py

@@ -0,0 +1,24 @@
+"""Utility functions for training.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+from torch.utils.checkpoint import checkpoint
+
+
+def checkpoint_wrapper(self, fn, *args):
+    """Helper function that applies checkpointing.
+
+    If enabled applies grad checkpointing, otherwise just executes the function normally.
+    """
+    if not hasattr(self, "grad_checkpointing"):
+        raise AttributeError(
+            "Trying to apply grad checkpointing on a model that does not have a grad_checkpointing "
+            "attribute."
+        )
+
+    if self.grad_checkpointing:
+        return checkpoint(fn, *args, use_reentrant=False)
+    else:
+        return fn(*args)

sharp/utils/vis.py

@@ -0,0 +1,77 @@
+"""Utility functions for visualization.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import torch
+from matplotlib import pyplot as plt
+
+METRIC_DEPTH_MAX_CLAMP_METER = 50.0
+
+
+def colorize_depth(depth: torch.Tensor, val_max: float = 10.0) -> torch.Tensor:
+    """Colorize depth map."""
+    depth_channels = depth.shape[-3]
+
+    # When we have a general depth/disparity map, output the color map as is.
+    if depth_channels == 1:
+        return colorize_scalar_map(
+            depth.squeeze(-3), val_min=0.0, val_max=val_max, color_map="turbo"
+        )
+
+    # When we have a multi-layered depth/disparity map,
+    # we concatenate the color maps horizontally and output it.
+    else:
+        colored_depths = []
+        for c in range(depth_channels):
+            colored_depths.append(
+                colorize_scalar_map(
+                    depth[..., c, :, :], val_min=0.0, val_max=val_max, color_map="turbo"
+                )
+            )
+        return torch.cat(colored_depths, dim=-1)
+
+
+def colorize_alpha(alpha: torch.Tensor) -> torch.Tensor:
+    """Colorize alpha map."""
+    return colorize_scalar_map(alpha.squeeze(-3), val_min=0.0, val_max=1.0, color_map="coolwarm")
+
+
+def colorize_scalar_map(
+    scalar_map: torch.Tensor, val_min=0.0, val_max=1.0, color_map: str = "jet"
+) -> torch.Tensor:
+    """Colorize a scalar map of.
+
+    Args:
+        scalar_map: Map of with format BHW.
+        val_min: Minimu value to display.
+        val_max: Maximum value to display.
+        color_map: Which color map to use. Will be passed to matplotlob.
+
+    Returns:
+        A colorized image with format BHWC.
+    """
+    if scalar_map.ndim not in (2, 3, 4):
+        raise ValueError("Only scalar maps of 2 or 3 or 4 dimensions supported.")
+
+    cmap = plt.get_cmap(color_map)
+
+    scalar_map_np = scalar_map.detach().cpu().float().numpy()
+    scalar_map_np = (scalar_map_np - val_min) / (val_max - val_min)
+    scalar_map_np = np.clip(scalar_map_np, a_min=0.0, a_max=1.0)
+
+    color_map_np = cmap(scalar_map_np)[..., :3]
+    tensor = torch.as_tensor(color_map_np * 255.0, dtype=torch.uint8)
+
+    if tensor.ndim == 3:
+        return tensor.permute(2, 0, 1)
+    elif tensor.ndim == 4:
+        return tensor.permute(0, 3, 1, 2)
+    elif tensor.ndim == 5:
+        return tensor.permute(0, 1, 4, 2, 3)
+    else:
+        assert False, "Invalid tensor shape encountered."