File size: 4,802 Bytes
2b156ca | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 | # ---------------------------------------------------------------------
# Copyright (c) 2026 Qualcomm Technologies, Inc. and/or its subsidiaries.
# SPDX-License-Identifier: BSD-3-Clause
# ---------------------------------------------------------------------
from __future__ import annotations
from collections.abc import Callable
from dataclasses import dataclass
from typing import Any, cast
import matplotlib.pyplot as plt
import numpy as np
import numpy.typing as npt
import torch
from PIL import Image
from torchvision import transforms
from qai_hub_models.utils.image_processing import pil_resize_pad, undo_resize_pad
@dataclass
class DepthProPrediction:
"""Structured output of :class:`DepthProApp`.
``depth`` is metric depth in scene units (aligned via ``focal_length_px``);
``heatmap`` is a plasma-colored visualization of inverse depth for display;
``field_of_view`` is horizontal FoV in degrees; ``focal_length_px`` is
the pixel focal length derived from FoV and the original image width.
"""
depth: npt.NDArray[np.float32]
heatmap: Image.Image
field_of_view: float
focal_length_px: float
class DepthProApp:
"""End-to-end app for Apple DepthPro depth estimation.
Wraps a callable returning ``(predicted_depth, field_of_view)`` — either
the torch model or an on-device runner. Preprocessing resizes with
aspect-preserving padding to the network's 1536x1536 input; post-
processing mirrors HuggingFace's
``DepthProImageProcessorFast.post_process_depth_estimation`` (metric
scaling by ``width / focal_length_px``, then inversion of the canonical
inverse depth).
"""
def __init__(
self,
model: Callable[
[torch.Tensor], tuple[torch.Tensor, torch.Tensor]
],
input_height: int | None = None,
input_width: int | None = None,
) -> None:
self.model = model
if input_height is None or input_width is None:
get_input_spec = getattr(model, "get_input_spec", None)
if get_input_spec is None:
raise TypeError(
"DepthProApp needs input_height and input_width when the "
"provided model is not a BaseModel (has no get_input_spec)."
)
_, _, h, w = get_input_spec()["image"][0]
input_height = input_height if input_height is not None else h
input_width = input_width if input_width is not None else w
self.input_height = input_height
self.input_width = input_width
def predict(self, *args: Any, **kwargs: Any) -> DepthProPrediction:
return self.estimate_depth(*args, **kwargs)
def estimate_depth(self, image: Image.Image) -> DepthProPrediction:
"""Estimate depth, FoV, and focal length for a single image.
Parameters
----------
image
PIL image in any resolution / aspect ratio.
"""
resized_image, scale, padding = pil_resize_pad(
image, (self.input_height, self.input_width)
)
image_tensor = transforms.ToTensor()(resized_image).unsqueeze(0)
predicted_depth, field_of_view = self.model(image_tensor)
# Horizontal FoV -> focal length in pixels of the *original* image.
# Matches HF's post_process_depth_estimation.
orig_width = float(image.size[0])
fov_deg = field_of_view.detach().float().view(-1)
focal_length_px = 0.5 * orig_width / torch.tan(
0.5 * torch.deg2rad(fov_deg)
)
# Metric scaling of canonical inverse depth (again, from HF).
depth_scaled = predicted_depth * (
orig_width / focal_length_px
).view(-1, 1, 1)
# (B, 1, H, W) shape is what undo_resize_pad expects.
depth_map = undo_resize_pad(
depth_scaled.unsqueeze(1), image.size, scale, padding
)
# Canonical inverse depth -> metric depth. Clamp mirrors HF.
depth_map = 1.0 / torch.clamp(depth_map, min=1e-4, max=1e4)
depth_np = cast(
npt.NDArray[np.float32],
depth_map.squeeze().detach().cpu().numpy().astype(np.float32),
)
# Visualize inverse depth so closer objects appear brighter, matching
# the shared depth-estimation demo convention.
inv = 1.0 / np.maximum(depth_np, 1e-6)
inv_norm = inv / max(inv.max(), 1e-6)
heatmap = plt.cm.get_cmap("plasma")(inv_norm)[..., :3]
heatmap_image = Image.fromarray((heatmap * 255).astype(np.uint8))
return DepthProPrediction(
depth=depth_np,
heatmap=heatmap_image,
field_of_view=float(fov_deg.squeeze().item()),
focal_length_px=float(focal_length_px.squeeze().item()),
)
|