lsf/depth.py

"""Depth estimation and dual-pixel handling.

What "dual-pixel depth" actually is
-----------------------------------
A dual-pixel (DP) sensor splits every photosite into two halves, giving two
slightly different sub-images from which a *disparity* (depth) map can be
computed — the trick behind Google Pixel "Portrait mode". DP data is exposed by
the *camera stack* (e.g. Android Camera2 `RAW`/`DEPTH` streams, or a depth file
embedded in the capture), **not** by a generic UVC webcam through OpenCV.

So in practice:

  * On a phone / device that publishes a depth or DP stream, feed that map into
    `DepthEstimator.from_stream(...)` and we sample it at each landmark.
  * On a normal laptop webcam there is no DP data. We *detect that*, and
    optionally fall back to a monocular depth model (MiDaS) if `torch` is
    installed. Otherwise we transparently use MediaPipe's own relative z.

Either way, the recogniser keeps working; depth only refines the per-landmark
z so movements toward/away from the camera are measured more accurately.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Optional

import numpy as np

from .landmarks import LandmarkFrame


@dataclass
class DepthInfo:
    """Where a frame's depth came from and the map itself (if any)."""

    source: str                      # "dual_pixel" | "monocular" | "mediapipe"
    available: bool                  # True if a real depth map backs this frame
    depth_map: Optional[np.ndarray]  # (H, W) float32, larger = farther, or None


def detect_dual_pixel_support(capture=None) -> tuple[bool, str]:
    """Best-effort check for a usable dual-pixel / depth stream.

    OpenCV's VideoCapture cannot surface dual-pixel data, so for a plain webcam
    this returns False with an explanation. The hook exists so a device backend
    that *does* provide depth can report support here.
    """
    if capture is None:
        return False, "No capture device supplied."

    # Some depth cameras expose CAP_PROP_OPENNI_* / a second stream; a UVC
    # webcam does not. We treat presence of a non-trivial depth generator as
    # support. This stays False for ordinary webcams.
    try:
        import cv2

        has_depth = capture.get(cv2.CAP_PROP_OPENNI_REGISTRATION) not in (-1, 0)
        if has_depth:
            return True, "Depth-capable capture detected."
    except Exception:
        pass
    return (
        False,
        "Dual-pixel/depth stream not exposed by this camera via OpenCV; "
        "using monocular fallback or MediaPipe z.",
    )


class DepthEstimator:
    """Produces a per-frame depth map and refines landmark z values.

    Priority order:
      1. An externally supplied dual-pixel/depth map (best).
      2. MiDaS monocular estimate, if `torch` + `timm` are installed and enabled.
      3. None — landmarks keep MediaPipe's relative z.
    """

    def __init__(self, *, enable_monocular: bool = False) -> None:
        self._supports_dp = False
        self._dp_reason = "uninitialised"
        self._midas = None
        self._midas_transform = None
        self._enable_monocular = enable_monocular
        if enable_monocular:
            self._try_load_midas()

    def probe(self, capture) -> None:
        """Inspect a capture device for dual-pixel support."""
        self._supports_dp, self._dp_reason = detect_dual_pixel_support(capture)

    @property
    def status(self) -> str:
        if self._supports_dp:
            return "dual_pixel"
        if self._midas is not None:
            return "monocular"
        return "mediapipe"

    @property
    def reason(self) -> str:
        return self._dp_reason

    # -- depth map sources ----------------------------------------------------
    def _try_load_midas(self) -> None:
        """Load MiDaS small from torch.hub. Silently no-ops if torch missing."""
        try:
            import torch  # noqa: F401

            self._midas = torch.hub.load("intel-isl/MiDaS", "MiDaS_small")
            self._midas.eval()
            transforms = torch.hub.load("intel-isl/MiDaS", "transforms")
            self._midas_transform = transforms.small_transform
        except Exception as exc:  # pragma: no cover - optional path
            self._midas = None
            self._dp_reason = f"Monocular depth unavailable: {exc}"

    def _monocular_depth(self, rgb_frame: np.ndarray) -> Optional[np.ndarray]:
        if self._midas is None:
            return None
        import torch

        with torch.no_grad():
            batch = self._midas_transform(rgb_frame)
            pred = self._midas(batch)
            pred = torch.nn.functional.interpolate(
                pred.unsqueeze(1),
                size=rgb_frame.shape[:2],
                mode="bicubic",
                align_corners=False,
            ).squeeze()
        depth = pred.cpu().numpy().astype(np.float32)
        # MiDaS returns inverse depth (larger = closer); invert so larger=farther.
        depth = depth.max() - depth
        return depth

    def depth_for_frame(
        self,
        rgb_frame: np.ndarray,
        external_depth: Optional[np.ndarray] = None,
    ) -> DepthInfo:
        """Return depth info for a frame, choosing the best available source."""
        if external_depth is not None:
            return DepthInfo("dual_pixel", True, external_depth.astype(np.float32))
        mono = self._monocular_depth(rgb_frame) if self._midas is not None else None
        if mono is not None:
            return DepthInfo("monocular", True, mono)
        return DepthInfo("mediapipe", False, None)

    # -- landmark refinement --------------------------------------------------
    @staticmethod
    def refine_landmarks(frame: LandmarkFrame, depth: DepthInfo) -> LandmarkFrame:
        """Overwrite each landmark's z with a metric depth sample when we have one.

        Landmark x/y are normalised [0, 1]; we map them to pixel coordinates to
        sample the depth map. Values are normalised to ~[0, 1] so the feature
        scale matches MediaPipe's native z range.
        """
        if not depth.available or depth.depth_map is None:
            return frame

        dmap = depth.depth_map
        h, w = dmap.shape[:2]
        dmin, dmax = float(dmap.min()), float(dmap.max())
        span = (dmax - dmin) or 1.0

        def sample(arr: Optional[np.ndarray]) -> Optional[np.ndarray]:
            if arr is None:
                return None
            out = arr.copy()
            xs = np.clip((arr[:, 0] * w).astype(int), 0, w - 1)
            ys = np.clip((arr[:, 1] * h).astype(int), 0, h - 1)
            out[:, 2] = (dmap[ys, xs] - dmin) / span
            return out

        return LandmarkFrame(
            pose=sample(frame.pose),
            face=sample(frame.face),
            left_hand=sample(frame.left_hand),
            right_hand=sample(frame.right_hand),
            raw=frame.raw,
        )