Upload depthpro_wrapper/depth_estimator.py

depthpro_wrapper/depth_estimator.py

@@ -0,0 +1,241 @@
+"""
+Core DepthPro depth estimation wrapper.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Optional, Union
+
+import numpy as np
+import torch
+from PIL import Image
+
+
+try:
+    from transformers import DepthProImageProcessorFast, DepthProForDepthEstimation
+except ImportError as exc:
+    raise ImportError(
+        "transformers>=4.40.0 is required. "
+        "Install it with:  pip install transformers[torch]"
+    ) from exc
+
+
+@dataclass
+class DepthResult:
+    """Container for depth estimation results."""
+
+    depth: np.ndarray
+    """(H, W) float32 metric depth map in meters."""
+
+    focal_length: float
+    """Estimated focal length in pixels (for the original image resolution)."""
+
+    field_of_view: float
+    """Estimated horizontal field of view in degrees."""
+
+    image: np.ndarray
+    """(H, W, 3) uint8 RGB image at original resolution."""
+
+    confidence: Optional[np.ndarray] = None
+    """(H, W) optional confidence / uncertainty map."""
+
+    @property
+    def height(self) -> int:
+        return self.depth.shape[0]
+
+    @property
+    def width(self) -> int:
+        return self.depth.shape[1]
+
+
+class DepthProEstimator:
+    """
+    High-level wrapper around Apple's DepthPro for metric depth estimation.
+
+    DepthPro is a single-image zero-shot metric depth estimator.  Unlike
+    relative-depth models (e.g. Depth Anything), it outputs **absolute-scale
+    metric depth in meters** and also estimates the **camera focal length**
+    and **field of view** automatically — no calibration required.
+
+    Parameters
+    ----------
+    model_name : str, default "apple/DepthPro-hf"
+        HuggingFace model id or local path.
+    device : str or torch.device, default "cuda:0"
+        PyTorch device.  CUDA strongly recommended for 952M-parameter ViT-L.
+    dtype : torch.dtype, default torch.float16
+        Inference dtype.  fp16 halves memory and is the default the model
+        was trained with; fp32 gives marginally higher precision.
+    """
+
+    _MODEL_INPUT_SIZE: int = 1536  # DepthPro always processes at 1536×1536
+
+    def __init__(
+        self,
+        model_name: str = "apple/DepthPro-hf",
+        device: Union[str, torch.device] = "cuda:0",
+        dtype: torch.dtype = torch.float16,
+    ):
+        self.device = torch.device(device)
+        self.dtype = dtype
+        self.model_name = model_name
+
+        if not torch.cuda.is_available() and self.device.type == "cuda":
+            raise RuntimeError(
+                "CUDA is not available but device='cuda' was requested. "
+                "DepthPro is a 952M ViT-L model; CPU inference will be extremely slow. "
+                "Pass device='cpu' explicitly if you really want this."
+            )
+
+        self._load_model()
+
+    # ------------------------------------------------------------------ #
+    #  Loading                                                           #
+    # ------------------------------------------------------------------ #
+
+    def _load_model(self) -> None:
+        """Load the processor and model from HF."""
+        self.processor = DepthProImageProcessorFast.from_pretrained(self.model_name)
+        self.model = DepthProForDepthEstimation.from_pretrained(
+            self.model_name,
+            torch_dtype=self.dtype,
+        ).to(self.device)
+        self.model.eval()
+
+    # ------------------------------------------------------------------ #
+    #  Public API                                                        #
+    # ------------------------------------------------------------------ #
+
+    @torch.no_grad()
+    def estimate(
+        self,
+        image: Union[str, Path, Image.Image, np.ndarray],
+        *,
+        return_confidence: bool = False,
+    ) -> DepthResult:
+        """
+        Run metric depth estimation on a single RGB image.
+
+        Parameters
+        ----------
+        image : str, Path, PIL.Image, or np.ndarray
+            Input RGB image.  If a path, loaded with PIL.
+        return_confidence : bool, default False
+            If True and the model provides a confidence map, include it in
+            the result.
+
+        Returns
+        -------
+        DepthResult
+            Dataclass containing:
+            * ``depth`` — (H, W) metric depth in meters
+            * ``focal_length`` — estimated focal length (px) at original res
+            * ``field_of_view`` — estimated horizontal FOV in degrees
+            * ``image`` — original RGB image as (H, W, 3) uint8
+        """
+        # ---- load / normalise image ---------------------------------------
+        pil_image = self._to_pil(image)
+        rgb_array = np.array(pil_image.convert("RGB"), dtype=np.uint8)
+
+        # ---- preprocess & forward pass ------------------------------------
+        inputs = self.processor(images=pil_image, return_tensors="pt")
+        inputs = {k: v.to(self.device) for k, v in inputs.items()}
+
+        outputs = self.model(**inputs)
+
+        # ---- post-process back to original resolution ---------------------
+        post = self.processor.post_process_depth_estimation(
+            outputs,
+            target_sizes=[(pil_image.height, pil_image.width)],
+        )[0]
+
+        depth = post["predicted_depth"]  # torch.Tensor [H, W]
+        focal_length = post["focal_length"]
+        fov = post["field_of_view"]
+
+        # DepthPro returns focal_length for the 1536×1536 processed image.
+        # Scale it to the original image resolution.
+        focal_original = focal_length * (pil_image.width / self._MODEL_INPUT_SIZE)
+
+        depth_np = depth.cpu().float().numpy()
+        confidence_np = None
+        if return_confidence and "confidence" in post:
+            confidence_np = post["confidence"].cpu().float().numpy()
+
+        return DepthResult(
+            depth=depth_np,
+            focal_length=focal_original.item(),
+            field_of_view=fov.item(),
+            image=rgb_array,
+            confidence=confidence_np,
+        )
+
+    @torch.no_grad()
+    def estimate_batch(
+        self,
+        images: list,
+        *,
+        return_confidence: bool = False,
+    ) -> list[DepthResult]:
+        """
+        Run depth estimation on a batch of images.
+
+        Parameters
+        ----------
+        images : list of str, Path, PIL.Image, or np.ndarray
+        return_confidence : bool
+
+        Returns
+        -------
+        list[DepthResult]
+        """
+        pil_images = [self._to_pil(img) for img in images]
+        rgb_arrays = [np.array(p.convert("RGB"), dtype=np.uint8) for p in pil_images]
+
+        # Processor batch handling
+        inputs = self.processor(images=pil_images, return_tensors="pt")
+        inputs = {k: v.to(self.device) for k, v in inputs.items()}
+
+        outputs = self.model(**inputs)
+
+        posts = self.processor.post_process_depth_estimation(
+            outputs,
+            target_sizes=[(p.height, p.width) for p in pil_images],
+        )
+
+        results = []
+        for post, pil_img, rgb in zip(posts, pil_images, rgb_arrays):
+            depth = post["predicted_depth"].cpu().float().numpy()
+            focal = post["focal_length"] * (pil_img.width / self._MODEL_INPUT_SIZE)
+            fov = post["field_of_view"]
+            conf = None
+            if return_confidence and "confidence" in post:
+                conf = post["confidence"].cpu().float().numpy()
+
+            results.append(DepthResult(
+                depth=depth,
+                focal_length=focal.item(),
+                field_of_view=fov.item(),
+                image=rgb,
+                confidence=conf,
+            ))
+        return results
+
+    # ------------------------------------------------------------------ #
+    #  Helpers                                                           #
+    # ------------------------------------------------------------------ #
+
+    @staticmethod
+    def _to_pil(image: Union[str, Path, Image.Image, np.ndarray]) -> Image.Image:
+        """Normalise input to a PIL RGB image."""
+        if isinstance(image, (str, Path)):
+            return Image.open(str(image)).convert("RGB")
+        if isinstance(image, np.ndarray):
+            if image.dtype != np.uint8:
+                image = (image * 255).clip(0, 255).astype(np.uint8)
+            return Image.fromarray(image).convert("RGB")
+        if isinstance(image, Image.Image):
+            return image.convert("RGB")
+        raise TypeError(f"Unsupported image type: {type(image)}")