| """Monocular depth estimation.""" |
|
|
| from __future__ import annotations |
|
|
| import threading |
| from pathlib import Path |
|
|
| import cv2 |
| import numpy as np |
|
|
|
|
| def _pick_device() -> str: |
| """Best available inference device.""" |
| try: |
| import torch |
| except ImportError: |
| return "cpu" |
| try: |
| if torch.cuda.is_available(): |
| return "cuda" |
| if torch.backends.mps.is_available(): |
| return "mps" |
| except (AttributeError, RuntimeError): |
| return "cpu" |
| return "cpu" |
|
|
|
|
| class DepthEstimator: |
| """Monocular depth from a YOLO depth model.""" |
|
|
| def __init__(self, model_path: str | Path, imgsz: int = 384) -> None: |
| from ultralytics import YOLO |
|
|
| model_path = Path(model_path) |
| if not model_path.exists(): |
| raise FileNotFoundError(f"Depth model not found: {model_path}") |
| self.device = _pick_device() |
| self.model = YOLO(str(model_path)) |
| self.imgsz = imgsz |
| self._lo: float | None = None |
| self._hi: float | None = None |
| self._bounds_alpha = 0.08 |
|
|
| def __call__(self, frame: np.ndarray) -> np.ndarray: |
| """Raw depth map, same size as frame.""" |
| result = self.model.predict( |
| frame, imgsz=self.imgsz, device=self.device, verbose=False |
| )[0] |
| depth = result.depth.data |
| if hasattr(depth, "cpu"): |
| depth = depth.cpu().numpy() |
| depth = np.asarray(depth, dtype=np.float32) |
| if depth.ndim == 3: |
| depth = depth[0] |
| if depth.shape[:2] != frame.shape[:2]: |
| depth = cv2.resize(depth, (frame.shape[1], frame.shape[0]), |
| interpolation=cv2.INTER_LINEAR) |
| return depth |
|
|
| def normalize(self, depth: np.ndarray) -> np.ndarray: |
| """Stable 0..1 map, 0 near, 1 far.""" |
| finite = depth[np.isfinite(depth)] |
| if finite.size == 0: |
| return np.full(depth.shape, 0.5, dtype=np.float32) |
| lo, hi = float(np.percentile(finite, 2.0)), float(np.percentile(finite, 98.0)) |
| if self._lo is None or self._hi is None: |
| self._lo, self._hi = lo, hi |
| else: |
| self._lo += self._bounds_alpha * (lo - self._lo) |
| self._hi += self._bounds_alpha * (hi - self._hi) |
| span = max(self._hi - self._lo, 1e-6) |
| norm = np.clip((depth - self._lo) / span, 0.0, 1.0) |
| return np.nan_to_num(norm, nan=0.5).astype(np.float32) |
|
|
| @staticmethod |
| def colorize(depth_norm: np.ndarray) -> np.ndarray: |
| """Turbo colormap of a normalized depth map.""" |
| u8 = (np.clip(depth_norm, 0.0, 1.0) * 255.0).astype(np.uint8) |
| return cv2.applyColorMap(255 - u8, cv2.COLORMAP_TURBO) |
|
|
| @staticmethod |
| def sample(depth_map: np.ndarray, x: float, y: float, radius: int = 9, |
| percentile: float = 20.0, default: float = 0.0) -> float: |
| """Nearest surface around a point.""" |
| h, w = depth_map.shape[:2] |
| xi = int(np.clip(round(float(x)), 0, w - 1)) |
| yi = int(np.clip(round(float(y)), 0, h - 1)) |
| x0, x1 = max(0, xi - radius), min(w, xi + radius + 1) |
| y0, y1 = max(0, yi - radius), min(h, yi + radius + 1) |
| patch = depth_map[y0:y1, x0:x1] |
| patch = patch[np.isfinite(patch)] |
| if patch.size == 0: |
| return default |
| return float(np.percentile(patch, percentile)) |
|
|
|
|
| class DepthWorker: |
| """Background depth estimation thread.""" |
|
|
| def __init__(self, model_path: str | Path, imgsz: int = 384, |
| input_width: int = 640) -> None: |
| self.estimator = DepthEstimator(model_path, imgsz) |
| self.device = self.estimator.device |
| self.input_width = input_width |
| self._pending: np.ndarray | None = None |
| self._target: tuple[int, int] | None = None |
| self._metric: np.ndarray | None = None |
| self._norm: np.ndarray | None = None |
| self._seq = 0 |
| self._lock = threading.Lock() |
| self._wake = threading.Event() |
| self._stop = threading.Event() |
| self._thread = threading.Thread(target=self._loop, daemon=True) |
| self._thread.start() |
|
|
| def submit(self, frame: np.ndarray) -> None: |
| """Queue the newest frame.""" |
| h, w = frame.shape[:2] |
| if w > self.input_width: |
| k = self.input_width / float(w) |
| small = cv2.resize(frame, (self.input_width, max(1, int(round(h * k)))), |
| interpolation=cv2.INTER_AREA) |
| else: |
| small = frame.copy() |
| with self._lock: |
| self._pending = small |
| self._target = (w, h) |
| self._wake.set() |
|
|
| def _loop(self) -> None: |
| """Consume frames until stopped.""" |
| while not self._stop.is_set(): |
| self._wake.wait(0.1) |
| self._wake.clear() |
| with self._lock: |
| frame, target = self._pending, self._target |
| self._pending = None |
| if frame is None or target is None: |
| continue |
| try: |
| metric = self.estimator(frame) |
| norm = self.estimator.normalize(metric) |
| except Exception as exc: |
| print(f"[!] depth failed: {exc}", flush=True) |
| self._stop.set() |
| return |
| if (metric.shape[1], metric.shape[0]) != target: |
| metric = cv2.resize(metric, target, interpolation=cv2.INTER_LINEAR) |
| norm = cv2.resize(norm, target, interpolation=cv2.INTER_LINEAR) |
| with self._lock: |
| self._metric = metric |
| self._norm = norm |
| self._seq += 1 |
|
|
| @property |
| def latest(self) -> np.ndarray | None: |
| """Newest metric depth map.""" |
| with self._lock: |
| return self._metric |
|
|
| @property |
| def latest_norm(self) -> np.ndarray | None: |
| """Newest normalized depth map.""" |
| with self._lock: |
| return self._norm |
|
|
| @property |
| def ready(self) -> bool: |
| """A depth map is available.""" |
| return self.latest is not None |
|
|
| @property |
| def frames(self) -> int: |
| """Number of finished estimations.""" |
| with self._lock: |
| return self._seq |
|
|
| def close(self) -> None: |
| """Stop the worker thread.""" |
| self._stop.set() |
| self._wake.set() |
| self._thread.join(timeout=1.0) |
|
|
