livekit-cameras/hand_raise.py

"""Cheap MoveNet Lightning hand-raise detector for the KMS wall.

OpenVINO CPU, FP32 single-pose, ~4 Hz per tile. Does not touch capture/encode.
Arc/iGPU stay unused (kmssink + H.264).
"""
from __future__ import annotations

import logging
import threading
import time
from collections import deque
from pathlib import Path
from typing import Callable, Optional

import numpy as np

log = logging.getLogger("cameras.display")

MODEL_PATH = Path(__file__).resolve().parent / "models" / "movenet_singlepose_lightning.onnx"
MODELS_DIR = Path(__file__).resolve().parent / "models"
INPUT_SIZE = 192
MOVENET = {
    "lightning": ("movenet_singlepose_lightning.onnx", 192),
    "thunder": ("movenet_singlepose_thunder.onnx", 256),
}


def resolve_movenet(
    name: str = "thunder",
    models_dir: Path | None = None,
) -> tuple[Path, int, str]:
    """lightning=192, thunder=256. Same Xenova FP32 ONNX family as production Lightning."""
    key = (name or "thunder").strip().lower()
    if key not in MOVENET:
        raise ValueError(f"DISPLAY_HAND_RAISE_MODEL must be lightning|thunder, got {name!r}")
    fname, size = MOVENET[key]
    root = Path(models_dir) if models_dir is not None else MODELS_DIR
    return root / fname, int(size), key

# MoveNet COCO-17, output is (y, x, score) in [0, 1], y grows downward.
NOSE = 0
L_SHOULDER, R_SHOULDER = 5, 6
L_ELBOW, R_ELBOW = 7, 8
L_WRIST, R_WRIST = 9, 10
MIN_SCORE = 0.20
# Wrist must sit this far above the shoulder (normalized image y).
LIFT = 0.05
FACE_RADIUS2 = 0.12 * 0.12


def letterbox_rgb(rgb: np.ndarray, size: int = INPUT_SIZE) -> np.ndarray:
    """Pad to square so MoveNet is not stretched (16:9 C920 tiles)."""
    import cv2

    size = int(size)
    h, w = rgb.shape[:2]
    if h < 1 or w < 1:
        return np.zeros((size, size, 3), dtype=np.uint8)
    scale = size / float(max(h, w))
    nw = max(1, int(round(w * scale)))
    nh = max(1, int(round(h * scale)))
    resized = cv2.resize(rgb, (nw, nh), interpolation=cv2.INTER_AREA)
    out = np.zeros((size, size, 3), dtype=np.uint8)
    y0 = (size - nh) // 2
    x0 = (size - nw) // 2
    out[y0:y0 + nh, x0:x0 + nw] = resized
    return out


def wrist_is_raised(kpts: np.ndarray) -> bool:
    """True if either wrist is clearly up — above the shoulder, or above the head."""
    if kpts is None or kpts.shape != (17, 3):
        return False
    k = kpts
    if k[NOSE, 2] < MIN_SCORE and k[L_SHOULDER, 2] < MIN_SCORE and k[R_SHOULDER, 2] < MIN_SCORE:
        return False
    return (
        _arm_up(k, L_WRIST, L_ELBOW, L_SHOULDER)
        or _arm_up(k, R_WRIST, R_ELBOW, R_SHOULDER)
        or _wrist_above_head(k)
    )


def _near_face(k: np.ndarray, wr: np.ndarray) -> bool:
    if k[NOSE, 2] < MIN_SCORE:
        return False
    dy = wr[0] - k[NOSE, 0]
    dx = wr[1] - k[NOSE, 1]
    return dx * dx + dy * dy < FACE_RADIUS2


def _wrist_above_head(k: np.ndarray) -> bool:
    if k[NOSE, 2] < MIN_SCORE:
        return False
    for wi in (L_WRIST, R_WRIST):
        wr = k[wi]
        if wr[2] < MIN_SCORE:
            continue
        if wr[0] >= k[NOSE, 0] - LIFT:
            continue
        if _near_face(k, wr):
            continue
        return True
    return False


def _arm_up(k: np.ndarray, wi: int, ei: int, si: int) -> bool:
    wr, el, sh = k[wi], k[ei], k[si]
    if wr[2] < MIN_SCORE or sh[2] < MIN_SCORE:
        return False
    if wr[0] >= sh[0] - LIFT:
        return False
    if el[2] >= MIN_SCORE and wr[0] >= el[0]:
        return False
    if _near_face(k, wr):
        return False
    return True


class RaiseLatch:
    """Vote `hits` in a sliding `window`; stay lit until `hold_s` without a hit."""

    def __init__(self, hits: int = 3, hold_s: float = 1.0, window: int = 5) -> None:
        self.hits = max(1, int(hits))
        self.hold_s = max(0.0, float(hold_s))
        self.window = max(self.hits, int(window))
        self._buf: deque[bool] = deque(maxlen=self.window)
        self.raised = False
        self.last_true = 0.0

    def update(self, pred: bool, now: float) -> bool:
        self._buf.append(bool(pred))
        if pred:
            self.last_true = now
        if sum(self._buf) >= self.hits:
            self.raised = True
        elif self.raised and (now - self.last_true) >= self.hold_s:
            self.raised = False
        return self.raised


class HandRaiseMonitor:
    """Background 2 Hz pose on the latest I420 tile per identity."""

    def __init__(
        self,
        enabled: bool = True,
        hz: float = 4.0,
        hold_s: float = 1.0,
        on_change: Optional[Callable[[str, bool], None]] = None,
        model_path: Optional[Path] = None,
        model: str = "thunder",
        input_size: Optional[int] = None,
    ) -> None:
        self.on_change = on_change
        self.hz = min(10.0, max(0.5, float(hz)))
        self.hold_s = float(hold_s)
        if model_path is not None:
            self._path = Path(model_path)
            self._variant = (model or "custom").strip().lower()
            self._input_size = int(input_size or (256 if "thunder" in self._path.name else 192))
        else:
            self._path, self._input_size, self._variant = resolve_movenet(model)
            if input_size is not None:
                self._input_size = int(input_size)
        self._lock = threading.Lock()
        self._latest: dict[str, tuple[bytes, int, int]] = {}
        self._latch: dict[str, RaiseLatch] = {}
        self._raised: set[str] = set()
        self._stop = threading.Event()
        self._thread: Optional[threading.Thread] = None
        self._request = None
        self._ov_input = None
        self.enabled = bool(enabled) and self._load()

    def _load(self) -> bool:
        if not self._path.is_file():
            log.warning("hand-raise model missing: %s", self._path)
            return False
        try:
            import openvino as ov

            core = ov.Core()
            model = core.read_model(str(self._path))
            compiled = core.compile_model(
                model,
                "CPU",
                {"INFERENCE_NUM_THREADS": 1, "NUM_STREAMS": 1},
            )
            self._request = compiled.create_infer_request()
            self._ov_input = compiled.input(0)
            log.info(
                "hand-raise MoveNet %s FP32 OpenVINO CPU hz=%.1f hold=%.1fs input=%d model=%s",
                self._variant, self.hz, self.hold_s, self._input_size, self._path.name,
            )
            return True
        except Exception as exc:  # noqa: BLE001
            log.warning("hand-raise model load failed: %s", exc)
            self._request = None
            return False

    def start(self) -> None:
        if not self.enabled or self._thread is not None:
            return
        self._thread = threading.Thread(
            target=self._run, name="hand-raise", daemon=True)
        self._thread.start()

    def stop(self) -> None:
        self._stop.set()
        t = self._thread
        if t is not None and t.is_alive():
            t.join(timeout=2.0)
        self._thread = None

    def offer(self, identity: str, i420: bytes | memoryview, width: int, height: int) -> None:
        if not self.enabled or width < 16 or height < 16:
            return
        blob = i420 if isinstance(i420, (bytes, bytearray)) else bytes(i420)
        with self._lock:
            self._latest[identity] = (blob, int(width), int(height))

    def forget(self, identity: str) -> None:
        with self._lock:
            self._latest.pop(identity, None)
            self._latch.pop(identity, None)
            was = identity in self._raised
            self._raised.discard(identity)
        if was and self.on_change is not None:
            try:
                self.on_change(identity, False)
            except Exception:  # noqa: BLE001
                log.exception("hand-raise on_change(%s, False)", identity)

    def raised(self) -> set[str]:
        with self._lock:
            return set(self._raised)

    def _run(self) -> None:
        import cv2

        period = 1.0 / self.hz
        while not self._stop.wait(period):
            with self._lock:
                snapshot = list(self._latest.items())
            now = time.monotonic()
            for ident, (blob, w, h) in snapshot:
                try:
                    pred = self._infer(cv2, blob, w, h)
                except Exception as exc:  # noqa: BLE001
                    log.warning("hand-raise infer %s: %s", ident, exc)
                    pred = False
                with self._lock:
                    latch = self._latch.get(ident)
                    if latch is None:
                        latch = RaiseLatch(hits=3, hold_s=self.hold_s, window=5)
                        self._latch[ident] = latch
                    was = latch.raised
                    now_raised = latch.update(pred, now)
                    if now_raised:
                        self._raised.add(ident)
                    else:
                        self._raised.discard(ident)
                if now_raised != was:
                    log.info("hand-raise %s %s", ident, "on" if now_raised else "off")
                    if self.on_change is not None:
                        try:
                            self.on_change(ident, now_raised)
                        except Exception:  # noqa: BLE001
                            log.exception("hand-raise on_change(%s, %s)", ident, now_raised)

    def infer_rgb(self, rgb: np.ndarray) -> np.ndarray:
        """MoveNet (17, 3) y,x,score from an HxWx3 RGB uint8 image."""
        if self._request is None:
            return np.zeros((17, 3), dtype=np.float32)
        small = letterbox_rgb(rgb, self._input_size)
        inp = np.expand_dims(np.ascontiguousarray(small, dtype=np.int32), 0)
        result = self._request.infer({self._ov_input: inp})
        return np.asarray(next(iter(result.values())), dtype=np.float32).reshape(17, 3)

    def _infer(self, cv2, blob: bytes, w: int, h: int) -> bool:
        need = w * h * 3 // 2
        if self._request is None or len(blob) < need:
            return False
        yuv = np.frombuffer(blob, dtype=np.uint8, count=need).reshape((h * 3 // 2, w))
        rgb = cv2.cvtColor(yuv, cv2.COLOR_YUV2RGB_I420)
        return wrist_is_raised(self.infer_rgb(rgb))