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 from .offline_benchmark import run_offline_benchmark 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.error_detail: 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, "error_detail": state.error_detail, "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() run_offline_benchmark(models, threshold, max_people) state.models_ready = True; LOG.info("official detector and pose checkpoints loaded") except Exception as exc: state.inference_fault = True; state.error_code = "MODEL_LOAD_FAILED" state.error_detail = f"{type(exc).__name__}: {str(exc)[:240]}"; 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; state.error_detail = 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 as exc: state.inference_fault = True; state.error_code = "INFERENCE_FAILED" state.error_detail = f"{type(exc).__name__}: {str(exc)[:240]}"; 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)