149 lines
6.3 KiB
Python
149 lines
6.3 KiB
Python
"""Resident native TensorRT inference with causal pose fusion and tracking."""
|
|
from __future__ import annotations
|
|
import time
|
|
from collections import defaultdict, deque
|
|
import numpy as np
|
|
from pose_tracking import BoxTracker
|
|
|
|
SIDES = ("left", "right")
|
|
|
|
class NativePoseEngine:
|
|
def __init__(self, args, native):
|
|
self.args = args
|
|
self.native = native
|
|
self.pose = native.pose
|
|
self.box_tracker = BoxTracker(max_age=30)
|
|
self.people = {}
|
|
self.timings = defaultdict(lambda: deque(maxlen=3000))
|
|
self.isolated_inferences = 0
|
|
|
|
def infer(self, frame: np.ndarray, frame_idx: int) -> dict:
|
|
from pose_geometry import HAND_SLICES
|
|
from pose_fusion import adaptive_smooth, build_global_hands, fuse_person
|
|
from pose_occlusion import Measurement, resolve_global_hand_assignments
|
|
from pose_tracking import (
|
|
LivePersonState,
|
|
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))
|
|
)
|
|
boxes, confidences = self.native.detect(frame[:detection_height])
|
|
self.timings["detector_ms"].append((time.perf_counter() - started) * 1000)
|
|
track_ids = self.box_tracker.update(boxes, frame_idx, (height, width))
|
|
|
|
# Retire disappeared people in long-running camera sessions.
|
|
for stale_id in list(self.people):
|
|
if stale_id not in self.box_tracker.tracks:
|
|
del self.people[stale_id]
|
|
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),
|
|
}
|