Spaces:
Paused
Paused
File size: 4,622 Bytes
728fc83 | 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 | """
models/depth_estimation.py
──────────────────────────
TransformersDepthLoader — wraps any HuggingFace depth-estimation
model that follows the transformers pipeline API (Depth Anything v2,
DPT-Large, MiDaS, ZoeDepth, etc.).
Inputs (run kwargs)
──────────────────────────────────────────────────────────────
image : PIL.Image (RGB) — the image to estimate depth for
Outputs (returned dict)
──────────────────────────────────────────────────────────────
depth_raw : np.ndarray (H×W float32, unnormalised)
depth_normalised : np.ndarray (H×W float32 in [0,1])
depth_colourmap : PIL.Image (false-colour for display)
depth_uint16 : PIL.Image (16-bit greyscale PNG for export)
model : str
"""
from __future__ import annotations
import logging
import os
from typing import Any
import numpy as np
from PIL import Image
from models.base_loader import BaseLoader
from utils.image_utils import normalise_depth, depth_to_colormap, depth_to_uint16, resize_image
logger = logging.getLogger(__name__)
def _hf_token() -> str | None:
return os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN") or None
class TransformersDepthLoader(BaseLoader):
def load(self) -> None:
if self._loaded:
return
logger.info("Loading depth model: %s", self.model_id)
from transformers import pipeline as hf_pipeline
token = _hf_token()
pipeline_kwargs: dict = {
"task": "depth-estimation",
"model": self.model_id,
"device": self.device,
}
if token:
pipeline_kwargs["token"] = token
try:
self.pipe = hf_pipeline(**pipeline_kwargs)
except (OSError, EnvironmentError) as hub_err:
err_str = str(hub_err).lower()
if any(k in err_str for k in ("not cached", "fetch metadata", "connection", "network",
"offline", "cannot reach", "name or service not known")):
logger.warning(
"Hub unreachable for depth model %s. Retrying with local_files_only=True …",
self.model_id,
)
try:
self.pipe = hf_pipeline(**{**pipeline_kwargs, "model": self.model_id},
local_files_only=True)
logger.info("Loaded depth model %s from local cache.", self.model_id)
except Exception as cache_err:
raise OSError(
f"Cannot load depth model {self.model_id}: Hub unreachable and no local cache.\n"
f" Hub error : {hub_err}\n"
f" Cache error: {cache_err}\n"
"Tip: ensure internet access on first run, or set HF_TOKEN for gated models."
) from cache_err
else:
raise
self._loaded = True
logger.info("Depth model ready: %s", self.model_id)
def run(self, **inputs: Any) -> dict[str, Any]:
if not self._loaded:
self.load()
image: Image.Image = inputs["image"]
image_size: int = int(self.kwargs.get("image_size", 518))
# Resize to the model's preferred resolution
image = resize_image(image, max_side=image_size).convert("RGB")
logger.info("Estimating depth image_size=%d model=%s", image_size, self.model_id)
result = self.pipe(image)
# HF depth-estimation pipeline returns {"depth": PIL.Image, "predicted_depth": tensor}
depth_tensor = result.get("predicted_depth")
if depth_tensor is not None:
import torch
depth_raw = depth_tensor.squeeze().detach().cpu().numpy().astype(np.float32)
else:
# Fallback: use the greyscale PIL image
depth_raw = np.array(result["depth"]).astype(np.float32)
depth_norm = normalise_depth(depth_raw)
depth_colour = depth_to_colormap(depth_raw)
depth_16 = depth_to_uint16(depth_raw)
return {
"depth_raw": depth_raw,
"depth_normalised": depth_norm,
"depth_colourmap": depth_colour,
"depth_uint16": depth_16,
"model": self.model_id,
}
|