File size: 5,055 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 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 | # ---------------------------------------------------------------------
# Copyright (c) 2026 Qualcomm Technologies, Inc. and/or its subsidiaries.
# SPDX-License-Identifier: BSD-3-Clause
# ---------------------------------------------------------------------
from __future__ import annotations
import torch
from transformers import DepthProForDepthEstimation
from typing_extensions import Self
from qai_hub_models.datasets.nyuv2 import NYUV2Dataset
from qai_hub_models.models._shared.depth_estimation.depth_evaluator import (
DepthEvaluator,
)
from qai_hub_models.utils.base_dataset import BaseDataset
from qai_hub_models.utils.base_evaluator import BaseEvaluator
from qai_hub_models.utils.base_model import BaseModel
from qai_hub_models.utils.input_spec import (
ColorFormat,
ImageMetadata,
InputSpec,
IoType,
OutputSpec,
TensorSpec,
)
MODEL_ID = "depth_pro"
MODEL_ASSET_VERSION = 1
DEFAULT_WEIGHTS = "apple/DepthPro-hf"
DEFAULT_INPUT_SIZE = 1536
class DepthProDepthEvaluator(DepthEvaluator):
"""Adapts the shared depth δ1 evaluator to DepthPro's (depth, fov) tuple.
DepthPro's canonical inverse-depth head is what the shared evaluator
already scale/shift-aligns against NYUv2 ground truth, so the FoV output
is ignored here — focal-length calibration is only needed for absolute
metric alignment, which the δ1 metric is invariant to.
"""
def add_batch(
self,
output: torch.Tensor | tuple[torch.Tensor, ...] | list[torch.Tensor],
gt: torch.Tensor,
) -> None:
if isinstance(output, (tuple, list)):
output = output[0]
if output.dim() == 3:
# depth arrives as (B, H, W); base evaluator expects (B, 1, H, W)
output = output.unsqueeze(1)
super().add_batch(output, gt)
class DepthPro(BaseModel):
"""Apple DepthPro monocular metric depth estimator, end-to-end.
Exposes two on-device outputs: canonical inverse depth at the network's
input resolution and a scalar horizontal field of view (degrees) per
image. Off-device post-processing (see ``DepthProApp``) converts these
into metric depth and a focal length in pixels, matching HuggingFace's
``DepthProImageProcessorFast.post_process_depth_estimation``.
"""
def __init__(self, model: torch.nn.Module) -> None:
super().__init__()
self.model = model.eval()
@classmethod
def from_pretrained(cls, ckpt: str = DEFAULT_WEIGHTS) -> Self:
net = DepthProForDepthEstimation.from_pretrained(ckpt)
return cls(net)
def forward(
self, image: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
"""Run DepthPro on `image`.
Parameters
----------
image
Shape ``[B, 3, 1536, 1536]`` RGB in ``[0, 1]``.
Returns
-------
predicted_depth
Shape ``[B, 1536, 1536]``. Canonical inverse depth (the raw head
output before FoV-based metric scaling and inversion).
field_of_view
Shape ``[B]``, horizontal field of view in degrees.
"""
# HF's DepthProImageProcessorFast rescales to [0, 1] then normalizes
# with mean=std=0.5, i.e. (image - 0.5) / 0.5 = 2 * image - 1.
pixel_values = image * 2.0 - 1.0
predicted_depth, field_of_view = self.model(
pixel_values, return_dict=False
)
return predicted_depth, field_of_view
def get_input_spec(
self,
batch_size: int = 1,
height: int = DEFAULT_INPUT_SIZE,
width: int = DEFAULT_INPUT_SIZE,
) -> InputSpec:
return {
"image": TensorSpec(
shape=(batch_size, 3, height, width),
dtype="float32",
io_type=IoType.IMAGE,
value_range=(0.0, 1.0),
image_metadata=ImageMetadata(color_format=ColorFormat.RGB),
apply_runtime_channel_reordering=True,
),
}
def get_output_spec(self) -> OutputSpec:
return {
"predicted_depth": TensorSpec(
io_type=IoType.TENSOR,
description=(
"Canonical inverse depth at the network's input "
"resolution. Invert and rescale by "
"width / focal_length_px for metric depth."
),
apply_runtime_channel_reordering=True,
),
"field_of_view": TensorSpec(
io_type=IoType.TENSOR,
description=(
"Horizontal field of view in degrees, one scalar per "
"image; used off-device to derive focal length."
),
),
}
def get_evaluator(self) -> BaseEvaluator:
return DepthProDepthEvaluator()
@classmethod
def get_eval_dataset_classes(cls) -> list[type[BaseDataset]]:
return [NYUV2Dataset]
def get_calibration_dataset_cls(self) -> type[BaseDataset]:
return NYUV2Dataset
|