feat(unraid): add Tail 2 GPU pose service

2026-08-15 09:29:01 -07:00
parent ee1df2cad5
commit 67f2adbe5b
17 changed files with 659 additions and 0 deletions
@@ -0,0 +1,2 @@
+__version__ = "0.1.0"
+
@@ -0,0 +1,91 @@
+from __future__ import annotations
+
+import asyncio
+import math
+import queue
+import threading
+from dataclasses import dataclass
+from typing import Any
+
+import numpy as np
+
+
+def finite01(value: float) -> float:
+    value = float(value)
+    return min(1.0, max(0.0, value)) if math.isfinite(value) else 0.0
+
+
+def normalize_bbox(bbox: list[float] | np.ndarray, width: int, height: int) -> list[float]:
+    return [finite01(bbox[0] / width), finite01(bbox[1] / height),
+            finite01(bbox[2] / width), finite01(bbox[3] / height)]
+
+
+def normalize_keypoints(points: np.ndarray, width: int, height: int, count: int) -> list[list[float]]:
+    out = [[finite01(p[0] / width), finite01(p[1] / height), finite01(p[2])] for p in points[:count]]
+    out.extend([[0.0, 0.0, 0.0] for _ in range(count - len(out))])
+    return out
+
+
+def inverse_letterbox(points: np.ndarray, scale: float, pad_x: float, pad_y: float) -> np.ndarray:
+    out = points.copy().astype(float)
+    out[..., 0] = (out[..., 0] - pad_x) / scale
+    out[..., 1] = (out[..., 1] - pad_y) / scale
+    return out
+
+
+class LatestQueue:
+    def __init__(self) -> None:
+        self._q: queue.Queue[Any] = queue.Queue(maxsize=1)
+
+    def put(self, item: Any) -> None:
+        try:
+            self._q.put_nowait(item)
+        except queue.Full:
+            try:
+                self._q.get_nowait()
+            except queue.Empty:
+                pass
+            self._q.put_nowait(item)
+
+    def get(self, timeout: float | None = None) -> Any:
+        return self._q.get(timeout=timeout)
+
+    def qsize(self) -> int:
+        return self._q.qsize()
+
+
+class BroadcastHub:
+    def __init__(self) -> None:
+        self._clients: set[asyncio.Queue] = set()
+        self._lock = threading.Lock()
+
+    def add(self) -> asyncio.Queue:
+        q: asyncio.Queue = asyncio.Queue(maxsize=1)
+        with self._lock:
+            self._clients.add(q)
+        return q
+
+    def remove(self, q: asyncio.Queue) -> None:
+        with self._lock:
+            self._clients.discard(q)
+
+    def publish_on_loop(self, loop: asyncio.AbstractEventLoop, message: dict) -> None:
+        def publish() -> None:
+            with self._lock:
+                clients = list(self._clients)
+            for q in clients:
+                if q.full():
+                    try:
+                        q.get_nowait()
+                    except asyncio.QueueEmpty:
+                        pass
+                q.put_nowait(message)
+        loop.call_soon_threadsafe(publish)
+
+
+@dataclass
+class Frame:
+    pixels: np.ndarray
+    received_ns: int
+    decode_ms: float
+
@@ -0,0 +1,181 @@
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+import os
+import statistics
+import threading
+import time
+import uuid
+from collections import deque
+
+import torch
+import uvicorn
+from fastapi import Depends, FastAPI, Header, HTTPException, WebSocket, WebSocketDisconnect
+from fastapi.responses import JSONResponse, PlainTextResponse
+from pynvml import (nvmlDeviceGetHandleByIndex, nvmlDeviceGetMemoryInfo,
+                    nvmlDeviceGetName, nvmlDeviceGetTemperature, nvmlInit,
+                    NVML_TEMPERATURE_GPU)
+
+from .core import BroadcastHub, LatestQueue
+from .models import Models
+from .ndi import NDIReceiver
+
+logging.basicConfig(level=os.getenv("LOG_LEVEL", "INFO"), format="%(asctime)s %(levelname)s %(name)s %(message)s")
+LOG = logging.getLogger("tail2")
+BODY17 = ["nose", "left_eye", "right_eye", "left_ear", "right_ear", "left_shoulder",
+          "right_shoulder", "left_elbow", "right_elbow", "left_wrist", "right_wrist",
+          "left_hip", "right_hip", "left_knee", "right_knee", "left_ankle", "right_ankle"]
+
+
+class State:
+    def __init__(self) -> None:
+        self.started = time.monotonic()
+        self.session_id = str(uuid.uuid4())
+        self.last_frame_ns = 0
+        self.width = self.height = 0
+        self.ndi_connected = False
+        self.models_ready = False
+        self.cuda_inference_ok = False
+        self.inference_fault = False
+        self.ws_source_ok = False
+        self.error_code: str | None = None
+        self.frame_id = 0
+        self.messages = 0
+        self.frames_dropped = 0
+        self.inference_ms: deque[float] = deque(maxlen=1000)
+        self.lock = threading.Lock()
+
+    def note_frame(self, now_ns: int, width: int, height: int) -> None:
+        with self.lock:
+            self.last_frame_ns, self.width, self.height = now_ns, width, height
+            self.ndi_connected = True
+
+
+state = State()
+frames = LatestQueue()
+hub = BroadcastHub()
+models = Models()
+token = open(os.environ["API_TOKEN_FILE"], encoding="utf-8").read().strip()
+config = json.load(open("/config/service.json", encoding="utf-8"))
+app = FastAPI(docs_url=None, redoc_url=None, openapi_url=None)
+
+
+def authorized(authorization: str | None = Header(default=None)) -> None:
+    if not authorization or not authorization.startswith("Bearer ") or not __import__("hmac").compare_digest(authorization[7:], token):
+        raise HTTPException(status_code=401, detail="unauthorized")
+
+
+def gpu_info() -> dict:
+    try:
+        nvmlInit(); handle = nvmlDeviceGetHandleByIndex(0); memory = nvmlDeviceGetMemoryInfo(handle)
+        name = nvmlDeviceGetName(handle)
+        if isinstance(name, bytes): name = name.decode()
+        return {"name": name, "cuda_available": torch.cuda.is_available(),
+                "memory_used_mb": round(memory.used / 1048576), "memory_total_mb": round(memory.total / 1048576),
+                "temperature_c": nvmlDeviceGetTemperature(handle, NVML_TEMPERATURE_GPU)}
+    except Exception:
+        return {"name": None, "cuda_available": torch.cuda.is_available(), "memory_used_mb": None,
+                "memory_total_mb": None, "temperature_c": None}
+
+
+def readiness() -> tuple[bool, float | None]:
+    age = (time.time_ns() - state.last_frame_ns) / 1e6 if state.last_frame_ns else None
+    ready = bool(state.ndi_connected and age is not None and age < config["stale_frame_ms"] and
+                 state.models_ready and state.cuda_inference_ok and not state.inference_fault and state.ws_source_ok)
+    return ready, age
+
+
+@app.get("/v1/health", dependencies=[Depends(authorized)])
+def health():
+    ready, age = readiness()
+    return {"schema": "tail2.health.v1", "status": "ok" if ready else "degraded", "ready": ready,
+            "server_time_unix_ns": time.time_ns(), "session_id": state.session_id, "gpu": gpu_info(),
+            "ndi": {"connected": bool(state.ndi_connected and age is not None and age < config["stale_frame_ms"]),
+                    "source": os.environ["NDI_SOURCE_NAME"], "source_ip": os.environ["NDI_SOURCE_IP"],
+                    "last_frame_age_ms": round(age, 3) if age is not None else None,
+                    "actual_width": state.width or None, "actual_height": state.height or None},
+            "models": {"detector_ready": state.models_ready, "pose_ready": state.models_ready},
+            "error_code": state.error_code, "uptime_seconds": round(time.monotonic() - state.started, 1)}
+
+
+@app.get("/v1/capabilities", dependencies=[Depends(authorized)])
+def capabilities():
+    return {"schema": "tail2.capabilities.v1", "protocol_version": 1, "stream_id": config["stream_id"],
+            "session_id": state.session_id,
+            "input": {"transport": "ndi", "source": os.environ["NDI_SOURCE_NAME"],
+                      "source_ip": os.environ["NDI_SOURCE_IP"], "nominal_width": config["nominal_width"],
+                      "nominal_height": config["nominal_height"], "coordinate_space": config["coordinate_space"],
+                      "horizontal_flip_applied": False, "rotation_deg": 0},
+            "detector": {"name": "rtmdet-m-person", "framework": "mmdetection", "config": models.det_config,
+                         "checkpoint": models.det_checkpoint, "checkpoint_sha256": models.sha256(models.det_checkpoint),
+                         "person_class_id": 0, "score_threshold": float(os.getenv("DETECTION_THRESHOLD", ".35")),
+                         "nms_iou_threshold": float(os.getenv("NMS_IOU_THRESHOLD", ".60"))},
+            "pose": {"name": "rtmw-l-256x192", "framework": "mmpose", "keypoint_schema": "coco_wholebody_133",
+                     "keypoint_order": "MMPose COCO WholeBody 133 official order", "body_schema": "coco17",
+                     "body17_names": BODY17, "config": models.pose_config, "checkpoint": models.pose_checkpoint,
+                     "checkpoint_sha256": models.sha256(models.pose_checkpoint)},
+            "limits": {"max_people": int(os.getenv("MAX_PEOPLE", "4")), "latest_frame_only": True, "frame_queue_size": 1}}
+
+
+@app.get("/metrics", response_class=PlainTextResponse, dependencies=[Depends(authorized)])
+def metrics():
+    ready, age = readiness(); vals = list(state.inference_ms)
+    p95 = sorted(vals)[max(0, int(len(vals) * .95) - 1)] if vals else 0
+    return (f"tail2_ready {int(ready)}\ntail2_ndi_connected {int(state.ndi_connected)}\n"
+            f"tail2_last_frame_age_ms {age or 0:.3f}\ntail2_messages_total {state.messages}\n"
+            f"tail2_inference_ms_p95 {p95:.3f}\ntail2_frame_queue_size {frames.qsize()}\n")
+
+
+@app.websocket("/v1/poses")
+async def poses(ws: WebSocket):
+    auth = ws.headers.get("authorization", "")
+    if not auth.startswith("Bearer ") or not __import__("hmac").compare_digest(auth[7:], token):
+        await ws.close(code=4401); return
+    await ws.accept(); q = hub.add()
+    try:
+        while True: await ws.send_json(await q.get())
+    except WebSocketDisconnect: pass
+    finally: hub.remove(q)
+
+
+def inference_worker(loop: asyncio.AbstractEventLoop) -> None:
+    threshold = float(os.getenv("DETECTION_THRESHOLD", ".35")); max_people = int(os.getenv("MAX_PEOPLE", "4"))
+    try:
+        models.load(); state.models_ready = True; LOG.info("official detector and pose checkpoints loaded")
+    except Exception:
+        state.inference_fault = True; state.error_code = "MODEL_LOAD_FAILED"; LOG.exception("model loading failed"); return
+    while True:
+        frame = frames.get(); started_ns = time.time_ns(); perf = time.perf_counter_ns()
+        try:
+            people, det_ms, pose_ms = models.infer(frame.pixels, threshold, max_people)
+            finished_ns = time.time_ns(); total_ms = (time.perf_counter_ns() - perf) / 1e6
+            state.frame_id += 1; state.cuda_inference_ok = True; state.inference_fault = False; state.error_code = None
+            msg = {"schema": "tail2.pose.v1", "stream_id": config["stream_id"], "session_id": state.session_id,
+                   "frame_id": state.frame_id, "source_received_unix_ns": frame.received_ns,
+                   "inference_started_unix_ns": started_ns, "inference_finished_unix_ns": finished_ns,
+                   "server_sent_unix_ns": time.time_ns(), "source_alive": True,
+                   "frame": {"width": frame.pixels.shape[1], "height": frame.pixels.shape[0], "rotation_deg": 0,
+                             "coordinate_space": "ndi_buffer_unmodified", "horizontal_flip_applied": False},
+                   "people": people, "timing": {"decode_ms": frame.decode_ms, "detector_ms": det_ms,
+                                                 "pose_ms": pose_ms, "total_inference_ms": total_ms,
+                                                 "source_to_sent_ms": (time.time_ns() - frame.received_ns) / 1e6},
+                   "error": None}
+            state.inference_ms.append(total_ms); state.messages += 1; state.ws_source_ok = True
+            hub.publish_on_loop(loop, msg)
+        except Exception:
+            state.inference_fault = True; state.error_code = "INFERENCE_FAILED"; LOG.exception("inference failed")
+
+
+@app.on_event("startup")
+async def startup() -> None:
+    loop = asyncio.get_running_loop()
+    threading.Thread(target=inference_worker, args=(loop,), name="inference", daemon=True).start()
+    NDIReceiver(frames, state).start()
+
+
+if __name__ == "__main__":
+    uvicorn.run(app, host=os.getenv("API_BIND", "0.0.0.0"), port=int(os.getenv("API_PORT", "18120")),
+                access_log=False, ws_max_size=65536)
+
@@ -0,0 +1,69 @@
+from __future__ import annotations
+
+import hashlib
+import os
+import time
+
+import numpy as np
+import torch
+
+from .core import normalize_bbox, normalize_keypoints
+
+
+class Models:
+    def __init__(self) -> None:
+        self.det_config = os.environ["DET_CONFIG"]
+        self.det_checkpoint = os.environ["DET_CHECKPOINT"]
+        self.pose_config = os.environ["POSE_CONFIG"]
+        self.pose_checkpoint = os.environ["POSE_CHECKPOINT"]
+        self.detector = self.pose = None
+
+    @staticmethod
+    def sha256(path: str) -> str:
+        digest = hashlib.sha256()
+        with open(path, "rb") as fh:
+            for chunk in iter(lambda: fh.read(8 * 1024 * 1024), b""):
+                digest.update(chunk)
+        return digest.hexdigest()
+
+    def load(self) -> None:
+        from mmdet.apis import init_detector
+        from mmpose.apis import init_model
+        self.detector = init_detector(self.det_config, self.det_checkpoint, device="cuda:0")
+        self.pose = init_model(self.pose_config, self.pose_checkpoint, device="cuda:0")
+
+    @torch.inference_mode()
+    def infer(self, frame: np.ndarray, threshold: float, max_people: int) -> tuple[list[dict], float, float]:
+        from mmdet.apis import inference_detector
+        from mmpose.apis import inference_topdown
+        height, width = frame.shape[:2]
+        t0 = time.perf_counter_ns()
+        result = inference_detector(self.detector, frame)
+        pred = result.pred_instances.cpu().numpy()
+        keep = (pred.labels == 0) & (pred.scores >= threshold)
+        boxes = pred.bboxes[keep]
+        scores = pred.scores[keep]
+        if len(boxes):
+            order = np.argsort(scores)[::-1][:max_people]
+            boxes, scores = boxes[order], scores[order]
+        det_ms = (time.perf_counter_ns() - t0) / 1e6
+        t1 = time.perf_counter_ns()
+        pose_results = inference_topdown(self.pose, frame, bboxes=boxes) if len(boxes) else []
+        pose_ms = (time.perf_counter_ns() - t1) / 1e6
+        people = []
+        for idx, (box, det_score, sample) in enumerate(zip(boxes, scores, pose_results)):
+            points = sample.pred_instances.keypoints[0]
+            kp_scores = sample.pred_instances.keypoint_scores[0]
+            combined = np.column_stack((points, kp_scores))
+            wholebody = normalize_keypoints(combined, width, height, 133)
+            people.append({
+                "detection_id": idx,
+                "bbox_xyxy_norm": normalize_bbox(box, width, height),
+                "det_score": float(det_score),
+                "keypoint_schema": "coco_wholebody_133",
+                "keypoints_norm": wholebody,
+                "body17": wholebody[:17],
+                "pose_score": float(np.mean(np.nan_to_num(kp_scores[:17]))),
+            })
+        return people, det_ms, pose_ms
+
@@ -0,0 +1,97 @@
+from __future__ import annotations
+
+import logging
+import os
+import socket
+import threading
+import time
+import uuid
+
+import cv2
+
+from .core import Frame, LatestQueue
+
+LOG = logging.getLogger("tail2.ndi")
+
+
+class NDIReceiver(threading.Thread):
+    """Receive only. This class intentionally exposes no NDI PTZ methods."""
+
+    daemon = True
+
+    def __init__(self, frames: LatestQueue, state) -> None:
+        super().__init__(name="ndi-receiver")
+        self.frames, self.state = frames, state
+        self.source_name = os.environ["NDI_SOURCE_NAME"]
+        self.source_ip = os.environ["NDI_SOURCE_IP"]
+        self.stop_event = threading.Event()
+
+    def stop(self) -> None:
+        self.stop_event.set()
+
+    def run(self) -> None:
+        while not self.stop_event.is_set():
+            try:
+                self._receive_session()
+            except Exception as exc:
+                self.state.ndi_connected = False
+                self.state.error_code = "NDI_DISCONNECTED"
+                LOG.warning("NDI receive session ended: %s", type(exc).__name__)
+            self.stop_event.wait(2.0)
+
+    def _receive_session(self) -> None:
+        from cyndilib.finder import Finder
+        from cyndilib.receiver import Receiver
+        from cyndilib.video_frame import VideoFrameSync
+        from cyndilib.wrapper.ndi_recv import RecvBandwidth, RecvColorFormat
+
+        # NDI SDK reads this file before finder initialization. It supplies the
+        # camera as an additional unicast discovery server without changing it.
+        ndi_dir = "/tmp/ndi"
+        os.makedirs(ndi_dir, exist_ok=True)
+        with open(f"{ndi_dir}/ndi-config.v1.json", "w", encoding="utf-8") as fh:
+            fh.write('{"ndi":{"networks":{"ips":"' + self.source_ip + '"}}}')
+        os.environ["NDI_CONFIG_DIR"] = ndi_dir
+
+        finder = Finder()
+        finder.open()
+        receiver = None
+        try:
+            deadline = time.monotonic() + 30
+            source = None
+            while not self.stop_event.is_set() and time.monotonic() < deadline:
+                finder.wait_for_sources(1)
+                names = finder.get_source_names()
+                if self.source_name in names:
+                    source = finder.get_source(self.source_name)
+                    break
+            if source is None:
+                raise RuntimeError("configured NDI source not discovered")
+            receiver = Receiver(color_format=RecvColorFormat.BGRX_BGRA,
+                                bandwidth=RecvBandwidth.highest)
+            video = VideoFrameSync()
+            receiver.frame_sync.set_video_frame(video)
+            receiver.set_source(source)
+            deadline = time.monotonic() + 15
+            while not receiver.is_connected() and time.monotonic() < deadline:
+                time.sleep(0.2)
+            if not receiver.is_connected():
+                raise RuntimeError("NDI source connection timeout")
+
+            self.state.session_id = str(uuid.uuid4())
+            while not self.stop_event.is_set() and receiver.is_connected():
+                started = time.perf_counter_ns()
+                receiver.frame_sync.capture_video()
+                if min(video.xres, video.yres) <= 0:
+                    time.sleep(0.005)
+                    continue
+                raw = video.get_array().reshape(video.yres, video.xres, 4)
+                pixels = cv2.cvtColor(raw, cv2.COLOR_BGRA2BGR)
+                now = time.time_ns()
+                self.frames.put(Frame(pixels.copy(), now, (time.perf_counter_ns() - started) / 1e6))
+                self.state.note_frame(now, video.xres, video.yres)
+        finally:
+            if receiver is not None and receiver.is_connected():
+                receiver.disconnect()
+            finder.close()
+