189 lines
9.7 KiB
Python
189 lines
9.7 KiB
Python
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": os.getenv("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)
|