305 lines
11 KiB
Python
305 lines
11 KiB
Python
"""Runtime pose tracking helpers."""
|
|
from __future__ import annotations
|
|
import math
|
|
from collections import deque
|
|
from dataclasses import dataclass, field
|
|
import numpy as np
|
|
from pose_geometry import HAND_SLICES
|
|
from pose_fusion import BODY_SIDE_INDICES, PersonTrack
|
|
from pose_occlusion import EndpointTrack, Measurement, calibrated_bone_lengths, reachable_region, solve_two_bone_elbow, temporal_recovery_measurement, update_endpoint_track
|
|
|
|
|
|
def bbox_iou(a: np.ndarray, b: np.ndarray) -> float:
|
|
x1 = max(float(a[0]), float(b[0]))
|
|
y1 = max(float(a[1]), float(b[1]))
|
|
x2 = min(float(a[2]), float(b[2]))
|
|
y2 = min(float(a[3]), float(b[3]))
|
|
intersection = max(x2 - x1, 0.0) * max(y2 - y1, 0.0)
|
|
union = (
|
|
max(float(a[2] - a[0]), 0.0) * max(float(a[3] - a[1]), 0.0)
|
|
+ max(float(b[2] - b[0]), 0.0) * max(float(b[3] - b[1]), 0.0)
|
|
- intersection
|
|
)
|
|
return intersection / max(union, 1.0)
|
|
|
|
|
|
class BoxTracker:
|
|
def __init__(self, max_age: int = 30):
|
|
self.max_age = max_age
|
|
self.next_id = 1
|
|
self.tracks: dict[int, dict] = {}
|
|
|
|
def update(
|
|
self,
|
|
boxes: list[list[float]],
|
|
frame_idx: int,
|
|
image_shape: tuple[int, int],
|
|
) -> list[int]:
|
|
height, width = image_shape
|
|
image_diagonal = math.hypot(width, height)
|
|
active_ids = [
|
|
track_id
|
|
for track_id, track in self.tracks.items()
|
|
if frame_idx - track["last_frame"] <= self.max_age
|
|
]
|
|
self.tracks = {track_id: self.tracks[track_id] for track_id in active_ids}
|
|
pairs = []
|
|
for detection_idx, box in enumerate(boxes):
|
|
box_array = np.asarray(box, dtype=np.float64)
|
|
center = (box_array[:2] + box_array[2:]) * 0.5
|
|
for track_id in active_ids:
|
|
prior = np.asarray(self.tracks[track_id]["box"], dtype=np.float64)
|
|
prior_center = (prior[:2] + prior[2:]) * 0.5
|
|
iou = bbox_iou(box_array, prior)
|
|
center_distance = float(np.linalg.norm(center - prior_center)) / image_diagonal
|
|
if iou >= 0.05 or center_distance <= 0.18:
|
|
score = 2.0 * iou + max(0.0, 1.0 - center_distance / 0.18)
|
|
pairs.append((score, detection_idx, track_id))
|
|
|
|
assigned_detections = set()
|
|
assigned_tracks = set()
|
|
result: list[int | None] = [None] * len(boxes)
|
|
for _, detection_idx, track_id in sorted(pairs, reverse=True):
|
|
if detection_idx in assigned_detections or track_id in assigned_tracks:
|
|
continue
|
|
assigned_detections.add(detection_idx)
|
|
assigned_tracks.add(track_id)
|
|
result[detection_idx] = track_id
|
|
|
|
for detection_idx, box in enumerate(boxes):
|
|
track_id = result[detection_idx]
|
|
if track_id is None:
|
|
track_id = self.next_id
|
|
self.next_id += 1
|
|
result[detection_idx] = track_id
|
|
self.tracks[track_id] = {"box": box, "last_frame": frame_idx}
|
|
return [int(track_id) for track_id in result]
|
|
|
|
|
|
@dataclass
|
|
class LivePersonState:
|
|
fusion: PersonTrack = field(default_factory=PersonTrack)
|
|
endpoints: dict[str, EndpointTrack] = field(
|
|
default_factory=lambda: {"left": EndpointTrack(), "right": EndpointTrack()}
|
|
)
|
|
bone_samples: dict[str, deque] = field(
|
|
default_factory=lambda: {
|
|
"left": deque(maxlen=120),
|
|
"right": deque(maxlen=120),
|
|
}
|
|
)
|
|
last_frame: int = -1
|
|
|
|
|
|
def bbox_diagonal(bbox: list[float]) -> float:
|
|
return max(math.hypot(bbox[2] - bbox[0], bbox[3] - bbox[1]), 1.0)
|
|
|
|
|
|
def update_bone_samples(
|
|
state: LivePersonState,
|
|
keypoints: np.ndarray,
|
|
metadata: dict,
|
|
bbox: list[float],
|
|
) -> None:
|
|
diagonal = bbox_diagonal(bbox)
|
|
for side, (shoulder, elbow, wrist) in BODY_SIDE_INDICES.items():
|
|
meta = metadata[side]
|
|
if meta["limb_state"] != "consensus" or not meta["distance_eligible"]:
|
|
continue
|
|
upper = float(np.linalg.norm(keypoints[shoulder] - keypoints[elbow])) / diagonal
|
|
forearm = float(np.linalg.norm(keypoints[elbow] - keypoints[wrist])) / diagonal
|
|
if 0.04 <= upper <= 0.65 and 0.04 <= forearm <= 0.65:
|
|
state.bone_samples[side].append((upper, forearm))
|
|
|
|
|
|
def current_bone_ratios(state: LivePersonState, side: str) -> tuple[float, float]:
|
|
samples = state.bone_samples[side]
|
|
if len(samples) < 10:
|
|
opposite = "right" if side == "left" else "left"
|
|
if len(state.bone_samples[opposite]) >= 10:
|
|
samples = state.bone_samples[opposite]
|
|
if not samples:
|
|
return 0.24, 0.22
|
|
values = np.asarray(samples, dtype=np.float64)
|
|
return float(np.median(values[:, 0])), float(np.median(values[:, 1]))
|
|
|
|
|
|
def update_live_v2(
|
|
frame_idx: int,
|
|
bbox: list[float],
|
|
selected_keypoints: np.ndarray,
|
|
smoothed_keypoints: np.ndarray,
|
|
fusion_metadata: dict,
|
|
state: LivePersonState,
|
|
assigned_measurements: dict[str, Measurement | None] | None,
|
|
identity_assignments: dict[str, dict] | None,
|
|
short_occlusion_frames: int,
|
|
medium_occlusion_frames: int,
|
|
) -> tuple[np.ndarray, dict]:
|
|
frame_gap = 0 if state.last_frame < 0 else max(frame_idx - state.last_frame - 1, 0)
|
|
for endpoint in state.endpoints.values():
|
|
endpoint.missed_frames += frame_gap
|
|
state.last_frame = frame_idx
|
|
update_bone_samples(state, selected_keypoints, fusion_metadata, bbox)
|
|
|
|
output = smoothed_keypoints.copy()
|
|
output_metadata = {}
|
|
diagonal = bbox_diagonal(bbox)
|
|
for side, (shoulder_idx, elbow_idx, wrist_idx) in BODY_SIDE_INDICES.items():
|
|
endpoint = state.endpoints[side]
|
|
meta = fusion_metadata[side]
|
|
start, end = HAND_SLICES[side]
|
|
measurement = (
|
|
assigned_measurements.get(side)
|
|
if assigned_measurements is not None
|
|
else None
|
|
)
|
|
if assigned_measurements is None and meta["alarm_eligible"]:
|
|
measurement = Measurement(
|
|
keypoints=selected_keypoints[start:end],
|
|
precise=bool(meta["distance_eligible"]),
|
|
provisional=not bool(meta["distance_eligible"]),
|
|
quality=float(meta.get("selected_hand_score", 0.0)),
|
|
)
|
|
if measurement is None:
|
|
measurement = temporal_recovery_measurement(
|
|
selected_keypoints,
|
|
meta,
|
|
side,
|
|
bbox,
|
|
endpoint,
|
|
medium_occlusion_frames,
|
|
)
|
|
returning_after_long = (
|
|
measurement is not None
|
|
and endpoint.missed_frames > short_occlusion_frames
|
|
)
|
|
wrist, hand, endpoint_state = update_endpoint_track(
|
|
endpoint,
|
|
measurement,
|
|
frame_idx,
|
|
short_occlusion_frames,
|
|
medium_occlusion_frames,
|
|
)
|
|
shoulder = smoothed_keypoints[shoulder_idx]
|
|
upper_ratio, forearm_ratio = current_bone_ratios(state, side)
|
|
upper_length, forearm_length = calibrated_bone_lengths(
|
|
upper_ratio, forearm_ratio, bbox
|
|
)
|
|
region = None
|
|
if wrist is not None and hand is not None:
|
|
reference_elbow = (
|
|
selected_keypoints[elbow_idx]
|
|
if returning_after_long or endpoint.last_elbow is None
|
|
else endpoint.last_elbow
|
|
)
|
|
use_ik = (
|
|
meta["limb_state"] != "consensus"
|
|
or endpoint_state
|
|
in ("provisional", "recovered", "reassigned", "predicted_short")
|
|
)
|
|
elbow = (
|
|
solve_two_bone_elbow(
|
|
shoulder,
|
|
wrist,
|
|
upper_length,
|
|
forearm_length,
|
|
reference_elbow,
|
|
)
|
|
if use_ik
|
|
else smoothed_keypoints[elbow_idx]
|
|
)
|
|
endpoint.last_elbow = elbow.copy()
|
|
direction = wrist - shoulder
|
|
if float(np.linalg.norm(direction)) > 1e-6:
|
|
endpoint.last_direction = direction / np.linalg.norm(direction)
|
|
output[shoulder_idx] = shoulder
|
|
output[elbow_idx] = elbow
|
|
output[wrist_idx] = wrist
|
|
output[start:end] = hand
|
|
uncertainty = (
|
|
0.025 * diagonal
|
|
if endpoint_state == "observed"
|
|
else (0.04 + 0.015 * endpoint.missed_frames) * diagonal
|
|
)
|
|
precise = endpoint_state == "observed"
|
|
else:
|
|
region = reachable_region(
|
|
shoulder,
|
|
upper_length,
|
|
forearm_length,
|
|
bbox,
|
|
endpoint,
|
|
endpoint_state,
|
|
short_occlusion_frames,
|
|
medium_occlusion_frames,
|
|
)
|
|
uncertainty = region["outer_radius_px"]
|
|
precise = False
|
|
output_metadata[side] = {
|
|
"state": endpoint_state,
|
|
"precise_distance": precise,
|
|
"fast_alarm_candidate": True,
|
|
"risk_mode": (
|
|
"precise_hand"
|
|
if endpoint_state == "observed"
|
|
else (
|
|
"motion_corridor_intersection"
|
|
if endpoint_state == "reachable_medium"
|
|
else (
|
|
"door_intersection_only"
|
|
if endpoint_state == "reachable_long"
|
|
else (
|
|
"reassigned_hand"
|
|
if endpoint_state == "reassigned"
|
|
else "temporal_hand"
|
|
)
|
|
)
|
|
)
|
|
),
|
|
"requires_door_intersection": endpoint_state
|
|
in ("reachable_medium", "reachable_long"),
|
|
"identity_assignment": (
|
|
identity_assignments.get(side)
|
|
if identity_assignments is not None
|
|
else None
|
|
),
|
|
"missed_frames": endpoint.missed_frames,
|
|
"uncertainty_radius_px": round(float(uncertainty), 3),
|
|
"bone_lengths_px": {
|
|
"upper": round(upper_length, 3),
|
|
"forearm": round(forearm_length, 3),
|
|
},
|
|
"reachable_region": region,
|
|
}
|
|
return output, output_metadata
|
|
|
|
|
|
def mask_other_people(
|
|
frame: np.ndarray,
|
|
boxes: list[list[float]],
|
|
target_idx: int,
|
|
) -> np.ndarray:
|
|
masked = frame.copy()
|
|
height, width = frame.shape[:2]
|
|
mask = np.zeros((height, width), dtype=bool)
|
|
target = boxes[target_idx]
|
|
target_area = max((target[2] - target[0]) * (target[3] - target[1]), 1.0)
|
|
for idx, bbox in enumerate(boxes):
|
|
if idx == target_idx:
|
|
continue
|
|
other_area = max((bbox[2] - bbox[0]) * (bbox[3] - bbox[1]), 1.0)
|
|
# The failure in this video is the small background person being
|
|
# absorbed by the much larger foreground person's 3:4 pose crop.
|
|
# Mask only markedly smaller people. Masking a larger occluder while
|
|
# estimating a small person removes genuine target pixels at overlap.
|
|
if other_area >= 0.5 * target_area:
|
|
continue
|
|
x1, y1, x2, y2 = (int(round(value)) for value in bbox)
|
|
x1, x2 = max(0, x1), min(width - 1, x2)
|
|
y1, y2 = max(0, y1), min(height - 1, y2)
|
|
mask[y1:y2 + 1, x1:x2 + 1] = True
|
|
|
|
masked[mask] = (124, 116, 104)
|
|
return masked
|