File size: 3,460 Bytes
52ccb53 457b373 52ccb53 457b373 ba5cb88 dd184e9 52ccb53 | 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 | """Hugging Face Inference Endpoints custom handler for monocular depth.
HF's managed task list does not expose "Depth Estimation", so we deploy
Depth-Anything-V2 through a custom handler instead. HF's default inference
container auto-detects an `EndpointHandler` class in this file and calls it per
request, bypassing the fixed task list.
Request : raw image bytes (Content-Type: image/*) OR JSON {"inputs": "<b64>"}.
Response : JSON {"depth": "<base64 PNG>", "width": W, "height": H} — an 8-bit
grayscale depth map (near = bright). This matches the AI engine's
`hf_client.depth()` contract, which decodes {"depth": <base64 png>}.
"""
from __future__ import annotations
import base64
import io
from typing import Any, Dict
import numpy as np
from PIL import Image
class EndpointHandler:
def __init__(self, model_dir: str = "", **kwargs: Any) -> None:
# HF's inference toolkit calls this as EndpointHandler(model_dir=...),
# so the first positional/keyword arg MUST be named `model_dir`.
# Lazy heavy imports so container build/health checks stay light.
import torch
from transformers import pipeline
self._torch = torch
device = 0 if torch.cuda.is_available() else -1
# This repo holds only the handler (no weights), so `model_dir` has no
# model config. Always load the depth model from its public repo id.
import os
model = os.environ.get("TARGET_MODEL_ID", "depth-anything/Depth-Anything-V2-Large-hf")
self.pipe = pipeline("depth-estimation", model=model, device=device)
def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
image = self._read_image(data)
with self._torch.inference_mode():
out = self.pipe(image)
depth = out["depth"] if isinstance(out, dict) else out
arr = self._to_uint8(depth, image.size)
buf = io.BytesIO()
Image.fromarray(arr, mode="L").save(buf, format="PNG")
b64 = base64.b64encode(buf.getvalue()).decode("ascii")
h, w = arr.shape
return {"depth": b64, "width": w, "height": h}
# --- helpers -----------------------------------------------------------
@staticmethod
def _read_image(data: Dict[str, Any]) -> Image.Image:
# HF passes raw request bytes under "inputs" (bytes) for binary payloads,
# or a base64 string when sent as JSON.
inputs = data.get("inputs", data)
if isinstance(inputs, (bytes, bytearray)):
raw = bytes(inputs)
elif isinstance(inputs, str):
raw = base64.b64decode(inputs)
elif isinstance(inputs, Image.Image):
return inputs.convert("RGB")
else:
raise ValueError("Unsupported input type for depth handler")
return Image.open(io.BytesIO(raw)).convert("RGB")
@staticmethod
def _to_uint8(depth: Any, size) -> np.ndarray:
"""Normalize a depth output (PIL image or array/tensor) to uint8 HxW."""
if isinstance(depth, Image.Image):
arr = np.asarray(depth, dtype=np.float32)
else:
arr = np.asarray(depth, dtype=np.float32)
if arr.ndim == 3:
arr = arr[..., 0]
lo, hi = float(arr.min()), float(arr.max())
if hi - lo < 1e-6:
arr = np.zeros_like(arr)
else:
arr = (arr - lo) / (hi - lo)
return (arr * 255.0).clip(0, 255).astype(np.uint8)
|