Create tools/self_check.py

tools/self_check.py

@@ -0,0 +1,155 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Quiet, one-shot startup self-check for HF Spaces.
+
+What it does:
+- loads SAM2Loader + MatAnyoneLoader (device from env or cuda/cpu auto)
+- runs a minimal first-frame path (synthetic frame) to validate
+- caches status in module state for later UI queries
+- does NOT print unless failure; logs via `BackgroundFX`/root logger
+
+Control via env:
+- DISABLE_SELF_CHECK=1         → skip entirely
+- SELF_CHECK_DEVICE=cpu|cuda   → override device
+- SELF_CHECK_TIMEOUT=seconds   → default 45
+"""
+
+from __future__ import annotations
+import os, time, threading, logging
+from typing import Optional, Dict, Any
+
+import numpy as np
+import cv2
+import torch
+
+# Import loaders and processor from your project
+from models.loaders.sam2_loader import SAM2Loader
+from models.loaders.matanyone_loader import MatAnyoneLoader
+from processing.two_stage.two_stage_processor import TwoStageProcessor
+
+logger = logging.getLogger("BackgroundFX") or logging.getLogger(__name__)
+
+# Module-level cache
+_SELF_CHECK_LOCK = threading.Lock()
+_SELF_CHECK_DONE = False
+_SELF_CHECK_OK = False
+_SELF_CHECK_MSG = "Self-check did not run yet."
+_SELF_CHECK_DURATION = 0.0
+
+def _pick_device() -> str:
+    dev = os.environ.get("SELF_CHECK_DEVICE", "").strip().lower()
+    if dev in ("cpu", "cuda"):
+        return dev
+    return "cuda" if torch.cuda.is_available() else "cpu"
+
+def _synth_frame(w=640, h=360) -> np.ndarray:
+    """
+    Create a simple BGR frame with a 'person-like' central blob over green.
+    We just need a plausible image; quality doesn’t matter for self-check.
+    """
+    img = np.zeros((h, w, 3), np.uint8)
+    # base: gray
+    img[:] = (40, 40, 40)
+    # put a green screen-like area on right to make chroma pass exercise
+    cv2.rectangle(img, (int(0.65*w), 0), (w, h), (0, 255, 0), -1)
+    # draw a central "person" blob
+    cx, cy = w//3, h//2
+    cv2.ellipse(img, (cx, cy-40), (35, 45), 0, 0, 360, (60, 60, 200), -1)  # head-ish
+    cv2.rectangle(img, (cx-40, cy-10), (cx+40, cy+80), (60, 60, 200), -1)  # torso-ish
+    return img
+
+def _run_once(timeout_s: float = 45.0) -> tuple[bool, str, float]:
+    t0 = time.time()
+    device = _pick_device()
+    try:
+        # 1) Load SAM2
+        sam = SAM2Loader(device=device).load("auto")
+        if sam is None:
+            return False, "SAM2 failed to load", time.time()-t0
+
+        # 2) Get synthetic frame
+        bgr = _synth_frame()
+        sam.set_image(bgr)
+        out = sam.predict(point_coords=None, point_labels=None)
+        masks = out.get("masks", None)
+        h, w = bgr.shape[:2]
+        if masks is None or len(masks) == 0:
+            logger.warning("Self-check: SAM2 returned no masks; accepting fallback.")
+            mask0 = np.ones((h, w), np.float32)
+        else:
+            mask0 = masks[0].astype(np.float32)
+            if mask0.shape != (h, w):
+                mask0 = cv2.resize(mask0, (w, h), interpolation=cv2.INTER_LINEAR)
+
+        # 3) Load MatAnyone stateful session
+        session = MatAnyoneLoader(device=device).load()
+        if session is None:
+            return False, "MatAnyone failed to load", time.time()-t0
+
+        # 4) Bootstrap (frame 0 must have a mask; fallback already ensured)
+        rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
+        alpha0 = session(rgb, mask0)
+        if not isinstance(alpha0, np.ndarray) or alpha0.shape != (h, w):
+            return False, f"MatAnyone alpha shape unexpected: {getattr(alpha0, 'shape', None)}", time.time()-t0
+
+        # 5) Minimal TwoStageProcessor wiring (no file IO, just instantiate)
+        _ = TwoStageProcessor(sam2_predictor=sam, matanyone_model=session)
+
+        return True, "OK", time.time()-t0
+
+    except Exception as e:
+        return False, f"Self-check error: {e}", time.time()-t0
+    finally:
+        # crude timeout enforcement info (the thread is joined by caller)
+        dur = time.time()-t0
+        if dur > timeout_s:
+            logger.warning(f"Self-check exceeded timeout {timeout_s:.1f}s (took {dur:.2f}s)")
+        return locals().get("sam", None) is not None and locals().get("session", None) is not None, \
+               locals().get("e", None) and f"Self-check error: {e}" or "OK", \
+               dur
+
+def _runner(timeout_s: float):
+    global _SELF_CHECK_DONE, _SELF_CHECK_OK, _SELF_CHECK_MSG, _SELF_CHECK_DURATION
+    ok, msg, dur = _run_once(timeout_s=timeout_s)
+    with _SELF_CHECK_LOCK:
+        _SELF_CHECK_DONE = True
+        _SELF_CHECK_OK = bool(ok and msg == "OK")
+        _SELF_CHECK_MSG = msg
+        _SELF_CHECK_DURATION = float(dur)
+    if _SELF_CHECK_OK:
+        logger.info(f"✅ Startup self-check OK in {dur:.2f}s")
+    else:
+        logger.error(f"❌ Startup self-check FAILED in {dur:.2f}s: {msg}")
+
+def launch_self_check_async(timeout_s: Optional[float] = None):
+    """
+    Fire-and-forget startup check. No effect if disabled or already started.
+    """
+    if os.environ.get("DISABLE_SELF_CHECK", "0") == "1":
+        logger.info("Self-check disabled via DISABLE_SELF_CHECK=1")
+        with _SELF_CHECK_LOCK:
+            global _SELF_CHECK_DONE, _SELF_CHECK_OK, _SELF_CHECK_MSG, _SELF_CHECK_DURATION
+            _SELF_CHECK_DONE = True
+            _SELF_CHECK_OK = True
+            _SELF_CHECK_MSG = "Disabled"
+            _SELF_CHECK_DURATION = 0.0
+        return
+
+    timeout_s = float(os.environ.get("SELF_CHECK_TIMEOUT", str(timeout_s or 45.0)))
+    # Only launch once
+    with _SELF_CHECK_LOCK:
+        if getattr(launch_self_check_async, "_started", False):
+            return
+        launch_self_check_async._started = True  # type: ignore[attr-defined]
+    th = threading.Thread(target=_runner, args=(timeout_s,), daemon=True)
+    th.start()
+
+def get_self_check_status() -> Dict[str, Any]:
+    with _SELF_CHECK_LOCK:
+        return {
+            "done": _SELF_CHECK_DONE,
+            "ok": _SELF_CHECK_OK,
+            "message": _SELF_CHECK_MSG,
+            "duration": _SELF_CHECK_DURATION,
+        }