"""Resident YOLO11s + RTMW-L pose stage extracted from the v9.9 pipeline.""" from __future__ import annotations import argparse import os import time from collections import defaultdict from pathlib import Path import numpy as np SIDES = ("left", "right") class LivePoseEngine: """Run the v9.9 YOLO/RTMW pose and temporal-fusion path without EdgeTAM.""" def __init__(self, args: argparse.Namespace): self._trt_dll_directory = None if args.pose_provider == "tensorrt" and os.name == "nt": try: import tensorrt_libs except ImportError as exc: raise RuntimeError( "--pose-provider tensorrt on Windows requires the " "tensorrt-cu12-libs package in the active Python environment." ) from exc trt_lib_dir = Path(tensorrt_libs.__file__).resolve().parent self._trt_dll_directory = os.add_dll_directory(str(trt_lib_dir)) os.environ["PATH"] = ( str(trt_lib_dir) + os.pathsep + os.environ.get("PATH", "") ) if args.pose_device.lower().startswith("cuda"): import onnxruntime as ort if "CUDAExecutionProvider" not in ort.get_available_providers(): raise RuntimeError( "--pose-device cuda requires onnxruntime-gpu. Uninstall " "onnxruntime and install onnxruntime-gpu==1.22.0." ) from rtmlib import RTMPose from ultralytics import YOLO from predict_optimized_pose_videos import BoxTracker self.args = args self.detector = YOLO(str(args.det_model)) self.pose = RTMPose( str(args.pose_model), model_input_size=(288, 384), backend="onnxruntime", device=args.pose_device, ) if args.pose_provider == "tensorrt": import onnxruntime as ort available = ort.get_available_providers() if "TensorrtExecutionProvider" not in available: raise RuntimeError( "--pose-provider tensorrt requires ONNX Runtime with " "TensorrtExecutionProvider. " f"Available providers: {available}" ) cache_dir = args.pose_trt_cache.expanduser().resolve() cache_dir.mkdir(parents=True, exist_ok=True) device_id = ( int(str(args.pose_device).split(":")[-1]) if ":" in str(args.pose_device) else 0 ) trt_options = { "device_id": device_id, "trt_fp16_enable": "True", "trt_engine_cache_enable": "True", "trt_engine_cache_path": str(cache_dir), "trt_timing_cache_enable": "True", "trt_timing_cache_path": str(cache_dir), } self.pose.session = ort.InferenceSession( str(args.pose_model.resolve()), providers=[ ("TensorrtExecutionProvider", trt_options), ("CUDAExecutionProvider", {"device_id": device_id}), ], ) active_providers = self.pose.session.get_providers() if not active_providers or active_providers[0] != "TensorrtExecutionProvider": raise RuntimeError( "RTMW TensorRT initialization failed instead of using the " "requested provider. Active providers: " f"{active_providers}" ) print(f"RTMW ONNX Runtime providers: {active_providers}") self.box_tracker = BoxTracker(max_age=30) self.people = {} self.timings = defaultdict(list) self.isolated_inferences = 0 def infer(self, frame: np.ndarray, frame_idx: int) -> dict: from compare_rtmw_ml import HAND_SLICES from fuse_rtmw_l_temporal import adaptive_smooth, build_global_hands, fuse_person from occlusion_aware_pose_v2 import Measurement, resolve_global_hand_assignments from predict_optimized_pose_videos import ( LivePersonState, filter_person_boxes, mask_other_people, update_live_v2, ) infer_started = time.perf_counter() height, width = frame.shape[:2] started = time.perf_counter() detection_height = max( 1, min(height, round(height * self.args.det_roi_y_max_ratio)) ) detection = self.detector.predict( frame[:detection_height], imgsz=self.args.det_imgsz, conf=self.args.det_conf, iou=0.30, classes=[0], device=self.args.det_device, verbose=False, )[0] self.timings["detector_ms"].append((time.perf_counter() - started) * 1000) boxes, confidences, _ = filter_person_boxes( detection, width, height, self.args.max_people ) track_ids = self.box_tracker.update(boxes, frame_idx, (height, width)) if boxes: started = time.perf_counter() raw_keypoints, raw_scores = self.pose(frame, bboxes=boxes) self.timings["raw_pose_ms"].append((time.perf_counter() - started) * 1000) isolated_keypoints = raw_keypoints.copy() isolated_scores = raw_scores.copy() started = time.perf_counter() for person_idx, bbox in enumerate(boxes): target_area = max((bbox[2] - bbox[0]) * (bbox[3] - bbox[1]), 1.0) needs_isolation = any( idx != person_idx and max((other[2] - other[0]) * (other[3] - other[1]), 1.0) < 0.5 * target_area for idx, other in enumerate(boxes) ) if not needs_isolation: continue masked = mask_other_people(frame, boxes, person_idx) keypoints, scores = self.pose(masked, bboxes=[bbox]) isolated_keypoints[person_idx] = keypoints[0] isolated_scores[person_idx] = scores[0] self.isolated_inferences += 1 self.timings["isolated_extra_ms"].append( (time.perf_counter() - started) * 1000 ) else: raw_keypoints = np.empty((0, 133, 2), dtype=np.float64) raw_scores = np.empty((0, 133), dtype=np.float64) isolated_keypoints = raw_keypoints.copy() isolated_scores = raw_scores.copy() selected_list, smoothed_list, metadata_list, states = [], [], [], [] for person_idx, (bbox, track_id) in enumerate(zip(boxes, track_ids)): state = self.people.setdefault(track_id, LivePersonState()) selected, metadata = fuse_person( raw_keypoints[person_idx], raw_scores[person_idx], isolated_keypoints[person_idx], isolated_scores[person_idx], bbox, state.fusion, ) selected_list.append(selected) smoothed_list.append(adaptive_smooth(selected, bbox, state.fusion)) metadata_list.append(metadata) states.append(state) measurements, endpoint_tracks = {}, {} for person_idx, (selected, metadata, state) in enumerate( zip(selected_list, metadata_list, states) ): for side, (start, end) in HAND_SLICES.items(): meta = metadata[side] measurements[(person_idx, side)] = ( Measurement( keypoints=selected[start:end], precise=bool(meta["distance_eligible"]), provisional=not bool(meta["distance_eligible"]), quality=float(meta.get("selected_hand_score", 0.0)), ) if meta["alarm_eligible"] else None ) endpoint_tracks[(person_idx, side)] = state.endpoints[side] measurements, assignments = resolve_global_hand_assignments( measurements, endpoint_tracks, boxes, 30 ) v2_list, v2_meta_list = [], [] for person_idx, (bbox, selected, smoothed, metadata, state) in enumerate( zip(boxes, selected_list, smoothed_list, metadata_list, states) ): v2, v2_meta = update_live_v2( frame_idx, bbox, selected, smoothed, metadata, state, {side: measurements[(person_idx, side)] for side in SIDES}, {side: assignments.get((person_idx, side)) for side in SIDES}, 12, 30, ) v2_list.append(v2) v2_meta_list.append(v2_meta) self.timings["pose_pipeline_total_ms"].append( (time.perf_counter() - infer_started) * 1000 ) return { "frame_index": frame_idx, "track_ids": track_ids, "boxes": np.round(np.asarray(boxes), 3).tolist() if boxes else [], "det_confidences": [round(value, 4) for value in confidences], "raw_keypoints": np.round(raw_keypoints, 3).tolist(), "temporal_selected_keypoints": np.round( np.asarray(selected_list), 3 ).tolist(), "temporal_keypoints": np.round(np.asarray(smoothed_list), 3).tolist(), "temporal_metadata": metadata_list, "v2_keypoints": np.round(np.asarray(v2_list), 3).tolist(), "v2_metadata": v2_meta_list, "global_hands": build_global_hands(selected_list, metadata_list), }