kai/infra-docker-configs
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add RTMLib CUDA backends and staged timing

Browse Source
This commit is contained in:
Codex
2026-08-15 10:58:02 -07:00
parent dba60667cb
commit 174ba875dd
8 changed files with 243 additions and 61 deletions
Download Patch File Download Diff File Expand all files Collapse all files
servers/unraid/tail2-pose-server/app/tail2/main.py
+31 -12
View File
@@ -75,11 +75,13 @@ def gpu_info() -> dict:
handle = nvmlDeviceGetHandleByIndex(physical_index); memory = nvmlDeviceGetMemoryInfo(handle)
name = nvmlDeviceGetName(handle)
if isinstance(name, bytes): name = name.decode()
return {"name": name, "cuda_available": torch.cuda.is_available(),
return {"name": name, "cuda_available": torch.cuda.is_available(), "physical_index": physical_index,
"memory_used_mb": round(memory.used / 1048576), "memory_total_mb": round(memory.total / 1048576),
"process_reserved_mb": round(torch.cuda.memory_reserved() / 1048576) if torch.cuda.is_available() else None,
"temperature_c": nvmlDeviceGetTemperature(handle, NVML_TEMPERATURE_GPU)}
except Exception:
return {"name": None, "cuda_available": torch.cuda.is_available(), "memory_used_mb": None,
return {"name": None, "cuda_available": torch.cuda.is_available(), "physical_index": None, "memory_used_mb": None,
"process_reserved_mb": None,
"memory_total_mb": None, "temperature_c": None}
@@ -112,15 +114,22 @@ def capabilities():
"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,
"model_backend": models.backend, "execution_providers": models.providers,
"runtime_inputs": models.input_metadata,
"detector": {"name": os.getenv("DETECTOR_NAME", "rtmdet-m-person"), "framework": "mmdetection" if models.backend == "openmmlab" else "rtmlib/onnxruntime", "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": os.getenv("POSE_NAME", "rtmw-l-256x192"), "framework": "mmpose", "keypoint_schema": "coco_wholebody_133",
"pose": {"name": os.getenv("POSE_NAME", "rtmw-l-256x192"), "framework": "mmpose" if models.backend == "openmmlab" else "rtmlib/onnxruntime", "keypoint_schema": "coco17" if models.backend == "rtmlib_body" else "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}}
"limits": {"max_people": int(os.getenv("MAX_PEOPLE", "4")), "latest_frame_only": True, "frame_queue_size": 1},
"timing_fields": ["decode_ms", "frame_copy_ms", "color_conversion_ms", "resize_letterbox_ms",
"cpu_to_gpu_ms", "detector_forward_ms", "detector_postprocess_ms",
"bbox_crop_affine_ms", "pose_forward_ms", "pose_decode_ms", "gpu_to_cpu_ms",
"json_serialization_ms", "websocket_enqueue_ms", "detector_ms", "pose_ms",
"total_ms", "total_inference_ms", "source_to_sent_ms"]}
@app.get("/metrics", response_class=PlainTextResponse, dependencies=[Depends(authorized)])
@@ -139,7 +148,7 @@ async def poses(ws: WebSocket):
await ws.close(code=4401); return
await ws.accept(); q = hub.add()
try:
while True: await ws.send_json(await q.get())
while True: await ws.send_text(await q.get())
except WebSocketDisconnect: pass
finally: hub.remove(q)
@@ -154,9 +163,9 @@ def inference_worker(loop: asyncio.AbstractEventLoop) -> None:
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()
frame, dropped = 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)
people, stage = 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
@@ -164,14 +173,24 @@ def inference_worker(loop: asyncio.AbstractEventLoop) -> None:
"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,
"dropped_since_last": dropped,
"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},
"people": people, "timing": {"decode_ms": frame.decode_ms, **stage,
"json_serialization_ms": 0.0, "websocket_enqueue_ms": 0.0,
"total_ms": 0.0, "total_inference_ms": total_ms,
"source_to_sent_ms": 0.0},
"error": None}
serialization_start = time.perf_counter_ns(); json.dumps(msg, separators=(",", ":"), allow_nan=False)
msg["timing"]["json_serialization_ms"] = (time.perf_counter_ns() - serialization_start) / 1e6
enqueue_start = time.perf_counter_ns()
msg["timing"]["websocket_enqueue_ms"] = (time.perf_counter_ns() - enqueue_start) / 1e6
msg["timing"]["total_ms"] = (time.perf_counter_ns() - perf) / 1e6
msg["server_sent_unix_ns"] = time.time_ns()
msg["timing"]["source_to_sent_ms"] = (msg["server_sent_unix_ns"] - frame.received_ns) / 1e6
encoded = json.dumps(msg, separators=(",", ":"), allow_nan=False)
state.inference_ms.append(total_ms); state.messages += 1; state.ws_source_ok = True
hub.publish_on_loop(loop, msg)
hub.publish_on_loop(loop, encoded)
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")