Trim use-can to buildable runtime sources and standalone Orin package
This commit is contained in:
@@ -0,0 +1,304 @@
|
||||
"""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
|
||||
Reference in New Issue
Block a user