vendored model

__init__.py

@@ -6,6 +6,14 @@ SHARP (Sharp Monocular View Synthesis) takes a single image and produces
 Based on Apple's SHARP model: https://arxiv.org/abs/2512.10685
 """
 
+import os
+import sys
+
+# Add this directory to sys.path so the vendored 'sharp' package can be imported
+_THIS_DIR = os.path.dirname(os.path.abspath(__file__))
+if _THIS_DIR not in sys.path:
+    sys.path.insert(0, _THIS_DIR)
+
 from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
 
 __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]

nodes/load_model.py

@@ -1,33 +1,14 @@
 """LoadSharpModel node for ComfyUI-Sharp."""
 
 import os
-import sys
 
 import torch
 
-# Add ml-sharp to path
-# Priority: SHARP_PATH env var > /home/shadeform/ml-sharp/src > relative path
-ML_SHARP_PATH = os.environ.get("SHARP_PATH")
-if not ML_SHARP_PATH:
-    # Try /home/shadeform/ml-sharp/src first
-    default_path = "/home/shadeform/ml-sharp/src"
-    if os.path.exists(default_path):
-        ML_SHARP_PATH = default_path
-    else:
-        # Fallback to relative path from this file
-        ML_SHARP_PATH = os.path.join(
-            os.path.dirname(os.path.dirname(os.path.dirname(__file__))),
-            "ml-sharp", "src"
-        )
-
-if ML_SHARP_PATH not in sys.path:
-    sys.path.insert(0, ML_SHARP_PATH)
-
 # Model cache
 _MODEL_CACHE = {}
 
-# Default model URL
-DEFAULT_MODEL_URL = "https://ml-site.cdn-apple.com/models/sharp/sharp_2572gikvuh.pt"
+# Default model URL (Hugging Face mirror)
+DEFAULT_MODEL_URL = "https://huggingface.co/apple/Sharp/resolve/main/sharp_2572gikvuh.pt"
 
 
 class LoadSharpModel:
@@ -45,7 +26,7 @@ class LoadSharpModel:
             "optional": {
                 "checkpoint_path": ("STRING", {
                     "default": "",
-                    "tooltip": "Path to .pt checkpoint. Leave empty to auto-download from Apple CDN."
+                    "tooltip": "Path to .pt checkpoint. Leave empty to auto-download from Hugging Face."
                 }),
             }
         }
@@ -79,7 +60,7 @@ class LoadSharpModel:
         if checkpoint_path and os.path.exists(checkpoint_path):
             state_dict = torch.load(checkpoint_path, weights_only=True)
         else:
-            print(f"[SHARP] Downloading model from {DEFAULT_MODEL_URL}")
+            print(f"[SHARP] Downloading model from Hugging Face...")
             state_dict = torch.hub.load_state_dict_from_url(DEFAULT_MODEL_URL, progress=True)
 
         # Create predictor

nodes/predict.py

@@ -1,7 +1,6 @@
 """SharpPredict node for ComfyUI-Sharp."""
 
 import os
-import sys
 import time
 from pathlib import Path
 
@@ -9,24 +8,6 @@ import numpy as np
 import torch
 import torch.nn.functional as F
 
-# Add ml-sharp to path
-# Priority: SHARP_PATH env var > /home/shadeform/ml-sharp/src > relative path
-ML_SHARP_PATH = os.environ.get("SHARP_PATH")
-if not ML_SHARP_PATH:
-    # Try /home/shadeform/ml-sharp/src first
-    default_path = "/home/shadeform/ml-sharp/src"
-    if os.path.exists(default_path):
-        ML_SHARP_PATH = default_path
-    else:
-        # Fallback to relative path from this file
-        ML_SHARP_PATH = os.path.join(
-            os.path.dirname(os.path.dirname(os.path.dirname(__file__))),
-            "ml-sharp", "src"
-        )
-
-if ML_SHARP_PATH not in sys.path:
-    sys.path.insert(0, ML_SHARP_PATH)
-
 # Try to import ComfyUI folder_paths for output directory
 try:
     import folder_paths

sharp/__init__.py

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

sharp/cli/__init__.py

@@ -0,0 +1,19 @@
+"""Command-line-interface to run SHARP model.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+import click
+
+from . import predict, render
+
+
+@click.group()
+def main_cli():
+    """Run inference for SHARP model."""
+    pass
+
+
+main_cli.add_command(predict.predict_cli, "predict")
+main_cli.add_command(render.render_cli, "render")

sharp/cli/predict.py

@@ -0,0 +1,206 @@
+"""Contains `sharp predict` CLI implementation.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+
+import click
+import numpy as np
+import torch
+import torch.nn.functional as F
+import torch.utils.data
+
+from sharp.models import (
+    PredictorParams,
+    RGBGaussianPredictor,
+    create_predictor,
+)
+from sharp.utils import io
+from sharp.utils import logging as logging_utils
+from sharp.utils.gaussians import (
+    Gaussians3D,
+    SceneMetaData,
+    save_ply,
+    unproject_gaussians,
+)
+
+from .render import render_gaussians
+
+LOGGER = logging.getLogger(__name__)
+
+DEFAULT_MODEL_URL = "https://ml-site.cdn-apple.com/models/sharp/sharp_2572gikvuh.pt"
+
+
+@click.command()
+@click.option(
+    "-i",
+    "--input-path",
+    type=click.Path(path_type=Path, exists=True),
+    help="Path to an image or containing a list of images.",
+    required=True,
+)
+@click.option(
+    "-o",
+    "--output-path",
+    type=click.Path(path_type=Path, file_okay=False),
+    help="Path to save the predicted Gaussians and renderings.",
+    required=True,
+)
+@click.option(
+    "-c",
+    "--checkpoint-path",
+    type=click.Path(path_type=Path, dir_okay=False),
+    default=None,
+    help="Path to the .pt checkpoint. If not provided, downloads the default model automatically.",
+    required=False,
+)
+@click.option(
+    "--render/--no-render",
+    "with_rendering",
+    is_flag=True,
+    default=False,
+    help="Whether to render trajectory for checkpoint.",
+)
+@click.option(
+    "--device",
+    type=str,
+    default="default",
+    help="Device to run on. ['cpu', 'mps', 'cuda']",
+)
+@click.option("-v", "--verbose", is_flag=True, help="Activate debug logs.")
+def predict_cli(
+    input_path: Path,
+    output_path: Path,
+    checkpoint_path: Path,
+    with_rendering: bool,
+    device: str,
+    verbose: bool,
+):
+    """Predict Gaussians from input images."""
+    logging_utils.configure(logging.DEBUG if verbose else logging.INFO)
+
+    extensions = io.get_supported_image_extensions()
+
+    image_paths = []
+    if input_path.is_file():
+        if input_path.suffix in extensions:
+            image_paths = [input_path]
+    else:
+        for ext in extensions:
+            image_paths.extend(list(input_path.glob(f"**/*{ext}")))
+
+    if len(image_paths) == 0:
+        LOGGER.info("No valid images found. Input was %s.", input_path)
+        return
+
+    LOGGER.info("Processing %d valid image files.", len(image_paths))
+
+    if device == "default":
+        if torch.cuda.is_available():
+            device = "cuda"
+        elif torch.mps.is_available():
+            device = "mps"
+        else:
+            device = "cpu"
+    LOGGER.info("Using device %s", device)
+
+    if with_rendering and device != "cuda":
+        LOGGER.warning("Can only run rendering with gsplat on CUDA. Rendering is disabled.")
+        with_rendering = False
+
+    # Load or download checkpoint
+    if checkpoint_path is None:
+        LOGGER.info("No checkpoint provided. Downloading default model from %s", DEFAULT_MODEL_URL)
+        state_dict = torch.hub.load_state_dict_from_url(DEFAULT_MODEL_URL, progress=True)
+    else:
+        LOGGER.info("Loading checkpoint from %s", checkpoint_path)
+        state_dict = torch.load(checkpoint_path, weights_only=True)
+
+    gaussian_predictor = create_predictor(PredictorParams())
+    gaussian_predictor.load_state_dict(state_dict)
+    gaussian_predictor.eval()
+    gaussian_predictor.to(device)
+
+    output_path.mkdir(exist_ok=True, parents=True)
+
+    for image_path in image_paths:
+        LOGGER.info("Processing %s", image_path)
+        image, _, f_px = io.load_rgb(image_path)
+        height, width = image.shape[:2]
+        intrinsics = torch.tensor(
+            [
+                [f_px, 0, (width - 1) / 2.0, 0],
+                [0, f_px, (height - 1) / 2.0, 0],
+                [0, 0, 1, 0],
+                [0, 0, 0, 1],
+            ],
+            device=device,
+            dtype=torch.float32,
+        )
+        gaussians = predict_image(gaussian_predictor, image, f_px, torch.device(device))
+
+        LOGGER.info("Saving 3DGS to %s", output_path)
+        save_ply(gaussians, f_px, (height, width), output_path / f"{image_path.stem}.ply")
+
+        if with_rendering:
+            output_video_path = (output_path / image_path.stem).with_suffix(".mp4")
+            LOGGER.info("Rendering trajectory to %s", output_video_path)
+
+            metadata = SceneMetaData(intrinsics[0, 0].item(), (width, height), "linearRGB")
+            render_gaussians(gaussians, metadata, output_video_path)
+
+
+@torch.no_grad()
+def predict_image(
+    predictor: RGBGaussianPredictor,
+    image: np.ndarray,
+    f_px: float,
+    device: torch.device,
+) -> Gaussians3D:
+    """Predict Gaussians from an image."""
+    internal_shape = (1536, 1536)
+
+    LOGGER.info("Running preprocessing.")
+    image_pt = torch.from_numpy(image.copy()).float().to(device).permute(2, 0, 1) / 255.0
+    _, height, width = image_pt.shape
+    disparity_factor = torch.tensor([f_px / width]).float().to(device)
+
+    image_resized_pt = F.interpolate(
+        image_pt[None],
+        size=(internal_shape[1], internal_shape[0]),
+        mode="bilinear",
+        align_corners=True,
+    )
+
+    # Predict Gaussians in the NDC space.
+    LOGGER.info("Running inference.")
+    gaussians_ndc = predictor(image_resized_pt, disparity_factor)
+
+    LOGGER.info("Running postprocessing.")
+    intrinsics = (
+        torch.tensor(
+            [
+                [f_px, 0, width / 2, 0],
+                [0, f_px, height / 2, 0],
+                [0, 0, 1, 0],
+                [0, 0, 0, 1],
+            ]
+        )
+        .float()
+        .to(device)
+    )
+    intrinsics_resized = intrinsics.clone()
+    intrinsics_resized[0] *= internal_shape[0] / width
+    intrinsics_resized[1] *= internal_shape[1] / height
+
+    # Convert Gaussians to metrics space.
+    gaussians = unproject_gaussians(
+        gaussians_ndc, torch.eye(4).to(device), intrinsics_resized, internal_shape
+    )
+
+    return gaussians

sharp/cli/render.py

@@ -0,0 +1,120 @@
+"""Contains `sharp render` CLI implementation.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+
+import click
+import torch
+import torch.utils.data
+
+from sharp.utils import camera, gsplat, io
+from sharp.utils import logging as logging_utils
+from sharp.utils.gaussians import Gaussians3D, SceneMetaData, load_ply
+
+LOGGER = logging.getLogger(__name__)
+
+
+@click.command()
+@click.option(
+    "-i",
+    "--input-path",
+    type=click.Path(exists=True, path_type=Path),
+    help="Path to the ply or a list of plys.",
+    required=True,
+)
+@click.option(
+    "-o",
+    "--output-path",
+    type=click.Path(path_type=Path, file_okay=False),
+    help="Path to save the rendered videos.",
+    required=True,
+)
+@click.option("-v", "--verbose", is_flag=True, help="Activate debug logs.")
+def render_cli(input_path: Path, output_path: Path, verbose: bool):
+    """Predict Gaussians from input images."""
+    logging_utils.configure(logging.DEBUG if verbose else logging.INFO)
+
+    if not torch.cuda.is_available():
+        LOGGER.error("Rendering a checkpoint requires CUDA.")
+        exit(1)
+
+    output_path.mkdir(exist_ok=True, parents=True)
+
+    params = camera.TrajectoryParams()
+
+    if input_path.suffix == ".ply":
+        scene_paths = [input_path]
+    elif input_path.is_dir():
+        scene_paths = list(input_path.glob("*.ply"))
+    else:
+        LOGGER.error("Input path must be either directory or single PLY file.")
+        exit(1)
+
+    for scene_path in scene_paths:
+        LOGGER.info("Rendering %s", scene_path)
+        gaussians, metadata = load_ply(scene_path)
+        render_gaussians(
+            gaussians=gaussians,
+            metadata=metadata,
+            params=params,
+            output_path=(output_path / scene_path.stem).with_suffix(".mp4"),
+        )
+
+
+def render_gaussians(
+    gaussians: Gaussians3D,
+    metadata: SceneMetaData,
+    output_path: Path,
+    params: camera.TrajectoryParams | None = None,
+) -> None:
+    """Render a single gaussian checkpoint file."""
+    (width, height) = metadata.resolution_px
+    f_px = metadata.focal_length_px
+
+    if params is None:
+        params = camera.TrajectoryParams()
+
+    if not torch.cuda.is_available():
+        raise RuntimeError("Rendering a checkpoint requires CUDA.")
+
+    device = torch.device("cuda")
+
+    intrinsics = torch.tensor(
+        [
+            [f_px, 0, (width - 1) / 2.0, 0],
+            [0, f_px, (height - 1) / 2.0, 0],
+            [0, 0, 1, 0],
+            [0, 0, 0, 1],
+        ],
+        device=device,
+        dtype=torch.float32,
+    )
+    camera_model = camera.create_camera_model(
+        gaussians, intrinsics, resolution_px=metadata.resolution_px
+    )
+
+    trajectory = camera.create_eye_trajectory(
+        gaussians, params, resolution_px=metadata.resolution_px, f_px=f_px
+    )
+    renderer = gsplat.GSplatRenderer(color_space=metadata.color_space)
+    video_writer = io.VideoWriter(output_path)
+
+    for _, eye_position in enumerate(trajectory):
+        camera_info = camera_model.compute(eye_position)
+        rendering_output = renderer(
+            gaussians.to(device),
+            extrinsics=camera_info.extrinsics[None].to(device),
+            intrinsics=camera_info.intrinsics[None].to(device),
+            image_width=camera_info.width,
+            image_height=camera_info.height,
+        )
+        color = (rendering_output.color[0].permute(1, 2, 0) * 255.0).to(dtype=torch.uint8)
+        depth = rendering_output.depth[0]
+        video_writer.add_frame(color, depth)
+    video_writer.close()

sharp/models/__init__.py

@@ -0,0 +1,79 @@
+"""Contains different Gaussian predictors.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+from __future__ import annotations
+
+from sharp.models.monodepth import (
+    create_monodepth_adaptor,
+    create_monodepth_dpt,
+)
+
+from .alignment import create_alignment
+from .composer import GaussianComposer
+from .gaussian_decoder import create_gaussian_decoder
+from .heads import DirectPredictionHead
+from .initializer import create_initializer
+from .params import PredictorParams
+from .predictor import RGBGaussianPredictor
+
+
+def create_predictor(params: PredictorParams) -> RGBGaussianPredictor:
+    """Create gaussian predictor model specified by name."""
+    if params.gaussian_decoder.stride < params.initializer.stride:
+        raise ValueError(
+            "We donot expected gaussian_decoder has higher resolution than initializer."
+        )
+
+    scale_factor = params.gaussian_decoder.stride // params.initializer.stride
+    gaussian_composer = GaussianComposer(
+        delta_factor=params.delta_factor,
+        min_scale=params.min_scale,
+        max_scale=params.max_scale,
+        color_activation_type=params.color_activation_type,
+        opacity_activation_type=params.opacity_activation_type,
+        color_space=params.color_space,
+        scale_factor=scale_factor,
+        base_scale_on_predicted_mean=params.base_scale_on_predicted_mean,
+    )
+    if params.num_monodepth_layers > 1 and params.initializer.num_layers != 2:
+        raise KeyError("We only support num_layers = 2 when num_monodepth_layers > 1.")
+
+    monodepth_model = create_monodepth_dpt(params.monodepth)
+    monodepth_adaptor = create_monodepth_adaptor(
+        monodepth_model,
+        params.monodepth_adaptor,
+        params.num_monodepth_layers,
+        params.sorting_monodepth,
+    )
+
+    if params.num_monodepth_layers == 2:
+        monodepth_adaptor.replicate_head(params.num_monodepth_layers)
+
+    gaussian_decoder = create_gaussian_decoder(
+        params.gaussian_decoder,
+        dims_depth_features=monodepth_adaptor.get_feature_dims(),
+    )
+    initializer = create_initializer(
+        params.initializer,
+    )
+    prediction_head = DirectPredictionHead(
+        feature_dim=gaussian_decoder.dim_out, num_layers=initializer.num_layers
+    )
+    decoder_dim = monodepth_model.decoder.dims_decoder[-1]
+    return RGBGaussianPredictor(
+        init_model=initializer,
+        feature_model=gaussian_decoder,
+        prediction_head=prediction_head,
+        monodepth_model=monodepth_adaptor,
+        gaussian_composer=gaussian_composer,
+        scale_map_estimator=create_alignment(params.depth_alignment, depth_decoder_dim=decoder_dim),
+    )
+
+
+__all__ = [
+    "PredictorParams",
+    "create_predictor",
+]

sharp/models/alignment.py

@@ -0,0 +1,126 @@
+"""Contains modules for different types of alignment.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from sharp.models.decoders import UNetDecoder
+from sharp.models.encoders import UNetEncoder
+from sharp.utils import math as math_utils
+
+from .params import AlignmentParams
+
+
+def create_alignment(
+    params: AlignmentParams, depth_decoder_dim: int | None = None
+) -> nn.Module | None:
+    """Create depth alignment."""
+    if depth_decoder_dim is None:
+        raise ValueError("Requires depth_decoder_dim for LearnedAlignment.")
+    alignment = LearnedAlignment(
+        depth_decoder_features=params.depth_decoder_features,
+        depth_decoder_dim=depth_decoder_dim,
+        steps=params.steps,
+        stride=params.stride,
+        base_width=params.base_width,
+        activation_type=params.activation_type,
+    )
+
+    if params.frozen:
+        alignment.requires_grad_(False)
+
+    return alignment
+
+
+class LearnedAlignment(nn.Module):
+    """Aligns tensors using a UNet."""
+
+    def __init__(
+        self,
+        steps: int = 4,
+        stride: int = 8,
+        base_width: int = 16,
+        depth_decoder_features: bool = False,
+        depth_decoder_dim: int = 256,
+        activation_type: math_utils.ActivationType = "exp",
+    ) -> None:
+        """Initialize LearnedAlignment.
+
+        Args:
+            steps: Number of steps in the UNet.
+            stride: Effective downsampling of the alignment module.
+            base_width: Base width of the UNet.
+            depth_decoder_features: Whether to use depth decoder features.
+            depth_decoder_dim: Dimension of the depth decoder features.
+            activation_type: Activation type for the alignment output.
+        """
+        super().__init__()
+        self.activation = math_utils.create_activation_pair(activation_type)
+        bias_value = self.activation.inverse(torch.tensor(1.0))
+
+        self.depth_decoder_features = depth_decoder_features
+        if depth_decoder_features:
+            dim_in = 2 + depth_decoder_dim
+        else:
+            dim_in = 2
+
+        def is_power_of_two(n: int) -> bool:
+            """Check if a number is a power of two."""
+            if n <= 0:
+                return False
+            return (n & (n - 1)) == 0
+
+        if not is_power_of_two(stride):
+            raise ValueError(f"Stride {stride} is not a power of two.")
+
+        steps_decoder = steps - int(math.log2(stride))
+        if steps_decoder < 1:
+            raise ValueError(f"{steps_decoder} must be greater or equal to 1.")
+        widths = [min(base_width << i, 1024) for i in range(steps + 1)]
+        self.encoder = UNetEncoder(dim_in=dim_in, width=widths, steps=steps, norm_num_groups=4)
+        self.decoder = UNetDecoder(
+            dim_out=widths[0], width=widths, steps=steps_decoder, norm_num_groups=4
+        )
+        self.conv_out = nn.Conv2d(widths[0], 1, 1, bias=True)
+        nn.init.zeros_(self.conv_out.weight)
+        nn.init.constant_(self.conv_out.bias, bias_value)
+
+    def forward(
+        self,
+        tensor_src: torch.Tensor,
+        tensor_tgt: torch.Tensor,
+        depth_decoder_features: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """Compute alignment map."""
+        # Since the tensors are usually given by depth which is >= 1.0, we invert
+        # the tensors to have them in a reasonable range.
+        tensor_src = 1.0 / tensor_src.clamp(min=1e-4)
+        tensor_tgt = 1.0 / tensor_tgt.clamp(min=1e-4)
+        tensor_input = torch.cat([tensor_src, tensor_tgt], dim=1)
+        if self.depth_decoder_features:
+            height, width = tensor_src.shape[-2:]
+            upsampled_encodings = F.interpolate(
+                depth_decoder_features,
+                size=(height, width),
+                mode="bilinear",
+            )
+            tensor_input = torch.cat([tensor_input, upsampled_encodings], dim=1)
+        features = self.encoder(tensor_input)
+        output = self.conv_out(self.decoder(features))
+        alignment_map_lowres = self.activation.forward(output)
+        if alignment_map_lowres.shape[-2:] != tensor_src.shape[-2]:
+            alignment_map = F.interpolate(
+                alignment_map_lowres,
+                size=tensor_src.shape[-2:],
+                mode="bilinear",
+                align_corners=False,
+            )
+        return alignment_map

sharp/models/blocks.py

@@ -0,0 +1,210 @@
+"""Contains reusable network components.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+from __future__ import annotations
+
+from typing import Literal
+
+import torch
+from torch import nn
+
+NormLayerName = Literal["noop", "batch_norm", "group_norm", "instance_norm"]
+UpsamplingMode = Literal["transposed_conv", "nearest", "bilinear"]
+
+
+def norm_layer_2d(num_features: int, norm_type: NormLayerName, num_groups: int = 8) -> nn.Module:
+    """Create normalization layer."""
+    if norm_type == "noop":
+        return nn.Identity()
+    elif norm_type == "batch_norm":
+        return nn.BatchNorm2d(num_features=num_features)
+    elif norm_type == "group_norm":
+        return nn.GroupNorm(num_channels=num_features, num_groups=num_groups)
+    elif norm_type == "instance_norm":
+        return nn.InstanceNorm2d(num_features=num_features)
+    else:
+        raise ValueError(f"Invalid normalization layer type: {norm_type}")
+
+
+def upsampling_layer(upsampling_mode: UpsamplingMode, scale_factor: int, dim_in: int) -> nn.Module:
+    """Create upsampling layer."""
+    if upsampling_mode == "transposed_conv":
+        return nn.ConvTranspose2d(
+            in_channels=dim_in,
+            out_channels=dim_in,
+            kernel_size=scale_factor,
+            stride=scale_factor,
+            padding=0,
+            bias=False,
+        )
+    elif upsampling_mode in ("nearest", "bilinear"):
+        return nn.Upsample(scale_factor=scale_factor, mode=upsampling_mode)
+    else:
+        raise ValueError(f"Invalid upsampling mode {upsampling_mode}.")
+
+
+class ResidualBlock(nn.Module):
+    """Generic implementation of residual blocks.
+
+    This implements a generic residual block from
+
+        He et al. - Identity Mappings in Deep Residual Networks (2016),
+        https://arxiv.org/abs/1603.05027
+
+    which can be further customized via factory functions.
+    """
+
+    def __init__(self, residual: nn.Module, shortcut: nn.Module | None = None) -> None:
+        """Initialize ResidualBlock."""
+        super().__init__()
+        self.residual = residual
+        self.shortcut = shortcut
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Apply residual block."""
+        delta_x = self.residual(x)
+
+        if self.shortcut is not None:
+            x = self.shortcut(x)
+
+        return x + delta_x
+
+
+def residual_block_2d(
+    dim_in: int,
+    dim_out: int,
+    dim_hidden: int | None = None,
+    actvn: nn.Module | None = None,
+    norm_type: NormLayerName = "noop",
+    norm_num_groups: int = 8,
+    dilation: int = 1,
+    kernel_size: int = 3,
+):
+    """Create a simple 2D residual block."""
+    if actvn is None:
+        actvn = nn.ReLU()
+
+    if dim_hidden is None:
+        dim_hidden = dim_out // 2
+
+    # Padding to maintain output size
+    # See https://pytorch.org/docs/stable/generated/torch.nn.Conv2d.html
+    padding = (dilation * (kernel_size - 1)) // 2
+
+    def _create_block(dim_in: int, dim_out: int) -> list[nn.Module]:
+        layers = [
+            norm_layer_2d(dim_in, norm_type, num_groups=norm_num_groups),
+            actvn,
+        ]
+
+        layers.append(
+            nn.Conv2d(
+                dim_in,
+                dim_out,
+                kernel_size=kernel_size,
+                stride=1,
+                dilation=dilation,
+                padding=padding,
+            )
+        )
+        return layers
+
+    residual = nn.Sequential(
+        *_create_block(dim_in, dim_hidden),
+        *_create_block(dim_hidden, dim_out),
+    )
+    shortcut = None
+
+    if dim_in != dim_out:
+        shortcut = nn.Conv2d(dim_in, dim_out, 1)
+
+    return ResidualBlock(residual, shortcut)
+
+
+class FeatureFusionBlock2d(nn.Module):
+    """Feature fusion for DPT."""
+
+    # We use the name "deconv" for backward compatibility. However, "deconv" can also
+    # refer to some other upsampling layer or a no-op.
+    deconv: nn.Module
+
+    def __init__(
+        self,
+        dim_in: int,
+        dim_out: int | None = None,
+        upsampling_mode: UpsamplingMode | None = None,
+        batch_norm: bool = False,
+    ):
+        """Initialize feature fusion block.
+
+        Args:
+            dim_in: Dimensions of input.
+            dim_out: Dimensions of output.
+            batch_norm: Whether to use batch normalization in resnet blocks.
+            upsampling_mode: What mode to use for upsampling. None if no upsampling
+                is required.
+        """
+        super().__init__()
+        if dim_out is None:
+            dim_out = dim_in
+        self.resnet1 = self._residual_block(dim_in, batch_norm)
+        self.resnet2 = self._residual_block(dim_in, batch_norm)
+
+        if upsampling_mode is not None:
+            self.deconv = upsampling_layer(upsampling_mode, scale_factor=2, dim_in=dim_in)
+        else:
+            self.deconv = nn.Sequential()
+
+        self.out_conv = nn.Conv2d(
+            dim_in,
+            dim_out,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            bias=True,
+        )
+
+        self.skip_add = nn.quantized.FloatFunctional()
+
+    def forward(self, x0: torch.Tensor, x1: torch.Tensor | None = None) -> torch.Tensor:
+        """Process and fuse input features."""
+        x = x0
+
+        if x1 is not None:
+            res = self.resnet1(x1)
+            x = self.skip_add.add(x, res)
+
+        x = self.resnet2(x)
+        x = self.deconv(x)
+        x = self.out_conv(x)
+
+        return x
+
+    @staticmethod
+    def _residual_block(num_features: int, batch_norm: bool):
+        """Create a residual block."""
+
+        def _create_block(dim: int, batch_norm: bool) -> list[nn.Module]:
+            layers = [
+                nn.ReLU(False),
+                nn.Conv2d(
+                    num_features,
+                    num_features,
+                    kernel_size=3,
+                    stride=1,
+                    padding=1,
+                    bias=not batch_norm,
+                ),
+            ]
+            if batch_norm:
+                layers.append(nn.BatchNorm2d(dim))
+            return layers
+
+        residual = nn.Sequential(
+            *_create_block(dim=num_features, batch_norm=batch_norm),
+            *_create_block(dim=num_features, batch_norm=batch_norm),
+        )
+        return ResidualBlock(residual)

sharp/models/composer.py

@@ -0,0 +1,251 @@
+"""Defines module to compose final Gaussians from base values and delta values.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+from __future__ import annotations
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from sharp.models.initializer import GaussianBaseValues
+from sharp.utils import math as math_utils
+from sharp.utils.color_space import ColorSpace, sRGB2linearRGB
+from sharp.utils.gaussians import Gaussians3D
+
+from .params import DeltaFactor
+
+
+def _get_scale_activation_constant(max_scale: float, min_scale: float) -> tuple[float, float]:
+    """Return constants for scale activation function."""
+    # To ensure for delta = 0, the value of scale_factor is 1 and the gradient is 1.
+    constant_a = (max_scale - min_scale) / (1 - min_scale) / (max_scale - 1)
+    constant_b = math_utils.inverse_sigmoid(
+        torch.tensor((1.0 - min_scale) / (max_scale - min_scale))
+    ).item()
+    return constant_a, constant_b
+
+
+class GaussianComposer(nn.Module):
+    """Converts base values and deltas into Gaussians."""
+
+    color_activation_type: math_utils.ActivationType
+    opacity_activation_type: math_utils.ActivationType
+
+    def __init__(
+        self,
+        delta_factor: DeltaFactor,
+        min_scale: float,
+        max_scale: float,
+        color_activation_type: math_utils.ActivationType,
+        opacity_activation_type: math_utils.ActivationType,
+        color_space: ColorSpace,
+        base_scale_on_predicted_mean: bool,
+        scale_factor: int = 1,
+    ) -> None:
+        """Initialize GaussianComposer.
+
+        Args:
+            delta_factor: Multiply delta offsets by this factor.
+            min_scale: The minimal scale factor for gaussian scale activation.
+            max_scale: The maximal scale factor for gaussian scale activation.
+            color_activation_type: Which activation function to use for colors.
+            opacity_activation_type: Which activation function to use for opacities.
+            color_space: Which color space is used in training.
+            scale_factor: The scale factor to upsample the delta_values before composition.
+            base_scale_on_predicted_mean: Whether to account z offsets for estimating base scale.
+        """
+        super().__init__()
+        self.delta_factor = delta_factor
+        self.max_scale = max_scale
+        self.min_scale = min_scale
+        self.color_activation_type = color_activation_type
+        self.opacity_activation_type = opacity_activation_type
+        self.color_space = color_space
+        self.scale_factor = scale_factor
+        self.base_scale_on_predicted_mean = base_scale_on_predicted_mean
+
+    def upsample_delta_value(self, delta: torch.Tensor, scale_factor: int = 1):
+        """Upsample the delta value.
+
+        Args:
+            delta: The delta values predicted by gaussian predictor.
+            scale_factor: The scale factor to upsample the delta_values.
+        """
+        (
+            batch_size,
+            num_channels,
+            num_layers,
+            image_height,
+            image_width,
+        ) = delta.shape
+        new_height = image_height * scale_factor
+        new_width = image_width * scale_factor
+        upsampled_delta = F.interpolate(
+            delta.view(batch_size, num_channels * num_layers, image_height, image_width),
+            scale_factor=scale_factor,
+        ).view(batch_size, num_channels, num_layers, new_height, new_width)
+        return upsampled_delta
+
+    def forward(
+        self,
+        delta: torch.Tensor,
+        base_values: GaussianBaseValues,
+        global_scale: torch.Tensor | None = None,
+        flatten_output: bool = True,
+    ) -> Gaussians3D:
+        """Combine predicted delta values with base gaussian values and apply activation function.
+
+        Args:
+            delta: The delta values predicted by gaussian predictor.
+            base_values: The gaussian base values.
+            global_scale: Global scale of Gaussians.
+            flatten_output: Flatten the gaussian parameters.
+
+        Returns:
+            The computed 3D Gaussians.
+        """
+        # Upsample the delta if delta and base_values have different strides.
+        scale_factor = self.scale_factor
+        # For triplane head, the delta has already been upsampled.
+        actual_scale_factor = base_values.mean_x_ndc.shape[-1] // delta.shape[-1]
+        if scale_factor != 1 and actual_scale_factor != 1:
+            delta = self.upsample_delta_value(delta, scale_factor)
+
+        mean_vectors = self._forward_mean(base_values, delta)
+
+        # Account for the change in base scale due to z offsets.
+        base_scales = (
+            (base_values.scales * base_values.mean_inverse_z_ndc * mean_vectors[:, 2:3, ...])
+            if self.base_scale_on_predicted_mean
+            else base_values.scales
+        )
+        singular_values = self._scale_activation(
+            base_scales,
+            delta[:, 3:6],
+            self.min_scale,
+            self.max_scale,
+        )
+        quaternions = self._quaternion_activation(base_values.quaternions, delta[:, 6:10])
+        colors = self._color_activation(base_values.colors, delta[:, 10:13])
+        opacities = self._opacity_activation(base_values.opacities, delta[:, 13])
+
+        if flatten_output:
+            # [B, C, N, H, W] -> [B, N, H, W, C].
+            # NOTE: opacities is [B, N, H, W] so it doesn't need to permute.
+            mean_vectors = mean_vectors.permute(0, 2, 3, 4, 1).flatten(1, 3)
+            singular_values = singular_values.permute(0, 2, 3, 4, 1).flatten(1, 3)
+            quaternions = quaternions.permute(0, 2, 3, 4, 1).flatten(1, 3)
+            colors = colors.permute(0, 2, 3, 4, 1).flatten(1, 3)
+            opacities = opacities.flatten(1, 3)
+
+        # Apply global scaling to convert Gaussians to metric space.
+        if global_scale is not None:
+            mean_vectors = global_scale[:, None, None] * mean_vectors
+            singular_values = global_scale[:, None, None] * singular_values
+
+        return Gaussians3D(
+            mean_vectors=mean_vectors,
+            singular_values=singular_values,
+            quaternions=quaternions,
+            colors=colors,
+            opacities=opacities,
+        )
+
+    def _forward_mean(self, base_values: GaussianBaseValues, delta: torch.Tensor) -> torch.Tensor:
+        # Concatenate base vectors and apply mean activation.
+        delta_factor = torch.tensor(
+            [self.delta_factor.xy, self.delta_factor.xy, self.delta_factor.z],
+            device=delta.device,
+        )[None, :, None, None, None]
+
+        dtype = base_values.mean_x_ndc.dtype
+        device = base_values.mean_x_ndc.device
+        target_shape = (1, 3, 1, 1, 1)
+        mean_x_mask = torch.tensor([1.0, 0.0, 0.0], dtype=dtype, device=device).reshape(
+            target_shape
+        )
+        mean_y_mask = torch.tensor([0.0, 1.0, 0.0], dtype=dtype, device=device).reshape(
+            target_shape
+        )
+        mean_z_mask = torch.tensor([0.0, 0.0, 1.0], dtype=dtype, device=device).reshape(
+            target_shape
+        )
+
+        mean_vectors_ndc = (
+            base_values.mean_x_ndc.repeat(target_shape) * mean_x_mask
+            + base_values.mean_y_ndc.repeat(target_shape) * mean_y_mask
+            + base_values.mean_inverse_z_ndc.repeat(target_shape) * mean_z_mask
+        )
+
+        mean_vectors = self._mean_activation(mean_vectors_ndc, delta_factor * delta[:, :3])
+        return mean_vectors
+
+    def _mean_activation(self, base: torch.Tensor, learned_delta: torch.Tensor) -> torch.Tensor:
+        """Mean activation function.
+
+        Args:
+            base: Tensor of shape [B, 3, H, W], where first two feature dimensions
+                (x,y) are in normalized device coordinates (NDC) where (-1, -1) is
+                the top, while the third dimension is inverse depth.
+            learned_delta: Tensor of shape [B, 3, H, W] with predicted delta values.
+
+        Returns:
+            Returns: The final mean vector after combining base and delta and applying nonlinearies.
+        """
+        xx = base[:, 0:1] + learned_delta[:, 0:1]
+        yy = base[:, 1:2] + learned_delta[:, 1:2]
+
+        a = base[:, 2:3]
+        b = learned_delta[:, 2:3]
+
+        # Original formula:
+        inverse_zz = F.softplus(math_utils.inverse_softplus(a) + b)
+        zz = 1.0 / (inverse_zz + 1e-3)
+
+        mean_vectors = torch.cat([zz * xx, zz * yy, zz], dim=1)
+        return mean_vectors
+
+    def _scale_activation(
+        self,
+        base: torch.Tensor,
+        learned_delta: torch.Tensor,
+        min_scale: float,
+        max_scale: float,
+    ) -> torch.Tensor:
+        constant_a, constant_b = _get_scale_activation_constant(max_scale, min_scale)
+        scale_factor = (max_scale - min_scale) * torch.sigmoid(
+            constant_a * self.delta_factor.scale * learned_delta + constant_b
+        ) + min_scale
+        return base * scale_factor
+
+    def _quaternion_activation(
+        self, base: torch.Tensor, learned_delta: torch.Tensor
+    ) -> torch.Tensor:
+        # No need to normalize the quaternions, since this is also done in rendering.
+        return base + self.delta_factor.quaternion * learned_delta
+
+    def _color_activation(self, base: torch.Tensor, learned_delta: torch.Tensor) -> torch.Tensor:
+        # For certain activation functions we need to clamp the base value to
+        # a supported range.
+        if self.color_activation_type == "sigmoid":
+            base = torch.clamp(base, min=0.01, max=0.99)
+        elif self.color_activation_type in ("exp", "softplus"):
+            base = torch.clamp(base, min=0.01)
+
+        activation = math_utils.create_activation_pair(self.color_activation_type)
+        colors: torch.Tensor = activation.forward(
+            activation.inverse(base) + self.delta_factor.color * learned_delta
+        )
+        # Convert gaussian color to linear if linearRGB colorspace is specified.
+        if self.color_space == "linearRGB":
+            colors = sRGB2linearRGB(colors)
+        return colors
+
+    def _opacity_activation(self, base: torch.Tensor, learned_delta: torch.Tensor) -> torch.Tensor:
+        activation = math_utils.create_activation_pair(self.opacity_activation_type)
+        return activation.forward(
+            activation.inverse(base) + self.delta_factor.opacity * learned_delta
+        )

sharp/models/decoders/__init__.py

@@ -0,0 +1,22 @@
+"""Contains different decoders for Gaussian predictor.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+from __future__ import annotations
+
+from .base_decoder import BaseDecoder
+from .monodepth_decoder import (
+    create_monodepth_decoder,
+)
+from .multires_conv_decoder import MultiresConvDecoder, UpsamplingMode
+from .unet_decoder import UNetDecoder
+
+__all__ = [
+    "BaseDecoder",
+    "UNetDecoder",
+    "MultiresConvDecoder",
+    "UpsamplingMode",
+    "create_monodepth_decoder",
+]

sharp/models/decoders/base_decoder.py

@@ -0,0 +1,21 @@
+"""Contains the base class for decoders.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+import abc
+from typing import List
+
+import torch
+from torch import nn
+
+
+class BaseDecoder(nn.Module, abc.ABC):
+    """Base decoder class."""
+
+    dim_out: int
+
+    @abc.abstractmethod
+    def forward(self, encodings: List[torch.Tensor]) -> torch.Tensor:
+        """Decode (multi-resolution) encodings."""

sharp/models/decoders/monodepth_decoder.py

@@ -0,0 +1,37 @@
+"""Contains factory function for loading/creating monodepth decoder.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+
+from __future__ import annotations
+
+from sharp.models.presets import (
+    MONODEPTH_ENCODER_DIMS_MAP,
+    ViTPreset,
+)
+
+from .multires_conv_decoder import MultiresConvDecoder
+
+
+def create_monodepth_decoder(
+    patch_encoder_preset: ViTPreset,
+    dims_decoder=None,
+) -> MultiresConvDecoder:
+    """Create DepthDensePredictionTransformer model.
+
+    Args:
+        patch_encoder_preset: The preset patch encoder architecture in SPN.
+        dims_decoder: The decoder architecture.
+    """
+    dims_encoder = MONODEPTH_ENCODER_DIMS_MAP[patch_encoder_preset]
+    if dims_decoder is None:
+        dims_decoder = dims_encoder[0]
+    if isinstance(dims_decoder, int):
+        dims_decoder = [dims_decoder]
+    decoder = MultiresConvDecoder(
+        dims_encoder=[dims_decoder[0]] + list(dims_encoder), dims_decoder=dims_decoder
+    )
+
+    return decoder

sharp/models/decoders/multires_conv_decoder.py

@@ -0,0 +1,116 @@
+"""Contains multi-res convolutional decoder.
+
+Implements the decoder for Vision Transformers for Dense Prediction, https://arxiv.org/abs/2103.13413
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+from __future__ import annotations
+
+from typing import Iterable
+
+import torch
+import torch.nn as nn
+
+from sharp.models.blocks import FeatureFusionBlock2d, UpsamplingMode
+from sharp.utils.training import checkpoint_wrapper
+
+from .base_decoder import BaseDecoder
+
+
+class MultiresConvDecoder(BaseDecoder):
+    """Decoder for multi-resolution encodings."""
+
+    def __init__(
+        self,
+        dims_encoder: Iterable[int],
+        dims_decoder: Iterable[int] | int,
+        grad_checkpointing: bool = False,
+        upsampling_mode: UpsamplingMode = "transposed_conv",
+    ):
+        """Initialize multiresolution convolutional decoder.
+
+        Args:
+            dims_encoder: Expected dims at each level from the encoder.
+            dims_decoder: Dim of decoder features.
+            grad_checkpointing: Whether to checkpoint gradient during training.
+            upsampling_mode: What method to use for upsampling.
+        """
+        super().__init__()
+        self.dims_encoder = list(dims_encoder)
+
+        if isinstance(dims_decoder, int):
+            self.dims_decoder = [dims_decoder] * len(self.dims_encoder)
+        else:
+            self.dims_decoder = list(dims_decoder)
+
+        if len(self.dims_decoder) != len(self.dims_encoder):
+            raise ValueError("Received dims_encoder and dims_decoder of different sizes.")
+
+        self.dim_out = self.dims_decoder[0]
+
+        num_encoders = len(self.dims_encoder)
+
+        # At the highest resolution, i.e. level 0, we apply projection w/ 1x1 convolution
+        # when the dimensions mismatch. Otherwise we do not do anything, which is
+        # the default behavior of monodepth.
+        conv0 = (
+            nn.Conv2d(self.dims_encoder[0], self.dims_decoder[0], kernel_size=1, bias=False)
+            if self.dims_encoder[0] != self.dims_decoder[0]
+            else nn.Identity()
+        )
+
+        convs = [conv0]
+        for i in range(1, num_encoders):
+            convs.append(
+                nn.Conv2d(
+                    self.dims_encoder[i],
+                    self.dims_decoder[i],
+                    kernel_size=3,
+                    stride=1,
+                    padding=1,
+                    bias=False,
+                )
+            )
+        self.convs = nn.ModuleList(convs)
+
+        fusions = []
+        for i in range(num_encoders):
+            fusions.append(
+                FeatureFusionBlock2d(
+                    dim_in=self.dims_decoder[i],
+                    dim_out=self.dims_decoder[i - 1] if i != 0 else self.dim_out,
+                    upsampling_mode=upsampling_mode if i != 0 else None,
+                    batch_norm=False,
+                )
+            )
+        self.fusions = nn.ModuleList(fusions)
+
+        self.grad_checkpointing = grad_checkpointing
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, is_enabled=True):
+        """Enable grad checkpointing."""
+        self.grad_checkpointing = is_enabled
+
+    def forward(self, encodings: list[torch.Tensor]) -> torch.Tensor:
+        """Decode the multi-resolution encodings."""
+        num_levels = len(encodings)
+        num_encoders = len(self.dims_encoder)
+
+        if num_levels != num_encoders:
+            raise ValueError(
+                f"Encoder output levels={num_levels} at runtime "
+                f"mismatch with expected levels={num_encoders}."
+            )
+
+        # Project features of different encoder dims to the same decoder dim.
+        # Fuse features from the lowest resolution (num_levels-1)
+        # to the highest (0).
+        features = self.convs[-1](encodings[-1])
+        features = checkpoint_wrapper(self, self.fusions[-1], features)
+        for i in range(num_levels - 2, -1, -1):
+            features_i = self.convs[i](encodings[i])
+            features = checkpoint_wrapper(self, self.fusions[i], features, features_i)
+        return features

sharp/models/decoders/unet_decoder.py

@@ -0,0 +1,113 @@
+"""Contains the UNet decoder.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+from __future__ import annotations
+
+from typing import List
+
+import torch
+from torch import nn
+
+from sharp.models.blocks import (
+    NormLayerName,
+    norm_layer_2d,
+    residual_block_2d,
+)
+
+from .base_decoder import BaseDecoder
+
+
+class UNetDecoder(BaseDecoder):
+    """Decoder of UNet model."""
+
+    def __init__(
+        self,
+        dim_out: int,
+        width: List[int] | int,
+        steps: int = 5,
+        norm_type: NormLayerName = "group_norm",
+        norm_num_groups=8,
+        blocks_per_layer=2,
+    ) -> None:
+        """Initialize UNet Decoder.
+
+        Args:
+            dim_out: The number of output channels.
+            width: Width of last input feature map from encoder
+                or the width list of all input feature maps from encoder.
+            steps: The number of upsampling steps.
+            norm_type: Which kind of normalization layer to use.
+            norm_num_groups: How many groups to use for group norm (if relevant).
+            blocks_per_layer: How many blocks per layer to use.
+        """
+        super().__init__()
+
+        if blocks_per_layer < 1:
+            raise ValueError("blocks_per_layer must be greater or equal to one.")
+
+        self.dim_out = dim_out
+
+        self.convs_up = nn.ModuleList()
+
+        self.output_dims: list[int]
+        # If only one number is specified, we assume each layer will double the channel dimension.
+        if isinstance(width, int):
+            self.input_dims = [width >> i for i in range(0, steps + 1)]
+        else:
+            self.input_dims = width[::-1][: steps + 1]
+
+        for i_step in range(steps):
+            input_width = self.input_dims[i_step]
+            current_width = self.input_dims[i_step + 1]
+            convs_up_i = nn.Sequential(
+                nn.Upsample(scale_factor=2),
+                residual_block_2d(
+                    input_width * (1 if i_step == 0 else 2),
+                    current_width,
+                    norm_type=norm_type,
+                    norm_num_groups=norm_num_groups,
+                ),
+                *[
+                    residual_block_2d(
+                        current_width,
+                        current_width,
+                        norm_type=norm_type,
+                        norm_num_groups=norm_num_groups,
+                    )
+                    for _ in range(blocks_per_layer - 1)
+                ],
+            )
+            self.convs_up.append(convs_up_i)
+            input_width = 2 * current_width
+            current_width //= 2
+
+        last_width = self.input_dims[-1]
+        self.conv_out = nn.Sequential(
+            norm_layer_2d(last_width * 2, norm_type, num_groups=norm_num_groups),
+            nn.ReLU(),
+            nn.Conv2d(last_width * 2, dim_out, 1),
+            norm_layer_2d(dim_out, norm_type, num_groups=norm_num_groups),
+            nn.ReLU(),
+        )
+
+    def forward(self, features: list[torch.Tensor]) -> torch.Tensor:
+        """Apply UNet to image.
+
+        Args:
+            features: The input multi-level feature map from encoder.
+
+        Returns:
+            The output feature map.
+        """
+        i_feature_layer = len(features) - 1
+        out = self.convs_up[0](features[i_feature_layer])
+        i_feature_layer -= 1
+        for conv_up in self.convs_up[1:]:  # type: ignore
+            out = conv_up(torch.cat([out, features[i_feature_layer]], dim=1))
+            i_feature_layer -= 1
+        out = self.conv_out(torch.cat([out, features[i_feature_layer]], dim=1))
+
+        return out

sharp/models/encoders/__init__.py

@@ -0,0 +1,24 @@
+"""Contains different encoders for Gaussian predictor.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+from sharp.models.encoders.base_encoder import BaseEncoder
+
+from .monodepth_encoder import (
+    MonodepthFeatureEncoder,
+    create_monodepth_encoder,
+)
+from .spn_encoder import SlidingPyramidNetwork
+from .unet_encoder import UNetEncoder
+from .vit_encoder import create_vit
+
+__all__ = [
+    "create_vit",
+    "BaseEncoder",
+    "UNetEncoder",
+    "SlidingPyramidNetwork",
+    "MonodepthFeatureEncoder",
+    "create_monodepth_encoder",
+]

sharp/models/encoders/base_encoder.py

@@ -0,0 +1,25 @@
+"""Contains the base class for encoders.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+import abc
+
+import torch
+from torch import nn
+
+
+class BaseEncoder(nn.Module, abc.ABC):
+    """Base encoder class."""
+
+    dim_in: int
+    output_dims: list[int]
+
+    @abc.abstractmethod
+    def forward(self, image: torch.Tensor) -> list[torch.Tensor]:
+        """Encode input image into multi-resolution encodings."""
+
+    def internal_resolution(self) -> int:
+        """Internal resolution of the encoder."""
+        return 1536

sharp/models/encoders/monodepth_encoder.py

@@ -0,0 +1,123 @@
+"""Contains Dense Transformer Prediction architecture.
+
+Implements a variant of Vision Transformers for Dense Prediction, https://arxiv.org/abs/2103.13413
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+from sharp.models.presets import (
+    MONODEPTH_ENCODER_DIMS_MAP,
+    MONODEPTH_HOOK_IDS_MAP,
+    ViTPreset,
+)
+
+from .base_encoder import BaseEncoder
+from .spn_encoder import SlidingPyramidNetwork
+from .vit_encoder import create_vit
+
+
+def create_monodepth_encoder(
+    patch_encoder_preset: ViTPreset,
+    image_encoder_preset: ViTPreset,
+    use_patch_overlap: bool = True,
+    last_encoder: int = 256,
+) -> SlidingPyramidNetwork:
+    """Creates DepthDensePredictionTransformer model.
+
+    Args:
+        patch_encoder_preset: The preset patch encoder architecture in SPN.
+        image_encoder_preset: The preset image encoder architecture in SPN.
+        use_patch_overlap: Whether to use overlap between patches in SPN.
+        last_encoder: last number of encoder features.
+    """
+    dims_encoder = [last_encoder] + MONODEPTH_ENCODER_DIMS_MAP[patch_encoder_preset]
+    patch_encoder_block_ids = MONODEPTH_HOOK_IDS_MAP[patch_encoder_preset]
+
+    patch_encoder = create_vit(
+        preset=patch_encoder_preset,
+        intermediate_features_ids=patch_encoder_block_ids,
+        # We always need to output intermediate features for assembly.
+    )
+    image_encoder = create_vit(
+        preset=image_encoder_preset,
+        intermediate_features_ids=None,
+    )
+
+    encoder = SlidingPyramidNetwork(
+        dims_encoder=dims_encoder,
+        patch_encoder=patch_encoder,
+        image_encoder=image_encoder,
+        use_patch_overlap=use_patch_overlap,
+    )
+
+    return encoder
+
+
+class ProjectionModule(nn.Module):
+    """Apply projection of features."""
+
+    def __init__(self, dims_in: list[int], dims_out: list[int]) -> None:
+        """Initialize projection module."""
+        super().__init__()
+        if len(dims_in) != len(dims_out):
+            raise ValueError("Length of dims_in must be same as length of dims_out.")
+        self.convs = nn.ModuleList(
+            [nn.Conv2d(dim_in, dim_out, 1) for dim_in, dim_out in zip(dims_in, dims_out)]
+        )
+
+    def forward(self, encodings: list[torch.Tensor]) -> list[torch.Tensor]:
+        """Apply projection module."""
+        if len(encodings) != len(self.convs):
+            raise ValueError("Number of encodings must be equal to number of projections.")
+        return [conv(encoding) for conv, encoding in zip(self.convs, encodings)]
+
+
+class MonodepthFeatureEncoder(BaseEncoder):
+    """A wrapper around monodepth network to extract features."""
+
+    def __init__(
+        self,
+        monodepth_encoder: SlidingPyramidNetwork,
+        output_dims: list[int] | None = None,
+        freeze_projection: bool = False,
+    ) -> None:
+        """Initialize MonodepthFeatureExtractor."""
+        super().__init__()
+
+        self.encoder = monodepth_encoder
+
+        # The monodepth network returns two feature maps for the first entry in
+        # backbone.encoder.dims_encoder.
+        monodepth_dims = self.encoder.dims_encoder
+        monodepth_dims = monodepth_dims
+
+        if output_dims is not None:
+            if not len(output_dims) == len(monodepth_dims):
+                raise ValueError(
+                    "When set, number of output dimensions must be equal to output "
+                    f"dimensions of monodepth model {len(monodepth_dims)}."
+                )
+
+            self.projection = ProjectionModule(monodepth_dims, output_dims)
+            self.output_dims = output_dims
+        else:
+            self.projection = nn.Identity()
+            self.output_dims = monodepth_dims
+
+        if freeze_projection:
+            self.projection.requires_grad_(False)
+
+    def forward(self, input_features: torch.Tensor) -> list[torch.Tensor]:
+        """Extract multi-resolution features."""
+        encodings = self.encoder(input_features[:, :3].contiguous())
+        return self.projection(encodings)
+
+    def internal_resolution(self) -> int:
+        """Internal resolution of the encoder."""
+        return self.encoder.internal_resolution()

sharp/models/encoders/spn_encoder.py

@@ -0,0 +1,369 @@
+"""Contains Sliding Pyramid Network architecture.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Iterable
+
+import torch
+import torch.fx
+import torch.nn as nn
+import torch.nn.functional as F
+
+from sharp.utils.training import checkpoint_wrapper
+
+from .base_encoder import BaseEncoder
+from .vit_encoder import TimmViT
+
+# torch.fx.wrap is used here to mark functions as leaf nodes during symbolic tracing
+# ensuring they are not traced but seen as atomic operation. In short, symbolic tracing
+# struggles with native python functions and conditional flows.
+non_traceable_ops = ("len", "int")
+for op in non_traceable_ops:
+    torch.fx.wrap(op)
+
+
+class SlidingPyramidNetwork(BaseEncoder):
+    """Sliding Pyramid Network.
+
+    An encoder aimed at creating multi-resolution encodings from Vision Transformers.
+
+    Reference: Bochkovskii et al. - "Depth pro: Sharp monocular metric depth in less
+               than a second." (ICLR 2024)
+    """
+
+    def __init__(
+        self,
+        dims_encoder: Iterable[int],
+        patch_encoder: TimmViT,
+        image_encoder: TimmViT,
+        use_patch_overlap: bool = True,
+    ):
+        """Initialize Sliding Pyramid Network.
+
+        The framework
+            1. creates an image pyramid,
+            2. generates overlapping patches with a sliding window at each pyramid level,
+            3. creates batched encodings via vision transformer backbones,
+            4. produces multi-resolution encodings.
+
+        Args:
+            dims_encoder: Dimensions of the encoder at different layers.
+            patch_encoder: Backbone used for highres part of the pyramid.
+            image_encoder: Backbone used for lowres part of the pyramid.
+            use_patch_overlap: Whether to use overlap between patches in SPN.
+        """
+        super().__init__()
+
+        self.dim_in = patch_encoder.dim_in
+
+        self.dims_encoder = list(dims_encoder)
+        self.patch_encoder = patch_encoder
+        self.image_encoder = image_encoder
+
+        base_embed_dim = patch_encoder.embed_dim
+        lowres_embed_dim = image_encoder.embed_dim
+        self.patch_size = patch_encoder.internal_resolution()
+
+        self.grad_checkpointing = False
+        self.use_patch_overlap = use_patch_overlap
+
+        # Retrieve intermediate feature ids registered in create_monodepth_encoder.
+        self.patch_intermediate_features_ids = patch_encoder.intermediate_features_ids
+        if (
+            not isinstance(self.patch_intermediate_features_ids, list)
+            or not len(self.patch_intermediate_features_ids) == 4
+        ):
+            raise ValueError("Patch intermediate feature ids must be a 4-item list.")
+
+        self.image_intermediate_features_ids = image_encoder.intermediate_features_ids
+
+        def _create_project_upsample_block(
+            dim_in: int,
+            dim_out: int,
+            upsample_layers: int,
+            dim_intermediate=None,
+        ) -> nn.Module:
+            if dim_intermediate is None:
+                dim_intermediate = dim_out
+            # Projection.
+            blocks = [
+                nn.Conv2d(
+                    in_channels=dim_in,
+                    out_channels=dim_intermediate,
+                    kernel_size=1,
+                    stride=1,
+                    padding=0,
+                    bias=False,
+                )
+            ]
+
+            # Upsampling.
+            blocks += [
+                nn.ConvTranspose2d(
+                    in_channels=dim_intermediate if i == 0 else dim_out,
+                    out_channels=dim_out,
+                    kernel_size=2,
+                    stride=2,
+                    padding=0,
+                    bias=False,
+                )
+                for i in range(upsample_layers)
+            ]
+
+            return nn.Sequential(*blocks)
+
+        self.upsample_latent0 = _create_project_upsample_block(
+            dim_in=base_embed_dim,
+            dim_out=self.dims_encoder[0],
+            upsample_layers=3,
+            dim_intermediate=self.dims_encoder[1],
+        )
+        self.upsample_latent1 = _create_project_upsample_block(
+            dim_in=base_embed_dim, dim_out=self.dims_encoder[1], upsample_layers=2
+        )
+
+        self.upsample0 = _create_project_upsample_block(
+            dim_in=base_embed_dim, dim_out=self.dims_encoder[2], upsample_layers=1
+        )
+        self.upsample1 = _create_project_upsample_block(
+            dim_in=base_embed_dim, dim_out=self.dims_encoder[3], upsample_layers=1
+        )
+        self.upsample2 = _create_project_upsample_block(
+            dim_in=base_embed_dim, dim_out=self.dims_encoder[4], upsample_layers=1
+        )
+
+        self.upsample_lowres = nn.ConvTranspose2d(
+            in_channels=lowres_embed_dim,
+            out_channels=self.dims_encoder[4],
+            kernel_size=2,
+            stride=2,
+            padding=0,
+            bias=True,
+        )
+        self.fuse_lowres = nn.Conv2d(
+            in_channels=(self.dims_encoder[4] + self.dims_encoder[4]),
+            out_channels=self.dims_encoder[4],
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            bias=True,
+        )
+
+    def internal_resolution(self) -> int:
+        """Return the full image size of the SPN network."""
+        return self.patch_size * 4
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, is_enabled=True):
+        """Enable grad checkpointing."""
+        self.grad_checkpointing = is_enabled
+        self.patch_encoder.set_grad_checkpointing(is_enabled)
+        self.image_encoder.set_grad_checkpointing(is_enabled)
+
+    @torch.jit.ignore
+    def set_requires_grad_(self, patch_encoder: bool, image_encoder: bool):
+        """Set requires grad for separate components."""
+        self.patch_encoder.requires_grad_(patch_encoder)
+        self.image_encoder.requires_grad_(image_encoder)
+
+        # Always freeze the unused TimmViT head to exclude it from the calculation of
+        # trainable parameters.
+        self.patch_encoder.head.requires_grad_(False)
+        self.image_encoder.head.requires_grad_(False)
+
+        # These upsamplers only affect patch encoder's feature maps.
+        self.upsample_latent0.requires_grad_(patch_encoder)
+        self.upsample_latent1.requires_grad_(patch_encoder)
+        self.upsample0.requires_grad_(patch_encoder)
+        self.upsample1.requires_grad_(patch_encoder)
+        self.upsample2.requires_grad_(patch_encoder)
+
+        # This upsampler affects only image encoder's feature map.
+        self.upsample_lowres.requires_grad_(image_encoder)
+
+        # This fuser affects both image and patch encoders.
+        self.fuse_lowres.requires_grad_(image_encoder or patch_encoder)
+
+    def _create_pyramid(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Creates a 3-level image pyramid."""
+        # Original resolution: 1536 by default.
+        x0 = x
+
+        # Middle resolution: 768 by default.
+        x1 = F.interpolate(x, size=None, scale_factor=0.5, mode="bilinear", align_corners=False)
+
+        # Low resolution: 384 by default, corresponding to the backbone resolution.
+        x2 = F.interpolate(x, size=None, scale_factor=0.25, mode="bilinear", align_corners=False)
+
+        return x0, x1, x2
+
+    def forward(self, x: torch.Tensor) -> list[torch.Tensor]:
+        """Encode input at multiple resolutions."""
+        batch_size = x.shape[0]
+
+        # Step 0: create a 3-level image pyramid.
+        x0, x1, x2 = self._create_pyramid(x)
+
+        if self.use_patch_overlap:
+            # Step 1: split to create batched overlapped mini-images at the ViT
+            # resolution.
+            # 5x5 @ 384x384 at the highest resolution (1536x1536).
+            x0_patches = split(x0, overlap_ratio=0.25, patch_size=self.patch_size)
+            # 3x3 @ 384x384 at the middle resolution (768x768).
+            x1_patches = split(x1, overlap_ratio=0.5, patch_size=self.patch_size)
+            # 1x1 # 384x384 at the lowest resolution (384x384).
+            x2_patches = x2
+            padding = 3
+        else:
+            # Step 1: split to create batched overlapped mini-images at the ViT
+            # resolution.
+            # 4x4 @ 384x384 at the highest resolution (1536x1536).
+            x0_patches = split(x0, overlap_ratio=0.0, patch_size=self.patch_size)
+            # 2x2 @ 384x384 at the middle resolution (768x768).
+            x1_patches = split(x1, overlap_ratio=0.0, patch_size=self.patch_size)
+            # 1x1 # 384x384 at the lowest resolution (384x384).
+            x2_patches = x2
+            padding = 0
+        x0_tile_size = x0_patches.shape[0]
+
+        # Concatenate all the sliding window patches and form a batch of size
+        # (35=5x5+3x3+1x1) or (21=4x4+2x2+1x1).
+        x_pyramid_patches = torch.cat(
+            (x0_patches, x1_patches, x2_patches),
+            dim=0,
+        )
+
+        # Run the ViT model and get the result of large batch size.
+        #
+        # For the retrieval of intermediate features forward hooks are more concise,
+        # but they are not well compatible with symbolic tracing because attributes
+        # of submodules can be lost during tracing. Therefore, forward hooks may not
+        # be preserved during graph transformation, leading to unexpected behavior.
+        # To avoid such issues it is safer not to use them because they are not
+        # essential here.
+        x_pyramid_encodings, patch_intermediate_features = self.patch_encoder(x_pyramid_patches)
+
+        # Step 3: merging.
+        # Merge highres latent encoding.
+        # NOTE: list type check has completed in init.
+        x_latent0_encodings = self.patch_encoder.reshape_feature(
+            patch_intermediate_features[self.patch_intermediate_features_ids[0]]  # type:ignore[index]
+        )
+        x_latent0_features = merge(
+            x_latent0_encodings[: batch_size * x0_tile_size],
+            batch_size=batch_size,
+            padding=padding,
+        )
+
+        x_latent1_encodings = self.patch_encoder.reshape_feature(
+            patch_intermediate_features[self.patch_intermediate_features_ids[1]]  # type:ignore[index]
+        )
+        x_latent1_features = merge(
+            x_latent1_encodings[: batch_size * x0_tile_size],
+            batch_size=batch_size,
+            padding=padding,
+        )
+
+        # Split the 35 batch size from pyramid encoding back into 5x5+3x3+1x1.
+        x0_encodings, x1_encodings, x2_encodings = torch.split(
+            x_pyramid_encodings,
+            [len(x0_patches), len(x1_patches), len(x2_patches)],
+            dim=0,
+        )
+
+        # 96x96 feature maps by merging 5x5 @ 24x24 patches with overlaps.
+        x0_features = merge(x0_encodings, batch_size=batch_size, padding=padding)
+
+        # 48x84 feature maps by merging 3x3 @ 24x24 patches with overlaps.
+        x1_features = merge(x1_encodings, batch_size=batch_size, padding=2 * padding)
+
+        # 24x24 feature maps.
+        x2_features = x2_encodings
+
+        # Apply the image encoder.
+        x_lowres_features, image_intermediate_features = self.image_encoder(x2_patches)
+
+        # Upsample feature maps.
+        x_latent0_features = checkpoint_wrapper(self, self.upsample_latent0, x_latent0_features)
+        x_latent1_features = checkpoint_wrapper(self, self.upsample_latent1, x_latent1_features)
+
+        x0_features = checkpoint_wrapper(self, self.upsample0, x0_features)
+        x1_features = checkpoint_wrapper(self, self.upsample1, x1_features)
+        x2_features = checkpoint_wrapper(self, self.upsample2, x2_features)
+
+        x_lowres_features = checkpoint_wrapper(self, self.upsample_lowres, x_lowres_features)
+        x_lowres_features = checkpoint_wrapper(
+            self, self.fuse_lowres, torch.cat((x2_features, x_lowres_features), dim=1)
+        )
+
+        output = [
+            x_latent0_features,
+            x_latent1_features,
+            x0_features,
+            x1_features,
+            x_lowres_features,
+        ]
+
+        return output
+
+
+# It seems that torch.fx.wrap can only be applied to functions, not methods.
+# Hence, split and merge were converted into functions to be marked as atomic
+# operations for symbolic tracing.
+@torch.fx.wrap
+def split(image: torch.Tensor, overlap_ratio: float = 0.25, patch_size: int = 384) -> torch.Tensor:
+    """Split the input into small patches with sliding window."""
+    patch_stride = int(patch_size * (1 - overlap_ratio))
+
+    image_size = image.shape[-1]
+    steps = int(math.ceil((image_size - patch_size) / patch_stride)) + 1
+
+    x_patch_list = []
+    for j in range(steps):
+        j0 = j * patch_stride
+        j1 = j0 + patch_size
+
+        for i in range(steps):
+            i0 = i * patch_stride
+            i1 = i0 + patch_size
+            x_patch_list.append(image[..., j0:j1, i0:i1])
+
+    return torch.cat(x_patch_list, dim=0)
+
+
+# Decorator marking function as an atomic operator for symbolic tracing.
+@torch.fx.wrap
+def merge(image_patches: torch.Tensor, batch_size: int, padding: int = 3) -> torch.Tensor:
+    """Merge the patched input into a image with sliding window."""
+    steps = int(math.sqrt(image_patches.shape[0] // batch_size))
+
+    idx = 0
+
+    output_list = []
+    for j in range(steps):
+        output_row_list = []
+        for i in range(steps):
+            output = image_patches[batch_size * idx : batch_size * (idx + 1)]
+
+            if padding != 0:
+                if j != 0:
+                    output = output[..., padding:, :]
+                if i != 0:
+                    output = output[..., :, padding:]
+                if j != steps - 1:
+                    output = output[..., :-padding, :]
+                if i != steps - 1:
+                    output = output[..., :, :-padding]
+
+            output_row_list.append(output)
+            idx += 1
+
+        output_row = torch.cat(output_row_list, dim=-1)
+        output_list.append(output_row)
+    output = torch.cat(output_list, dim=-2)
+    return output

sharp/models/encoders/unet_encoder.py

@@ -0,0 +1,117 @@
+"""Contains backbone models for feature extraction from RGBD input.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+from __future__ import annotations
+
+from typing import List
+
+import torch
+from torch import nn
+
+from sharp.models.blocks import (
+    NormLayerName,
+    norm_layer_2d,
+    residual_block_2d,
+)
+
+from .base_encoder import BaseEncoder
+
+
+class UNetEncoder(BaseEncoder):
+    """Encoder of UNet model."""
+
+    def __init__(
+        self,
+        dim_in: int,
+        width: List[int] | int,
+        steps: int = 6,
+        norm_type: NormLayerName = "group_norm",
+        norm_num_groups=8,
+        blocks_per_layer=2,
+    ) -> None:
+        """Initialize UNet Encoder.
+
+        Args:
+            dim_in: The number of input channels.
+            width: Width multiplicator of intermediate layers or the width list of all layers.
+            steps: The number of downsampling steps.
+            norm_type: Which kind of normalization layer to use.
+            norm_num_groups: How many groups to use for group norm (if relevant).
+            blocks_per_layer: How many residual blocks per layer to use.
+        """
+        super().__init__()
+
+        if blocks_per_layer < 1:
+            raise ValueError("blocks_per_layer must be greater or equal to one.")
+
+        self.dim_in = dim_in
+        self.width = width
+        self.num_steps = steps
+
+        self.convs_down = nn.ModuleList()
+
+        self.output_dims: list[int]
+        # If only one number is specified, we assume each layer will double the channel dimension.
+        if isinstance(width, int):
+            self.output_dims = [width << i for i in range(0, steps + 1)]
+        else:
+            if len(width) != (steps + 1):
+                raise ValueError("Length of width should match the steps for UNetEncoder.")
+            self.output_dims = width
+
+        self.conv_in = nn.Sequential(
+            nn.Conv2d(self.dim_in, self.output_dims[0], 3, stride=1, padding=1),
+            norm_layer_2d(self.output_dims[0], norm_type, num_groups=norm_num_groups),
+            nn.ReLU(),
+        )
+
+        for i_step in range(steps):
+            input_width = self.output_dims[i_step]
+            current_width = self.output_dims[i_step + 1]
+            convs_down_i = nn.Sequential(
+                nn.AvgPool2d(2, stride=2),
+                residual_block_2d(
+                    input_width,
+                    current_width,
+                    norm_type=norm_type,
+                    norm_num_groups=norm_num_groups,
+                ),
+                *[
+                    residual_block_2d(
+                        current_width,
+                        current_width,
+                        norm_type=norm_type,
+                        norm_num_groups=norm_num_groups,
+                    )
+                    for _ in range(blocks_per_layer - 1)
+                ],
+            )
+            self.convs_down.append(convs_down_i)
+
+    def forward(self, input: torch.Tensor) -> list[torch.Tensor]:
+        """Apply UNet Encoder to image.
+
+        Args:
+            input: The input image.
+
+        Returns:
+            The output multi-level feature map from encoder.
+        """
+        features = []
+
+        feat_i = self.conv_in(input)
+        features.append(feat_i)
+
+        for conv_down in self.convs_down:
+            feat_i = conv_down(feat_i)
+            features.append(feat_i)
+
+        return features
+
+    @property
+    def out_width(self) -> int:
+        """Compute the output width for UNet decoder."""
+        return self.output_dims[-1]

sharp/models/encoders/vit_encoder.py

@@ -0,0 +1,111 @@
+"""Contains factory functions to build and load ViT.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+from __future__ import annotations
+
+import logging
+
+import timm
+import torch
+
+from sharp.models.presets.vit import VIT_CONFIG_DICT, ViTConfig, ViTPreset
+
+LOGGER = logging.getLogger(__name__)
+
+
+class TimmViT(timm.models.VisionTransformer):
+    """Contains TIMM implementation for Vanilla ViT."""
+
+    def __init__(self, config: ViTConfig):
+        """Initialize ViT from TIMM implementation."""
+        # Handle mlp layers.
+        mlp_layer = timm.layers.GluMlp if config.mlp_mode == "glu" else timm.layers.Mlp
+
+        super().__init__(
+            in_chans=config.in_chans,
+            embed_dim=config.embed_dim,
+            depth=config.depth,
+            num_heads=config.num_heads,
+            init_values=config.init_values,
+            img_size=config.img_size,
+            patch_size=config.patch_size,
+            num_classes=config.num_classes,
+            mlp_ratio=config.mlp_ratio,
+            qkv_bias=config.qkv_bias,
+            global_pool=config.global_pool,
+            mlp_layer=mlp_layer,
+        )
+
+        # Required for extracting intermediate features.
+        self.dim_in = config.in_chans
+        self.intermediate_features_ids = config.intermediate_features_ids
+
+    def reshape_feature(self, embeddings: torch.Tensor):
+        """Discard class token and reshape 1D feature map to a 2D grid."""
+        batch_size, seq_len, channel = embeddings.shape
+
+        height, width = self.patch_embed.grid_size
+
+        # Remove class token.
+        if self.num_prefix_tokens:
+            embeddings = embeddings[:, self.num_prefix_tokens :, :]
+
+        # Shape: (batch, height, width, dim) -> (batch, dim, height, width)
+        embeddings = embeddings.reshape(batch_size, height, width, channel).permute(0, 3, 1, 2)
+        return embeddings
+
+    def forward(self, input_tensor: torch.Tensor) -> tuple[torch.Tensor, dict[int, torch.Tensor]]:
+        """Override forwarding with intermediate features.
+
+        Adapted from timm ViT.
+
+        Returns:
+            Output features and list of features from intermediate layers (patch encoder only).
+        """
+        intermediate_features = {}
+
+        x = self.patch_embed(input_tensor)
+        batch_size, seq_len, _ = x.shape
+
+        x = self._pos_embed(x)
+        x = self.patch_drop(x)
+        x = self.norm_pre(x)
+
+        for idx, block in enumerate(self.blocks):
+            x = block(x)
+            if self.intermediate_features_ids is not None and idx in self.intermediate_features_ids:
+                intermediate_features[idx] = x
+        x = self.norm(x)
+
+        x = self.reshape_feature(x)
+        return x, intermediate_features
+
+    def internal_resolution(self) -> int:
+        """Return the internal image size of the network."""
+        if isinstance(self.patch_embed.img_size, tuple):
+            return self.patch_embed.img_size[0]
+        else:
+            return self.patch_embed.img_size
+
+
+def create_vit(
+    config: ViTConfig | None = None,
+    preset: ViTPreset | None = "dinov2l16_384",
+    intermediate_features_ids: list[int] | None = None,
+) -> TimmViT:
+    """Factory function for creating a ViT model."""
+    if config is not None:
+        LOGGER.info("Using user-defined config.")
+    else:
+        if preset is None:
+            raise ValueError("User-defined config and preset cannot be both None.")
+        LOGGER.info("Using preset ViT %s.", preset)
+        config = VIT_CONFIG_DICT[preset]
+
+    config.intermediate_features_ids = intermediate_features_ids
+    model = TimmViT(config)
+    LOGGER.debug(model)
+    return model

sharp/models/gaussian_decoder.py

@@ -0,0 +1,267 @@
+"""Contains Dense Transformer Prediction architecture.
+
+Implements a variant of Vision Transformers for Dense Prediction, https://arxiv.org/abs/2103.13413
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+from __future__ import annotations
+
+from typing import NamedTuple
+
+import torch
+import torch.nn as nn
+
+from sharp.models.blocks import (
+    FeatureFusionBlock2d,
+    NormLayerName,
+    residual_block_2d,
+)
+from sharp.models.decoders import BaseDecoder, MultiresConvDecoder
+from sharp.models.params import DPTImageEncoderType, GaussianDecoderParams
+
+
+def create_gaussian_decoder(
+    params: GaussianDecoderParams, dims_depth_features: list[int]
+) -> GaussianDensePredictionTransformer:
+    """Create gaussian_decoder model specified by gaussian_decoder_name."""
+    decoder = MultiresConvDecoder(
+        dims_depth_features,
+        params.dims_decoder,
+        grad_checkpointing=params.grad_checkpointing,
+        upsampling_mode=params.upsampling_mode,
+    )
+
+    return GaussianDensePredictionTransformer(
+        decoder=decoder,
+        dim_in=params.dim_in,
+        dim_out=params.dim_out,
+        stride_out=params.stride,
+        norm_type=params.norm_type,
+        norm_num_groups=params.norm_num_groups,
+        use_depth_input=params.use_depth_input,
+        grad_checkpointing=params.grad_checkpointing,
+        image_encoder_type=params.image_encoder_type,
+        image_encoder_params=params,
+    )
+
+
+def _create_project_upsample_block(
+    dim_in: int,
+    dim_out: int,
+    upsample_layers: int,
+    dim_intermediate: int | None = None,
+) -> nn.Module:
+    if dim_intermediate is None:
+        dim_intermediate = dim_out
+    # Projection.
+    blocks = [
+        nn.Conv2d(
+            in_channels=dim_in,
+            out_channels=dim_intermediate,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            bias=False,
+        )
+    ]
+
+    # Upsampling.
+    blocks += [
+        nn.ConvTranspose2d(
+            in_channels=dim_intermediate if i == 0 else dim_out,
+            out_channels=dim_out,
+            kernel_size=2,
+            stride=2,
+            padding=0,
+            bias=False,
+        )
+        for i in range(upsample_layers)
+    ]
+
+    return nn.Sequential(*blocks)
+
+
+class ImageFeatures(NamedTuple):
+    """Image feature extracted from decoder."""
+
+    texture_features: torch.Tensor
+    geometry_features: torch.Tensor
+
+
+class SkipConvBackbone(nn.Module):
+    """A wrapper around a conv layer that behaves like a BaseBackbone."""
+
+    def __init__(self, dim_in: int, dim_out: int, kernel_size: int, stride_out: int):
+        """Initialize SkipConvBackbone."""
+        super().__init__()
+        self.stride_out = stride_out
+        if stride_out == 1 and kernel_size != 1:
+            raise ValueError("We only support kernel_size = 1 if stride_out is 1.")
+        padding: int = (kernel_size - 1) // 2
+        self.conv = nn.Conv2d(
+            dim_in, dim_out, kernel_size=kernel_size, stride=stride_out, padding=padding
+        )
+
+    def forward(
+        self,
+        input_features: torch.Tensor,
+        encodings: list[torch.Tensor] | None = None,
+    ) -> ImageFeatures:
+        """Apply SkipConvBackbone to image."""
+        output = self.conv(input_features)
+        return ImageFeatures(
+            texture_features=output,
+            geometry_features=output,
+        )
+
+    @property
+    def stride(self) -> int:
+        """Effective downsampling stride."""
+        return self.stride_out
+
+
+class GaussianDensePredictionTransformer(nn.Module):
+    """Dense Prediction Transformer for Gaussian.
+
+    Reuse monodepth decoded features for processing.
+    """
+
+    norm_type: NormLayerName
+
+    def __init__(
+        self,
+        decoder: BaseDecoder,
+        dim_in: int,
+        dim_out: int,
+        stride_out: int,
+        image_encoder_params: GaussianDecoderParams,
+        image_encoder_type: DPTImageEncoderType = "skip_conv",
+        norm_type: NormLayerName = "group_norm",
+        norm_num_groups: int = 8,
+        use_depth_input: bool = True,
+        grad_checkpointing: bool = False,
+    ):
+        """Initialize Dense Prediction Transformer for Gaussian.
+
+        Args:
+            decoder: Decoder to decode features.
+            monodepth_decoder: Optional monodepth decoder to fuse monodepth decoded features.
+            dim_in: Input dimension.
+            dim_out: Final output dimension.
+            stride_out: Stride of output feature map.
+            image_encoder_params: The backbone parameters to configurate the image encoder.
+            image_encoder_type: Type of image encoder to use.
+            encoder: Encoder to generate features using monodepth model.
+            norm_type: Type of norm layers.
+            norm_num_groups: Num groups for norm layers.
+            use_depth_input: Whether to use depth input.
+            grad_checkpointing: Whether to use gradient checkpointing.
+        """
+        super().__init__()
+
+        self.decoder = decoder
+        self.dim_in = dim_in
+        self.dim_out = dim_out
+        self.stride_out = stride_out
+        self.norm_type = norm_type
+        self.norm_num_groups = norm_num_groups
+        self.use_depth_input = use_depth_input
+        self.grad_checkpointing = grad_checkpointing
+        self.image_encoder_type = image_encoder_type
+
+        # Adopt an image encoder to lift dimension to monodepth feature and
+        # resize to be the same resolution as the decoder output.
+        dim_in = self.dim_in if use_depth_input else self.dim_in - 1
+        image_encoder_params.dim_in = dim_in
+        image_encoder_params.dim_out = decoder.dim_out
+        self.image_encoder = self._create_image_encoder(image_encoder_params, stride_out)
+
+        self.fusion = FeatureFusionBlock2d(decoder.dim_out)
+
+        if stride_out == 1:
+            self.upsample = _create_project_upsample_block(
+                decoder.dim_out,
+                decoder.dim_out,
+                upsample_layers=1,
+            )
+        elif stride_out == 2:
+            self.upsample = nn.Identity()
+        else:
+            raise ValueError("We only support stride is 1 or 2 for DPT backbone.")
+
+        self.texture_head = self._create_head(dim_decoder=decoder.dim_out, dim_out=self.dim_out)
+        self.geometry_head = self._create_head(dim_decoder=decoder.dim_out, dim_out=self.dim_out)
+
+    def _create_head(self, dim_decoder: int, dim_out: int) -> nn.Module:
+        return nn.Sequential(
+            residual_block_2d(
+                dim_in=dim_decoder,
+                dim_out=dim_decoder,
+                dim_hidden=dim_decoder // 2,
+                norm_type=self.norm_type,
+                norm_num_groups=self.norm_num_groups,
+            ),
+            residual_block_2d(
+                dim_in=dim_decoder,
+                dim_hidden=dim_decoder // 2,
+                dim_out=dim_decoder,
+                norm_type=self.norm_type,
+                norm_num_groups=self.norm_num_groups,
+            ),
+            nn.ReLU(),
+            nn.Conv2d(dim_decoder, dim_out, kernel_size=1, stride=1),
+            nn.ReLU(),
+        )
+
+    def _create_image_encoder(
+        self, image_encoder_params: GaussianDecoderParams, stride_out: int
+    ) -> nn.Module:
+        """Create image encoder and return based on parameters."""
+        if self.image_encoder_type == "skip_conv":
+            # Use kernel_size = 1 only if stride_out is 1.
+            return SkipConvBackbone(
+                image_encoder_params.dim_in,
+                image_encoder_params.dim_out,
+                kernel_size=3 if stride_out != 1 else 1,
+                stride_out=stride_out,
+            )
+        elif self.image_encoder_type == "skip_conv_kernel2":
+            return SkipConvBackbone(
+                image_encoder_params.dim_in,
+                image_encoder_params.dim_out,
+                kernel_size=stride_out,
+                stride_out=stride_out,
+            )
+        else:
+            raise ValueError(f"Unsupported image encoder type: {self.image_encoder_type}")
+
+    def forward(self, input_features: torch.Tensor, encodings: list[torch.Tensor]) -> ImageFeatures:
+        """Run monodepth and fuse features with input image to predict Gaussians.
+
+        Args:
+            input_features: The input features to use.
+            encodings: Feature encodings (e.g. from monodepth network).
+        """
+        features = self.decoder(encodings).contiguous()
+        features = self.upsample(features)
+
+        if self.use_depth_input:
+            skip_features = self.image_encoder(input_features).texture_features
+        else:
+            skip_features = self.image_encoder(input_features[:, :3].contiguous())
+        features = self.fusion(features, skip_features)
+
+        texture_features = self.texture_head(features)
+        geometry_features = self.geometry_head(features)
+
+        return ImageFeatures(
+            texture_features=texture_features,  # type: ignore
+            geometry_features=geometry_features,  # type: ignore
+        )
+
+    @property
+    def stride(self) -> int:
+        """Internal stride of GaussianDensePredictionTransformer."""
+        return self.stride_out

sharp/models/heads.py

@@ -0,0 +1,53 @@
+"""Contains decoder head for direct prediction of delta values.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+from .gaussian_decoder import ImageFeatures
+
+
+class DirectPredictionHead(nn.Module):
+    """Decodes features into delta values using convolutions."""
+
+    def __init__(self, feature_dim: int, num_layers: int) -> None:
+        """Initialize DirectGaussianPredictor.
+
+        Args:
+            feature_dim: Number of input features.
+            num_layers: The number of layers of Gaussians to predict.
+        """
+        super().__init__()
+        self.num_layers = num_layers
+
+        # 14 is 3 means, 3 scales, 4 quaternions, 3 colors and 1 opacity
+        self.geometry_prediction_head = nn.Conv2d(feature_dim, 3 * num_layers, 1)
+        self.geometry_prediction_head.weight.data.zero_()
+        assert self.geometry_prediction_head.bias is not None
+        self.geometry_prediction_head.bias.data.zero_()
+
+        self.texture_prediction_head = nn.Conv2d(feature_dim, (14 - 3) * num_layers, 1)
+        self.texture_prediction_head.weight.data.zero_()
+        assert self.texture_prediction_head.bias is not None
+        self.texture_prediction_head.bias.data.zero_()
+
+    def forward(self, image_features: ImageFeatures) -> torch.Tensor:
+        """Predict deltas for 3D Gaussians.
+
+        Args:
+            image_features: Image features from decoder.
+
+        Returns:
+            The predicted deltas for Gaussian attributes.
+        """
+        delta_values_geometry = self.geometry_prediction_head(image_features.geometry_features)
+        delta_values_texture = self.texture_prediction_head(image_features.texture_features)
+        delta_values_geometry = delta_values_geometry.unflatten(1, (3, self.num_layers))
+        delta_values_texture = delta_values_texture.unflatten(1, (14 - 3, self.num_layers))
+        delta_values = torch.cat([delta_values_geometry, delta_values_texture], dim=1)
+        return delta_values

sharp/models/initializer.py

@@ -0,0 +1,297 @@
+"""Contains modules to initialize Gaussians from RGBD.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+from __future__ import annotations
+
+from typing import NamedTuple
+
+import torch
+from torch import nn
+
+from .params import ColorInitOption, DepthInitOption, InitializerParams
+
+
+def create_initializer(params: InitializerParams) -> nn.Module:
+    """Create inpainter."""
+    return MultiLayerInitializer(
+        num_layers=params.num_layers,
+        stride=params.stride,
+        base_depth=params.base_depth,
+        scale_factor=params.scale_factor,
+        disparity_factor=params.disparity_factor,
+        color_option=params.color_option,
+        first_layer_depth_option=params.first_layer_depth_option,
+        rest_layer_depth_option=params.rest_layer_depth_option,
+        normalize_depth=params.normalize_depth,
+        feature_input_stop_grad=params.feature_input_stop_grad,
+    )
+
+
+class GaussianBaseValues(NamedTuple):
+    """Base values for gaussian predictor.
+
+    We predict x and y in normalized device coordinates (NDC) where (-1, -1) is the top
+    left corner and (1, 1) the bottom right corner. The last component of
+    mean_vectors_ndc is inverse depth.
+    """
+
+    mean_x_ndc: torch.Tensor
+    mean_y_ndc: torch.Tensor
+    mean_inverse_z_ndc: torch.Tensor
+
+    scales: torch.Tensor
+    quaternions: torch.Tensor
+    colors: torch.Tensor
+    opacities: torch.Tensor
+
+
+class InitializerOutput(NamedTuple):
+    """Output of initializer."""
+
+    # Gaussian base values.
+    gaussian_base_values: GaussianBaseValues
+
+    # Feature input to the Gaussian predictor.
+    feature_input: torch.Tensor
+
+    # Global scale to unscale output.
+    global_scale: torch.Tensor | None = None
+
+
+class MultiLayerInitializer(nn.Module):
+    """Initialize Gaussians with multilayer representation.
+
+    The returned tensors have the shape
+
+        batch_size x dim x num_layers x height x width
+
+    where dim indicates the dimensionality of the property.
+    Some of the dimensions might be set to 1 for efficiency reasons.
+    """
+
+    def __init__(
+        self,
+        num_layers: int,
+        stride: int,
+        base_depth: float,
+        scale_factor: float,
+        disparity_factor: float,
+        color_option: ColorInitOption = "first_layer",
+        first_layer_depth_option: DepthInitOption = "surface_min",
+        rest_layer_depth_option: DepthInitOption = "surface_min",
+        normalize_depth: bool = True,
+        feature_input_stop_grad: bool = True,
+    ) -> None:
+        """Initialize MultilayerInitializer.
+
+        Args:
+            stride: The downsample rate of output feature map.
+            base_depth: The depth of the first layer (after the foreground
+                layer if use_depth=True).
+            scale_factor: Multiply scale of Gaussians by this factor.
+            disparity_factor: Factor to convert inverse depth to disparity.
+            num_layers: How many layers of Gaussians to predict.
+            color_option: Which color option to initialize the multi-layer gaussians.
+            first_layer_depth_option: Which depth option to initialize the first layer of gaussians.
+            rest_layer_depth_option: Which depth option to initialize the rest layers of gaussians.
+            normalize_depth: # Whether to normalize depth to [DepthTransformParam.depth_min,
+                DepthTransformParam.depth_max).
+            feature_input_stop_grad: Whether to not propagate gradients through feature inputs.
+        """
+        super().__init__()
+        self.num_layers = num_layers
+        self.stride = stride
+        self.base_depth = base_depth
+        self.scale_factor = scale_factor
+        self.disparity_factor = disparity_factor
+        self.color_option = color_option
+        self.first_layer_depth_option = first_layer_depth_option
+        self.rest_layer_depth_option = rest_layer_depth_option
+        self.normalize_depth = normalize_depth
+        self.feature_input_stop_grad = feature_input_stop_grad
+
+    def prepare_feature_input(self, image: torch.Tensor, depth: torch.Tensor) -> torch.Tensor:
+        """Prepare the feature input to the Guassian predictor."""
+        if self.feature_input_stop_grad:
+            image = image.detach()
+            depth = depth.detach()
+
+        normalized_disparity = self.disparity_factor / depth
+        features_in = torch.cat([image, normalized_disparity], dim=1)
+        features_in = 2.0 * features_in - 1.0
+        return features_in
+
+    def forward(self, image: torch.Tensor, depth: torch.Tensor) -> InitializerOutput:
+        """Construct Gaussian base values and prepare feature input.
+
+        Args:
+            image: The image to process.
+            depth: The corresponding depth map from the monodepth network.
+
+        Returns:
+            The base value for Gaussians.
+        """
+        image = image.contiguous()
+        depth = depth.contiguous()
+        device = depth.device
+        batch_size, _, image_height, image_width = depth.shape
+        base_height, base_width = (
+            image_height // self.stride,
+            image_width // self.stride,
+        )
+        # global_scale is the inverse of the depth_factor, which is used to rescale
+        # the depth such that it is numerically stable for training.
+        global_scale: torch.Tensor | None = None
+        if self.normalize_depth:
+            depth, depth_factor = _rescale_depth(depth)
+            global_scale = 1.0 / depth_factor
+
+        def _create_disparity_layers(num_layers: int = 1) -> torch.Tensor:
+            """Create multiple disparity layers."""
+            disparity = torch.linspace(1.0 / self.base_depth, 0.0, num_layers + 1, device=device)
+            return disparity[None, None, :-1, None, None].repeat(
+                batch_size, 1, 1, base_height, base_width
+            )
+
+        def _create_surface_layer(
+            depth: torch.Tensor,
+            depth_pooling_mode: str,
+        ) -> torch.Tensor:
+            """Create multiple surface layers."""
+            disparity = 1.0 / depth
+            if depth_pooling_mode == "min":
+                disparity = torch.max_pool2d(disparity, self.stride, self.stride)
+            elif depth_pooling_mode == "max":
+                disparity = -torch.max_pool2d(-disparity, self.stride, self.stride)
+            else:
+                raise ValueError(f"Invalid depth pooling mode {depth_pooling_mode}.")
+
+            return disparity[:, :, None, :, :]
+
+        # Input disparity dimensions:
+        #   (batch_size, num_channels in (1, 2), height, width)
+
+        # Output disparity dimensions:
+        #   (batch_size, num_channels=1, num_layers in (1, 2), height, width)
+        if self.first_layer_depth_option == "surface_min":
+            first_disparity = _create_surface_layer(depth[:, 0:1], "min")
+        elif self.first_layer_depth_option == "surface_max":
+            first_disparity = _create_surface_layer(depth[:, 0:1], "max")
+        elif self.first_layer_depth_option in ("base_depth", "linear_disparity"):
+            first_disparity = _create_disparity_layers()
+        else:
+            raise ValueError(f"Unknown depth init option: {self.first_layer_depth_option}.")
+
+        if self.num_layers == 1:
+            disparity = first_disparity
+        else:  # Fill in the rest layers.
+            following_depth = depth if depth.shape[1] == 1 else depth[:, 1:]
+            if self.rest_layer_depth_option == "surface_min":
+                following_disparity = _create_surface_layer(following_depth, "min")
+            elif self.rest_layer_depth_option == "surface_max":
+                following_disparity = _create_surface_layer(following_depth, "max")
+            elif self.rest_layer_depth_option == "base_depth":
+                following_disparity = torch.cat(
+                    [_create_disparity_layers() for i in range(self.num_layers - 1)],
+                    dim=2,
+                )
+            elif self.rest_layer_depth_option == "linear_disparity":
+                following_disparity = _create_disparity_layers(self.num_layers - 1)
+            else:
+                raise ValueError(f"Unknown depth init option: {self.rest_layer_depth_option}.")
+
+            disparity = torch.cat([first_disparity, following_disparity], dim=2)
+
+        # Prepare base values.
+        base_x_ndc, base_y_ndc = _create_base_xy(depth, self.stride, self.num_layers)
+        disparity_scale_factor = 2 * self.scale_factor * self.stride / float(image_width)
+        base_scales = _create_base_scale(disparity, disparity_scale_factor)
+
+        base_quaternions = torch.tensor([1.0, 0.0, 0.0, 0.0], device=device)
+        base_quaternions = base_quaternions[None, :, None, None, None]
+
+        # Initializing the opacitiy this way ensures that the initial transmittance
+        # is approximately
+        #
+        #       1 / e ~= (1 - 1 / self.num_layers)**self.num_layers
+        #
+        # and hence independent of the number of layers.
+        #
+        base_opacities = torch.tensor([min(1.0 / self.num_layers, 0.5)], device=device)
+        base_colors = torch.empty(
+            batch_size, 3, self.num_layers, base_height, base_width, device=device
+        ).fill_(0.5)
+        # Dimensions: (batch_size, num_channels, num_layers, height, width)
+        if self.color_option == "none":
+            pass
+        elif self.color_option == "first_layer":
+            base_colors[:, :, 0] = torch.nn.functional.avg_pool2d(image, self.stride, self.stride)
+        elif self.color_option == "all_layers":
+            temp = torch.nn.functional.avg_pool2d(image, self.stride, self.stride)
+            base_colors = temp[:, :, None, :, :].repeat(1, 1, self.num_layers, 1, 1)
+        else:
+            raise ValueError(f"Unknown color init option: {self.color_option}.")
+
+        features_in = self.prepare_feature_input(image, depth)
+        base_gaussians = GaussianBaseValues(
+            mean_x_ndc=base_x_ndc,
+            mean_y_ndc=base_y_ndc,
+            mean_inverse_z_ndc=disparity,
+            scales=base_scales,
+            quaternions=base_quaternions,
+            colors=base_colors,
+            opacities=base_opacities,
+        )
+
+        return InitializerOutput(
+            gaussian_base_values=base_gaussians,
+            feature_input=features_in,
+            global_scale=global_scale,
+        )
+
+
+def _create_base_xy(
+    depth: torch.Tensor, stride: int, num_layers: int
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Create base x and y coordinates for the gaussians in NDC space."""
+    device = depth.device
+    batch_size, _, image_height, image_width = depth.shape
+    xx = torch.arange(0.5 * stride, image_width, stride, device=device)
+    yy = torch.arange(0.5 * stride, image_height, stride, device=device)
+    xx = 2 * xx / image_width - 1.0
+    yy = 2 * yy / image_height - 1.0
+
+    xx, yy = torch.meshgrid(xx, yy, indexing="xy")
+    base_x_ndc = xx[None, None, None].repeat(batch_size, 1, num_layers, 1, 1)
+    base_y_ndc = yy[None, None, None].repeat(batch_size, 1, num_layers, 1, 1)
+
+    return base_x_ndc, base_y_ndc
+
+
+def _create_base_scale(disparity: torch.Tensor, disparity_scale_factor: float) -> torch.Tensor:
+    """Create base scale for the gaussians."""
+    inverse_disparity = torch.ones_like(disparity) / disparity
+    base_scales = inverse_disparity * disparity_scale_factor
+    return base_scales
+
+
+def _rescale_depth(
+    depth: torch.Tensor, depth_min: float = 1.0, depth_max: float = 1e2
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Rescale a depth image tensor.
+
+    Args:
+        depth: The depth tensor to transform.
+        depth_min: The min depth to scale depth to.
+        depth_max: The max clamp depth after scaling.
+
+    Returns:
+        The rescaled depth and rescale factor.
+    """
+    current_depth_min = depth.flatten(depth.ndim - 3).min(dim=-1).values
+    depth_factor = depth_min / (current_depth_min + 1e-6)
+    depth = (depth * depth_factor[..., None, None, None]).clamp(max=depth_max)
+    return depth, depth_factor

sharp/models/monodepth.py

@@ -0,0 +1,268 @@
+"""Contains Dense Transformer Prediction architecture.
+
+Implements a variant of Vision Transformers for Dense Prediction, https://arxiv.org/abs/2103.13413
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+from __future__ import annotations
+
+import copy
+from typing import NamedTuple, Tuple
+
+import torch
+import torch.nn as nn
+
+from sharp.models import normalizers
+from sharp.models.decoders import MultiresConvDecoder, create_monodepth_decoder
+from sharp.models.encoders import (
+    SlidingPyramidNetwork,
+    create_monodepth_encoder,
+)
+from sharp.utils import module_surgery
+
+from .params import MonodepthAdaptorParams, MonodepthParams
+
+DimsDecoder = Tuple[int, int, int, int, int]
+
+
+class MonodepthDensePredictionTransformer(nn.Module):
+    """Dense Prediction Transformer for monodepth.
+
+    Attach the disparity prediction head for monodepth prediction.
+    """
+
+    def __init__(
+        self,
+        encoder: SlidingPyramidNetwork,
+        decoder: MultiresConvDecoder,
+        last_dims: tuple[int, int],
+    ):
+        """Initialize Dense Prediction Transformer.
+
+        Args:
+            encoder: The SlidingPyramidTransformer backbone.
+            decoder: The MultiresConvDecoder decoder.
+            last_dims: The dimension for the last convolution layers.
+        """
+        super().__init__()
+
+        self.normalizer = normalizers.AffineRangeNormalizer(
+            input_range=(0, 1), output_range=(-1, 1)
+        )
+        self.encoder = encoder
+        self.decoder = decoder
+
+        dim_decoder = decoder.dim_out
+        self.head = nn.Sequential(
+            nn.Conv2d(dim_decoder, dim_decoder // 2, kernel_size=3, stride=1, padding=1),
+            nn.ConvTranspose2d(
+                in_channels=dim_decoder // 2,
+                out_channels=dim_decoder // 2,
+                kernel_size=2,
+                stride=2,
+                padding=0,
+                bias=True,
+            ),
+            nn.Conv2d(
+                dim_decoder // 2,
+                last_dims[0],
+                kernel_size=3,
+                stride=1,
+                padding=1,
+            ),
+            nn.ReLU(True),
+            nn.Conv2d(last_dims[0], last_dims[1], kernel_size=1, stride=1, padding=0),
+            nn.ReLU(),
+        )
+
+        # Set the final convoultion layer's bias to be 0.
+        self.head[4].bias.data.fill_(0)
+
+        self.grad_checkpointing = False
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, is_enabled=True):
+        """Enable grad checkpointing."""
+        self.grad_checkpointing = is_enabled
+        self.encoder.set_grad_checkpointing(self.grad_checkpointing)
+        self.decoder.set_grad_checkpointing(self.grad_checkpointing)
+
+    def forward(self, image: torch.Tensor) -> torch.Tensor:
+        """Decode by projection and fusion of multi-resolution encodings."""
+        encodings = self.encoder(self.normalizer(image))
+        num_encoder_features = len(self.encoder.dims_encoder)
+        features = self.decoder(encodings[:num_encoder_features])
+        disparity = self.head(features)
+        return disparity
+
+    def internal_resolution(self) -> int:
+        """Return the internal image size of the network."""
+        return self.encoder.internal_resolution()
+
+
+def create_monodepth_dpt(
+    params: MonodepthParams | None = None,
+) -> MonodepthDensePredictionTransformer:
+    """Creates DepthDensePredictionTransformer model.
+
+    Args:
+        params: Parameters of monodepth network.
+
+    Returns:
+        The configured monodepth DPT.
+    """
+    if params is None:
+        params = MonodepthParams()
+    encoder: SlidingPyramidNetwork = create_monodepth_encoder(
+        params.patch_encoder_preset,
+        params.image_encoder_preset,
+        use_patch_overlap=params.use_patch_overlap,
+        last_encoder=params.dims_decoder[0],
+    )
+
+    decoder: MultiresConvDecoder = create_monodepth_decoder(
+        params.patch_encoder_preset, params.dims_decoder
+    )
+
+    monodepth_model = MonodepthDensePredictionTransformer(
+        encoder=encoder, decoder=decoder, last_dims=(32, 1)
+    )
+
+    # By default, we don't train the monodepth model.
+    # However, we allow to selectively unfreeze parts of the network.
+    monodepth_model.requires_grad_(False)
+
+    monodepth_model.encoder.set_requires_grad_(
+        patch_encoder=params.unfreeze_patch_encoder,
+        image_encoder=params.unfreeze_image_encoder,
+    )
+    monodepth_model.decoder.requires_grad_(params.unfreeze_decoder)
+    monodepth_model.head.requires_grad_(params.unfreeze_head)
+
+    if not params.unfreeze_norm_layers:
+        module_surgery.freeze_norm_layer(monodepth_model)
+
+    monodepth_model.set_grad_checkpointing(params.grad_checkpointing)
+
+    return monodepth_model
+
+
+class MonodepthOutput(NamedTuple):
+    """Output of the monodepth model."""
+
+    # Disparity output from the monodepth model.
+    disparity: torch.Tensor
+    # Multi-level features from monodepth encoder.
+    encoder_features: list[torch.Tensor]
+    # Single-level feature from monodepth decoder.
+    decoder_features: torch.Tensor
+    # List of monodepth features to be used in gaussian predictor.
+    output_features: list[torch.Tensor]
+    # List of intermediate encoder features to be used in distillation.
+    intermediate_features: list[torch.Tensor] = []
+
+
+class MonodepthWithEncodingAdaptor(nn.Module):
+    """Monodepth model with feature maps."""
+
+    def __init__(
+        self,
+        monodepth_predictor: MonodepthDensePredictionTransformer,
+        return_encoder_features: bool,
+        return_decoder_features: bool,
+        num_monodepth_layers: int,
+        sorting_monodepth: bool,
+    ):
+        """Initialize MonodepthWithEncodingAdaptor.
+
+        Args:
+            monodepth_predictor: The monodepth model.
+            return_encoder_features: Whether to return encoder features from monodepth model.
+            return_decoder_features: Whether to return decoder features from monodepth model.
+            num_monodepth_layers: How many layers the monodepth model predicts.
+            sorting_monodepth: Whether to sort the monodepth output (for two layer monodepth).
+        """
+        super().__init__()
+        self.monodepth_predictor = monodepth_predictor
+        self.return_encoder_features = return_encoder_features
+        self.return_decoder_features = return_decoder_features
+        self.num_monodepth_layers = num_monodepth_layers
+        self.sorting_monodepth = sorting_monodepth
+
+    def forward(self, image: torch.Tensor) -> MonodepthOutput:
+        """Process image and return disparity and feature maps."""
+        inputs = self.monodepth_predictor.normalizer(image)
+        encoder_output = self.monodepth_predictor.encoder(inputs)
+
+        num_encoder_features = len(self.monodepth_predictor.encoder.dims_encoder)
+
+        # NOTE: whether intermediate features are empty have already been decided
+        # in monodepth_predictor during create_monodepth_dpt.
+        encoder_features = encoder_output[:num_encoder_features]
+        intermediate_features = encoder_output[num_encoder_features:]
+        decoder_features = self.monodepth_predictor.decoder(encoder_features)
+        disparity = self.monodepth_predictor.head(decoder_features)
+
+        # We cannot use disparity.shape[1], otherwise the tracer will fail.
+        if self.num_monodepth_layers == 2 and self.sorting_monodepth:
+            first_layer_disparity = disparity.max(dim=1, keepdims=True).values
+            second_layer_disparity = disparity.min(dim=1, keepdims=True).values
+            disparity = torch.cat([first_layer_disparity, second_layer_disparity], dim=1)
+
+        output_features = []
+        if self.return_encoder_features:
+            output_features.extend(encoder_features)
+
+        if self.return_decoder_features:
+            output_features.append(decoder_features)
+
+        return MonodepthOutput(
+            disparity=disparity,
+            encoder_features=encoder_features,
+            decoder_features=decoder_features,
+            output_features=output_features,
+            intermediate_features=intermediate_features,
+        )
+
+    def get_feature_dims(self) -> list[int]:
+        """Return dimensions of output feature maps."""
+        dims = []
+        if self.return_encoder_features:
+            dims.extend(self.monodepth_predictor.encoder.dims_encoder)
+
+        if self.return_decoder_features:
+            dims.append(self.monodepth_predictor.decoder.dim_out)
+
+        return dims
+
+    def internal_resolution(self) -> int:
+        """Return the internal image size of the network."""
+        return self.monodepth_predictor.internal_resolution()
+
+    def replicate_head(self, num_repeat: int):
+        """Replicate the last convolution layer (head[4] in DPT) for multi layer depth."""
+        conv_last = copy.deepcopy(self.monodepth_predictor.head[4])
+        self.monodepth_predictor.head[4].out_channels = num_repeat
+        self.monodepth_predictor.head[4].weight = nn.Parameter(
+            conv_last.weight.repeat(num_repeat, 1, 1, 1)
+        )
+        self.monodepth_predictor.head[4].bias = nn.Parameter(conv_last.bias.repeat(num_repeat))
+
+
+def create_monodepth_adaptor(
+    monodepth_predictor: MonodepthDensePredictionTransformer,
+    params: MonodepthAdaptorParams,
+    num_monodepth_layers: int,
+    sorting_monodepth: bool,
+) -> MonodepthWithEncodingAdaptor:
+    """Create an adaptor that returns both disparity and features."""
+    adaptor = MonodepthWithEncodingAdaptor(
+        monodepth_predictor=monodepth_predictor,
+        return_encoder_features=params.encoder_features,
+        return_decoder_features=params.decoder_features,
+        num_monodepth_layers=num_monodepth_layers,
+        sorting_monodepth=sorting_monodepth,
+    )
+    return adaptor

sharp/models/normalizers.py

@@ -0,0 +1,80 @@
+"""Contains an implementation of image normalizers for perceptual loss.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+from __future__ import annotations
+
+from typing import Sequence, Union
+
+import torch
+from torch import nn
+
+
+class MeanStdNormalizer(nn.Module):
+    """Normalizing image input by mean and std."""
+
+    mean: torch.Tensor
+    std_inv: torch.Tensor
+
+    def __init__(
+        self,
+        mean: Union[Sequence[float], torch.Tensor],
+        std: Union[Sequence[float], torch.Tensor],
+    ):
+        """Initialize MeanStdNormalizer."""
+        super(MeanStdNormalizer, self).__init__()
+        if not isinstance(mean, torch.Tensor):
+            mean = torch.as_tensor(mean).view(-1, 1, 1)
+        if not isinstance(std, torch.Tensor):
+            std = torch.as_tensor(std).view(-1, 1, 1)
+        self.register_buffer("mean", mean)
+        # We use inverse std to use a multiplication which is better supported by the hardware
+        self.register_buffer("std_inv", 1.0 / std)
+
+    def forward(self, image: torch.Tensor) -> torch.Tensor:
+        """Apply mean and std normalization over input image."""
+        return (image - self.mean) * self.std_inv
+
+
+class AffineRangeNormalizer(nn.Module):
+    """Perform linear mapping to map input_range to output_range.
+
+    Output_range defaults to (0, 1).
+    """
+
+    def __init__(
+        self,
+        input_range: tuple[float, float],
+        output_range: tuple[float, float] = (0, 1),
+    ):
+        """Initialize AffineRangeNormalizer."""
+        super().__init__()
+        input_min, input_max = input_range
+        output_min, output_max = output_range
+        if input_max <= input_min:
+            raise ValueError(f"Invalid input_range: {input_range}")
+        if output_max <= output_min:
+            raise ValueError(f"Invalid output_range: {output_range}")
+
+        self.scale = (output_max - output_min) / (input_max - input_min)
+        self.bias = output_min - input_min * self.scale
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Apply affine range normalization over input image."""
+        if self.scale != 1.0:
+            x = x * self.scale
+
+        if self.bias != 0.0:
+            x = x + self.bias
+
+        return x
+
+
+class MobileNetNormalizer(AffineRangeNormalizer):
+    """Image normalization in mobilenet."""
+
+    def __init__(self, input_range: tuple[float, float] = (0, 1)):
+        """Initialize MobileNetNormalizer."""
+        super().__init__(input_range=input_range, output_range=(-1, 1))

sharp/models/params.py

@@ -0,0 +1,203 @@
+"""Contains params for backbone.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+import dataclasses
+from typing import Literal
+
+import sharp.utils.math as math_utils
+from sharp.models.blocks import NormLayerName, UpsamplingMode
+from sharp.models.presets import ViTPreset
+from sharp.utils.color_space import ColorSpace
+
+DimsDecoder = tuple[int, int, int, int, int]
+DPTImageEncoderType = Literal["skip_conv", "skip_conv_kernel2"]
+
+ColorInitOption = Literal[
+    "none",  # Initialize as gray.
+    "first_layer",  # Initialize the first layer with input image, other layers with gray.
+    "all_layers",  # Initialize all layers with input image.
+]
+DepthInitOption = Literal[
+    # Initialize the layer of gaussian on surface using min pooling of input depth.
+    "surface_min",
+    # Initialize the layer of gaussian on surface using max pooling of input depth
+    "surface_max",
+    # Initialize the layer of gaussian on plane using base_depth depth.
+    "base_depth",
+    # Initialize the layer of gaussian on plane based on base_depth and index of layer.
+    "linear_disparity",
+]
+
+
+@dataclasses.dataclass
+class AlignmentParams:
+    """Parameters for depth alignment."""
+
+    kernel_size: int = 16
+    stride: int = 1
+    frozen: bool = False
+
+    # The following parameters are only used for LearnedAlignment.
+    # Number of steps in the UNet for LearnedAlignment.
+    steps: int = 4
+    # Activation type for LearnedAlignment.
+    activation_type: math_utils.ActivationType = "exp"
+    # Whether to use depth decoder features for LearnedAlignment.
+    depth_decoder_features: bool = False
+    # Base width of the UNet for LearnedAlignment.
+    base_width: int = 16
+
+
+@dataclasses.dataclass
+class DeltaFactor:
+    """Factors to multiply deltas with before activation.
+
+    These factors effectively selectively reduce the learning rate.
+    """
+
+    xy: float = 0.001
+    z: float = 0.001
+    color: float = 0.1  # We recommend 0.1 for linearRGB and 1.0 for sRGB.
+    opacity: float = 1.0
+    scale: float = 1.0
+    quaternion: float = 1.0
+
+
+@dataclasses.dataclass
+class InitializerParams:
+    """Parameters for initializer."""
+
+    # Common parameters.
+    # Multiply scale of Gaussians by this factor.
+    scale_factor: float = 1.0
+    # Factor to convert inverse depth to disparity.
+    disparity_factor: float = 1.0
+    # Stride of the initializer.
+    stride: int = 2
+
+    # Parameters that only affect MultiLayerInitializer.
+    # How many layers of Gaussians to predict (only available for MultiLayerInitializer).
+    num_layers: int = 2
+    # Which option to use for depth initialization.
+    first_layer_depth_option: DepthInitOption = "surface_min"
+    rest_layer_depth_option: DepthInitOption = "surface_min"
+    # Which option to use for color initialization.
+    color_option: ColorInitOption = "all_layers"
+    # Which depth value to use for depth layers.
+    base_depth: float = 10.0
+    # Deactivate gradient for feature inputs.
+    feature_input_stop_grad: bool = False
+    # Whether to normalize depth to [DepthTransformParam.depth_min,
+    # DepthTransformParam.depth_max).
+    normalize_depth: bool = True
+
+    # Output only the inpainted layer. In this case, num_layers = 1.
+    output_inpainted_layer_only: bool = False
+    # Whether to set the uninpainted region to zero opacities.
+    set_uninpainted_opacity_to_zero: bool = False
+    # Whether to concatenate the inpainting mask to the feature input.
+    concat_inpainting_mask: bool = False
+
+
+@dataclasses.dataclass
+class MonodepthParams:
+    """Parameters for monodepth network."""
+
+    patch_encoder_preset: ViTPreset = "dinov2l16_384"
+    image_encoder_preset: ViTPreset = "dinov2l16_384"
+
+    checkpoint_uri: str | None = None
+    unfreeze_patch_encoder: bool = False
+    unfreeze_image_encoder: bool = False
+    unfreeze_decoder: bool = False
+    unfreeze_head: bool = False
+    unfreeze_norm_layers: bool = False
+    grad_checkpointing: bool = False
+    use_patch_overlap: bool = True
+    dims_decoder: DimsDecoder = (256, 256, 256, 256, 256)
+
+
+@dataclasses.dataclass
+class MonodepthAdaptorParams:
+    """Parameters for monodepth network feature adaptor."""
+
+    encoder_features: bool = True
+    decoder_features: bool = False
+
+
+@dataclasses.dataclass
+class GaussianDecoderParams:
+    """Parameters for backbone with default values."""
+
+    dim_in: int = 5
+    dim_out: int = 32
+    # Which normalization to use in backbone.
+    norm_type: NormLayerName = "group_norm"
+    # How many groups to use for group normalization.
+    norm_num_groups: int = 8
+    # Stride of backbone.
+    stride: int = 2
+
+    patch_encoder_preset: ViTPreset = "dinov2l16_384"
+    image_encoder_preset: ViTPreset = "dinov2l16_384"
+
+    # Dimensionality of feature maps for DPT decoder.
+    dims_decoder: DimsDecoder = (128, 128, 128, 128, 128)
+
+    # Whether to use depth as input.
+    use_depth_input: bool = True
+
+    # Whether to enable gradient checkpointing for the backbone
+    grad_checkpointing: bool = False
+
+    # What mode to use for upsampling in decoder.
+    upsampling_mode: UpsamplingMode = "transposed_conv"
+
+    # The type of image encoder.
+    image_encoder_type: DPTImageEncoderType = "skip_conv_kernel2"
+
+
+@dataclasses.dataclass
+class PredictorParams:
+    """Parameters for predictors with default values."""
+
+    # Parameters for submodules.
+    initializer: InitializerParams = dataclasses.field(default_factory=InitializerParams)
+    monodepth: MonodepthParams = dataclasses.field(default_factory=MonodepthParams)
+    monodepth_adaptor: MonodepthAdaptorParams = dataclasses.field(
+        default_factory=MonodepthAdaptorParams
+    )
+    gaussian_decoder: GaussianDecoderParams = dataclasses.field(
+        default_factory=GaussianDecoderParams
+    )
+    # How to align depth map (only relevant for RGBGaussianPredictor).
+    depth_alignment: AlignmentParams = dataclasses.field(default_factory=AlignmentParams)
+
+    # Selectively reduce learning rate for different properties.
+    delta_factor: DeltaFactor = dataclasses.field(default_factory=DeltaFactor)
+    # The maximum scale of Gaussians relative to initial scale.
+    max_scale: float = 10.0
+    # The minimum scale of Gaussians relative to initial scale.
+    min_scale: float = 0.0
+    # Which normalization to use in prediction head.
+    norm_type: NormLayerName = "group_norm"
+    # How many groups to use for group normalization.
+    norm_num_groups: int = 8
+    # Whether to use predicted mean to sample triplane features.
+    use_predicted_mean: bool = False
+    # Which activation function to use for colors / opacities.
+    color_activation_type: math_utils.ActivationType = "sigmoid"
+    opacity_activation_type: math_utils.ActivationType = "sigmoid"
+    # Colorspace of the renderer ("linearRGB" or "sRGB").
+    color_space: ColorSpace = "linearRGB"
+    # A small value to avoid ill-conditioned splats
+    low_pass_filter_eps: float = 1e-2
+    # How many layer of depth does monodepth model predict.
+    num_monodepth_layers: int = 2
+    # Whether to sort the monodepth output (for two layer monodepth).
+    sorting_monodepth: bool = False
+    # Whether to account the z offsets for estimating base scale.
+    base_scale_on_predicted_mean: bool = True

sharp/models/predictor.py

@@ -0,0 +1,201 @@
+"""Contains definition of RGB-only gaussian predictor.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+from __future__ import annotations
+
+import logging
+
+import torch
+from torch import nn
+
+from sharp.models.monodepth import MonodepthWithEncodingAdaptor
+from sharp.utils.gaussians import Gaussians3D
+
+from .composer import GaussianComposer
+
+LOGGER = logging.getLogger(__name__)
+
+
+class DepthAlignment(nn.Module):
+    """Depth alignment in a dedicated nn.Module.
+
+    Wrap scale_map_estimator to perform the conditional logic in a separated torch
+    module outside the forward of RGBGaussianPredictor. This module can be then
+    excluded during symbolic tracing.
+    """
+
+    def __init__(self, scale_map_estimator: nn.Module | None):
+        """Initialize DepthAlignmentWrapper.
+
+        Args:
+            scale_map_estimator: Module to align monodepth to ground truth depth.
+        """
+        super().__init__()
+        self.scale_map_estimator = scale_map_estimator
+
+    def forward(
+        self,
+        monodepth: torch.Tensor,
+        depth: torch.Tensor,
+        depth_decoder_features: torch.Tensor | None = None,
+    ):
+        """Optionally align monodepth to ground truth with a local scale map.
+
+        Args:
+            monodepth: The monodepth model with intermediate features to use.
+            depth: Ground truth depth to align predicted depth to.
+            depth_decoder_features: The (optional) monodepth decoder features.
+        """
+        if depth is not None and self.scale_map_estimator is not None:
+            depth_alignment_map = self.scale_map_estimator(
+                monodepth[:, 0:1], depth, depth_decoder_features
+            )
+            monodepth = depth_alignment_map * monodepth
+        else:
+            # Some losses rely on the presence of an alignment map.
+            # We ensure that they can be computed by creating a fake alignment map.
+            depth_alignment_map = torch.ones_like(monodepth)
+        return monodepth, depth_alignment_map
+
+
+class RGBGaussianPredictor(nn.Module):
+    """Predicts 3D Gaussians from images."""
+
+    feature_model: nn.Module
+
+    def __init__(
+        self,
+        init_model: nn.Module,
+        monodepth_model: MonodepthWithEncodingAdaptor,
+        feature_model: nn.Module,
+        prediction_head: nn.Module,
+        gaussian_composer: GaussianComposer,
+        scale_map_estimator: nn.Module | None,
+    ) -> None:
+        """Initialize RGBGaussianPredictor.
+
+        Args:
+            init_model: A model mapping image and depth to base values.
+            monodepth_model: The monodepth model with intermediate features to use.
+            feature_model: The image2image model to predict Gaussians from.
+            prediction_head: Head to decode image features.
+            gaussian_composer: Module to compose final prediction from deltas and
+                base values.
+            scale_map_estimator: Module to align monodepth to ground truth depth.
+
+        Note:
+        ----
+            when monodepth_model is trainable, using local depth alignment can
+            result in the monodepth model losing its ability to predict shapes. It is
+            hence recommend to deactivate the corresponding flag.
+        """
+        super().__init__()
+        self.init_model = init_model
+        self.feature_model = feature_model
+        self.monodepth_model = monodepth_model
+        self.prediction_head = prediction_head
+        self.gaussian_composer = gaussian_composer
+        self.depth_alignment = DepthAlignment(scale_map_estimator)
+
+    def forward(
+        self,
+        image: torch.Tensor,
+        disparity_factor: torch.Tensor,
+        depth: torch.Tensor | None = None,
+    ) -> Gaussians3D:
+        """Predict 3D Gaussians.
+
+        Args:
+            image: The image to process.
+            disparity_factor: Factor to convert depth to disparities.
+            depth: Ground truth depth to align predicted depth to.
+
+        Returns:
+            The predicted 3D Gaussians.
+
+        Note:
+        ----
+        During training, it is recommended to feed an additional ground truth depth
+        map to the network to align the predicted depth to. During inference, it is
+        recommended to use depth_gt=None and use monodepth_disparity output from the
+        model instead to compute depth.
+        """
+        # Estimate depth and align to ground truth (if available).
+        monodepth_output = self.monodepth_model(image)
+        monodepth_disparity = monodepth_output.disparity
+
+        disparity_factor = disparity_factor[:, None, None, None]
+        monodepth = disparity_factor / monodepth_disparity.clamp(min=1e-4, max=1e4)
+
+        # In the model we apply additional alignment to provided ground truth depth
+        # as well as additional normalization.
+        #
+        # The overall graph looks as follows:
+        #
+        #     monodepth        depth    # Both monodepth and depth are metric here.
+        #         |              |
+        #         +------+-------+
+        #                |
+        #        +-------+--------+     # Optionally align monodepth to ground truth
+        #        |depth_alignement|     # with a local scale map.
+        #        +-------+--------+
+        #                |
+        #                v
+        #       monodepth (aligned)     # Monodepth is now aligned to ground truth.
+        #                |
+        #          +-----+----+         # Normalize depth and compute base gaussians.
+        #          |init_model|         # in these normalized coordinates.
+        #          +-----+----+
+        #                |
+        #                v
+        #   +------ init_output         # Init_output consists of features, base
+        #   |            |              # gaussians and a global scale.
+        #   |     +------+-----+
+        #   |     |main network|        # Compute delta values to base gaussians.
+        #   |     +------+-----+
+        #   |            |
+        #   |            V
+        #   |        delta_values       # The delta values are computed with normalized depth.
+        #   |            |
+        #   |    +-------+---------+
+        #   +--> |gaussian_composer|    # Add delta to base values and unscale gaussians.
+        #        +-------+---------+
+        #                |
+        #                v
+        #            gaussians          # The final Gaussians are metric again.
+        #
+
+        # The logic to decide whether to align monodepth to the ground truth is wrapped
+        # in a submodule 'DepthAlignement' to facilitate the symbolic tracing of the
+        # predictor. This way, the depth alignment submodule containing the conditional
+        # logic can be excluded during the tracing and the graph of the predictors is
+        # static.
+        monodepth, _ = self.depth_alignment(
+            monodepth,
+            depth,
+            monodepth_output.decoder_features,
+        )
+
+        init_output = self.init_model(image, monodepth)
+        image_features = self.feature_model(
+            init_output.feature_input, encodings=monodepth_output.output_features
+        )
+        delta_values = self.prediction_head(image_features)
+        gaussians = self.gaussian_composer(
+            delta=delta_values,
+            base_values=init_output.gaussian_base_values,
+            global_scale=init_output.global_scale,
+        )
+        return gaussians
+
+    def internal_resolution(self) -> int:
+        """Internal resolution."""
+        return self.monodepth_model.internal_resolution()
+
+    @property
+    def output_resolution(self) -> int:
+        """Output resolution of Gaussians."""
+        return self.internal_resolution() // 2

sharp/models/presets/__init__.py

@@ -0,0 +1,23 @@
+"""Contains presets for pretrained neural networks.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+from .monodepth import (
+    MONODEPTH_ENCODER_DIMS_MAP,
+    MONODEPTH_HOOK_IDS_MAP,
+)
+from .vit import (
+    VIT_CONFIG_DICT,
+    ViTConfig,
+    ViTPreset,
+)
+
+__all__ = [
+    "ViTConfig",
+    "ViTPreset",
+    "VIT_CONFIG_DICT",
+    "MONODEPTH_ENCODER_DIMS_MAP",
+    "MONODEPTH_HOOK_IDS_MAP",
+]

sharp/models/presets/monodepth.py

@@ -0,0 +1,21 @@
+"""Contains preset for monodepth modules.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+from __future__ import annotations
+
+from .vit import ViTPreset
+
+# Map the decoder configuration with the number of output channels
+# for each tensor from the decoder output.
+MONODEPTH_ENCODER_DIMS_MAP: dict[ViTPreset, list[int]] = {
+    # For publication
+    "dinov2l16_384": [256, 512, 1024, 1024],
+}
+
+MONODEPTH_HOOK_IDS_MAP: dict[ViTPreset, list[int]] = {
+    # For publication
+    "dinov2l16_384": [5, 11, 17, 23],
+}

sharp/models/presets/vit.py

@@ -0,0 +1,58 @@
+"""Contains preset for ViT modules.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+from typing import Literal
+
+ViTPreset = Literal["dinov2l16_384",]
+
+MLPMode = Literal["vanilla", "glu"]
+
+
+@dataclasses.dataclass
+class ViTConfig:
+    """Configuration for ViT."""
+
+    in_chans: int
+    embed_dim: int
+    depth: int
+    num_heads: int
+    init_values: float
+
+    img_size: int = 384
+    patch_size: int = 16
+
+    num_classes: int = 21841
+    mlp_ratio: float = 4.0
+    drop_rate: float = 0.0
+    attn_drop_rate: float = 0.0
+    drop_path_rate: float = 0.0
+    qkv_bias: bool = True
+    global_pool: str = "avg"
+
+    # Properties for timm_vit.
+    mlp_mode: MLPMode = "vanilla"
+
+    # Properties for SPN.
+    intermediate_features_ids: list[int] | None = None
+
+    def asdict(self):
+        """Convenience method to convert the class to a dict."""
+        return dataclasses.asdict(self)
+
+
+VIT_CONFIG_DICT: dict[ViTPreset, ViTConfig] = {
+    "dinov2l16_384": ViTConfig(
+        in_chans=3,
+        embed_dim=1024,
+        depth=24,
+        num_heads=16,
+        init_values=1e-5,
+        global_pool="",
+    ),
+}

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."