241 lines
9.7 KiB
Python
241 lines
9.7 KiB
Python
"""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),
|
|
}
|