Spaces:
Running
Running
| """ | |
| 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, | |
| } | |