Trim use-can to buildable runtime sources and standalone Orin package
This commit is contained in:
@@ -0,0 +1,530 @@
|
||||
"""Runtime pose fusion helpers."""
|
||||
from __future__ import annotations
|
||||
import math
|
||||
from dataclasses import dataclass, field
|
||||
import cv2
|
||||
import numpy as np
|
||||
from pose_geometry import BODY_WRISTS, HAND_EDGES_LOCAL, HAND_SLICES, normalized_keypoints
|
||||
|
||||
|
||||
BODY_SIDE_INDICES = {"left": (5, 7, 9), "right": (6, 8, 10)}
|
||||
|
||||
|
||||
SIDE_COLORS = {"left": (40, 40, 255), "right": (30, 230, 30)}
|
||||
|
||||
|
||||
HAND_COLORS = {"left": (255, 80, 220), "right": (255, 190, 40)}
|
||||
|
||||
|
||||
@dataclass
|
||||
class Candidate:
|
||||
source: str
|
||||
keypoints: np.ndarray
|
||||
scores: np.ndarray
|
||||
normalized: np.ndarray
|
||||
limb_score: float
|
||||
hand_score: float
|
||||
upper_length: float
|
||||
forearm_length: float
|
||||
hand_span: float
|
||||
wrist_link: float
|
||||
limb_valid: bool
|
||||
hand_valid: bool
|
||||
|
||||
|
||||
@dataclass
|
||||
class SideTrack:
|
||||
limb_previous: np.ndarray | None = None
|
||||
limb_velocity: np.ndarray | None = None
|
||||
hand_previous: np.ndarray | None = None
|
||||
hand_velocity: np.ndarray | None = None
|
||||
limb_source: str | None = None
|
||||
hand_source: str | None = None
|
||||
previous_hand_eligible: bool = False
|
||||
|
||||
|
||||
@dataclass
|
||||
class PersonTrack:
|
||||
sides: dict[str, SideTrack] = field(
|
||||
default_factory=lambda: {"left": SideTrack(), "right": SideTrack()}
|
||||
)
|
||||
smoothed: np.ndarray | None = None
|
||||
|
||||
|
||||
def bbox_diagonal(bbox: list[float]) -> float:
|
||||
return max(math.hypot(bbox[2] - bbox[0], bbox[3] - bbox[1]), 1.0)
|
||||
|
||||
|
||||
def build_candidate(
|
||||
source: str,
|
||||
keypoints: np.ndarray,
|
||||
scores: np.ndarray,
|
||||
bbox: list[float],
|
||||
side: str,
|
||||
) -> Candidate:
|
||||
normalized = normalized_keypoints(keypoints, bbox)
|
||||
shoulder, elbow, wrist = BODY_SIDE_INDICES[side]
|
||||
hand_start, hand_end = HAND_SLICES[side]
|
||||
hand = keypoints[hand_start:hand_end]
|
||||
diagonal = bbox_diagonal(bbox)
|
||||
upper_length = float(np.linalg.norm(normalized[shoulder] - normalized[elbow]))
|
||||
forearm_length = float(np.linalg.norm(normalized[elbow] - normalized[wrist]))
|
||||
hand_span = math.hypot(float(np.ptp(hand[:, 0])), float(np.ptp(hand[:, 1]))) / diagonal
|
||||
wrist_link = float(np.linalg.norm(keypoints[wrist] - hand[0])) / diagonal
|
||||
limb_score = float(np.mean(scores[[shoulder, elbow, wrist]]))
|
||||
hand_score = float(np.mean(scores[hand_start:hand_end]))
|
||||
limb_valid = 0.025 <= upper_length <= 0.8 and 0.025 <= forearm_length <= 0.8
|
||||
hand_valid = 0.008 <= hand_span <= 0.35 and wrist_link <= 0.12
|
||||
return Candidate(
|
||||
source=source,
|
||||
keypoints=keypoints,
|
||||
scores=scores,
|
||||
normalized=normalized,
|
||||
limb_score=limb_score,
|
||||
hand_score=hand_score,
|
||||
upper_length=upper_length,
|
||||
forearm_length=forearm_length,
|
||||
hand_span=hand_span,
|
||||
wrist_link=wrist_link,
|
||||
limb_valid=limb_valid,
|
||||
hand_valid=hand_valid,
|
||||
)
|
||||
|
||||
|
||||
def predicted(previous: np.ndarray | None, velocity: np.ndarray | None) -> np.ndarray | None:
|
||||
if previous is None:
|
||||
return None
|
||||
if velocity is None:
|
||||
return previous
|
||||
return previous + 0.65 * velocity
|
||||
|
||||
|
||||
def robust_distance(points: np.ndarray, target: np.ndarray | None) -> float:
|
||||
if target is None:
|
||||
return 0.0
|
||||
distances = np.linalg.norm(points - target, axis=1)
|
||||
return float(np.median(distances))
|
||||
|
||||
|
||||
def select_limb(
|
||||
candidates: list[Candidate],
|
||||
side: str,
|
||||
track: SideTrack | None,
|
||||
) -> tuple[Candidate, dict[str, float]]:
|
||||
indices = np.asarray(BODY_SIDE_INDICES[side])
|
||||
target = predicted(track.limb_previous, track.limb_velocity) if track else None
|
||||
scored: list[tuple[float, Candidate, float]] = []
|
||||
for candidate in candidates:
|
||||
temporal_cost = robust_distance(candidate.normalized[indices], target)
|
||||
geometry_bonus = 1.2 if candidate.limb_valid else -5.0
|
||||
switch_bonus = 0.35 if track and candidate.source == track.limb_source else 0.0
|
||||
score = candidate.limb_score + geometry_bonus + switch_bonus - 7.0 * temporal_cost
|
||||
scored.append((score, candidate, temporal_cost))
|
||||
score, selected, temporal_cost = max(scored, key=lambda item: item[0])
|
||||
return selected, {"score": score, "temporal_cost": temporal_cost}
|
||||
|
||||
|
||||
def hand_link_to_limb(
|
||||
candidate: Candidate,
|
||||
selected_limb: Candidate,
|
||||
bbox: list[float],
|
||||
side: str,
|
||||
) -> float:
|
||||
hand_start, _ = HAND_SLICES[side]
|
||||
wrist_idx = BODY_WRISTS[side]
|
||||
return float(
|
||||
np.linalg.norm(candidate.keypoints[hand_start] - selected_limb.keypoints[wrist_idx])
|
||||
) / bbox_diagonal(bbox)
|
||||
|
||||
|
||||
def select_hand(
|
||||
candidates: list[Candidate],
|
||||
selected_limb: Candidate,
|
||||
bbox: list[float],
|
||||
side: str,
|
||||
track: SideTrack | None,
|
||||
) -> tuple[Candidate, bool, dict[str, float]]:
|
||||
start, end = HAND_SLICES[side]
|
||||
target = predicted(track.hand_previous, track.hand_velocity) if track else None
|
||||
scored: list[tuple[float, Candidate, bool, float, float]] = []
|
||||
for candidate in candidates:
|
||||
temporal_cost = robust_distance(candidate.normalized[start:end], target)
|
||||
link = hand_link_to_limb(candidate, selected_limb, bbox, side)
|
||||
valid = candidate.hand_valid
|
||||
geometry_bonus = 3.0 if valid else -7.0
|
||||
switch_bonus = 0.45 if track and candidate.source == track.hand_source else 0.0
|
||||
limb_source_bonus = 0.15 if candidate.source == selected_limb.source else 0.0
|
||||
# A newly visible endpoint may disagree with the currently selected
|
||||
# (occlusion-driven) limb. Keep this as a soft cost so strong image
|
||||
# evidence can switch both the hand and its limb together.
|
||||
link_penalty = 1.5 * min(max(link - 0.12, 0.0), 0.2)
|
||||
score = (
|
||||
candidate.hand_score
|
||||
+ geometry_bonus
|
||||
+ switch_bonus
|
||||
+ limb_source_bonus
|
||||
- link_penalty
|
||||
- 8.0 * min(temporal_cost, 0.15)
|
||||
)
|
||||
scored.append((score, candidate, valid, temporal_cost, link))
|
||||
score, selected, valid, temporal_cost, link = max(scored, key=lambda item: item[0])
|
||||
return selected, valid, {
|
||||
"score": score,
|
||||
"temporal_cost": temporal_cost,
|
||||
"link_to_selected_limb": link,
|
||||
}
|
||||
|
||||
|
||||
def candidate_disagreement(
|
||||
raw: Candidate,
|
||||
isolated: Candidate,
|
||||
side: str,
|
||||
) -> tuple[float, float]:
|
||||
shoulder, elbow, wrist = BODY_SIDE_INDICES[side]
|
||||
start, end = HAND_SLICES[side]
|
||||
limb = float(
|
||||
np.median(
|
||||
np.linalg.norm(
|
||||
raw.normalized[[shoulder, elbow, wrist]]
|
||||
- isolated.normalized[[shoulder, elbow, wrist]],
|
||||
axis=1,
|
||||
)
|
||||
)
|
||||
)
|
||||
hand_distances = np.linalg.norm(
|
||||
raw.normalized[start:end] - isolated.normalized[start:end], axis=1
|
||||
)
|
||||
hand = float(np.percentile(hand_distances, 90))
|
||||
return limb, hand
|
||||
|
||||
|
||||
def suppress_duplicate_hands(
|
||||
fused: np.ndarray,
|
||||
metadata: dict,
|
||||
bbox: list[float],
|
||||
) -> None:
|
||||
left_start, left_end = HAND_SLICES["left"]
|
||||
right_start, right_end = HAND_SLICES["right"]
|
||||
left_center = np.mean(fused[left_start:left_end], axis=0)
|
||||
right_center = np.mean(fused[right_start:right_end], axis=0)
|
||||
separation = float(np.linalg.norm(left_center - right_center)) / bbox_diagonal(bbox)
|
||||
if separation >= 0.08:
|
||||
return
|
||||
weaker = min(
|
||||
("left", "right"),
|
||||
key=lambda side: metadata[side]["selected_hand_score"],
|
||||
)
|
||||
metadata[weaker]["hand_state"] = "duplicate"
|
||||
metadata[weaker]["distance_eligible"] = False
|
||||
metadata[weaker]["alarm_eligible"] = False
|
||||
metadata[weaker]["duplicate_hand_separation"] = round(separation, 5)
|
||||
|
||||
|
||||
def fuse_person(
|
||||
raw_keypoints: np.ndarray,
|
||||
raw_scores: np.ndarray,
|
||||
isolated_keypoints: np.ndarray,
|
||||
isolated_scores: np.ndarray,
|
||||
bbox: list[float],
|
||||
track: PersonTrack | None,
|
||||
) -> tuple[np.ndarray, dict]:
|
||||
fused = raw_keypoints.copy()
|
||||
metadata: dict[str, dict] = {}
|
||||
|
||||
for side in ("left", "right"):
|
||||
side_track = track.sides[side] if track else None
|
||||
raw = build_candidate("raw", raw_keypoints, raw_scores, bbox, side)
|
||||
isolated = build_candidate(
|
||||
"isolated", isolated_keypoints, isolated_scores, bbox, side
|
||||
)
|
||||
candidates = [raw, isolated]
|
||||
limb, limb_detail = select_limb(candidates, side, side_track)
|
||||
hand, hand_valid, hand_detail = select_hand(
|
||||
candidates, limb, bbox, side, side_track
|
||||
)
|
||||
|
||||
other_hand = isolated if hand.source == "raw" else raw
|
||||
endpoint_advantage = hand.hand_score - other_hand.hand_score
|
||||
if hand.source != limb.source and hand_valid and endpoint_advantage >= 0.6:
|
||||
limb = hand
|
||||
limb_detail["visible_endpoint_override"] = 1.0
|
||||
hand_detail["link_to_selected_limb"] = hand.wrist_link
|
||||
|
||||
body_indices = np.asarray(BODY_SIDE_INDICES[side])
|
||||
hand_start, hand_end = HAND_SLICES[side]
|
||||
fused[body_indices] = limb.keypoints[body_indices]
|
||||
fused[hand_start:hand_end] = hand.keypoints[hand_start:hand_end]
|
||||
|
||||
limb_disagreement, hand_disagreement = candidate_disagreement(raw, isolated, side)
|
||||
if limb_disagreement <= 0.05:
|
||||
limb_state = "consensus"
|
||||
else:
|
||||
limb_state = "inferred"
|
||||
|
||||
score_advantage = raw.hand_score - isolated.hand_score
|
||||
raw_temporally_stable = hand_detail["temporal_cost"] <= 0.08
|
||||
if not hand_valid:
|
||||
hand_state = "rejected"
|
||||
distance_eligible = False
|
||||
elif hand_disagreement <= 0.06:
|
||||
hand_state = "consensus"
|
||||
distance_eligible = True
|
||||
elif (
|
||||
hand.source == "raw"
|
||||
and score_advantage >= 0.6
|
||||
and raw_temporally_stable
|
||||
):
|
||||
hand_state = "visible_endpoint"
|
||||
distance_eligible = True
|
||||
elif hand.source == "isolated":
|
||||
hand_state = "inferred"
|
||||
distance_eligible = False
|
||||
else:
|
||||
hand_state = "conflict"
|
||||
distance_eligible = False
|
||||
fast_alarm_endpoint = (
|
||||
hand_valid and hand.source == "raw" and score_advantage >= 0.6
|
||||
)
|
||||
alarm_eligible = distance_eligible or fast_alarm_endpoint
|
||||
|
||||
metadata[side] = {
|
||||
"limb_source": limb.source,
|
||||
"hand_source": hand.source,
|
||||
"limb_state": limb_state,
|
||||
"hand_state": hand_state,
|
||||
"distance_eligible": distance_eligible,
|
||||
"alarm_eligible": alarm_eligible,
|
||||
"limb_disagreement": round(limb_disagreement, 5),
|
||||
"hand_disagreement": round(hand_disagreement, 5),
|
||||
"raw_hand_score": round(raw.hand_score, 4),
|
||||
"isolated_hand_score": round(isolated.hand_score, 4),
|
||||
"selected_hand_span": round(hand.hand_span, 5),
|
||||
"selected_hand_score": round(hand.hand_score, 4),
|
||||
"selected_hand_valid": hand_valid,
|
||||
"limb_selection": {key: round(value, 5) for key, value in limb_detail.items()},
|
||||
"hand_selection": {key: round(value, 5) for key, value in hand_detail.items()},
|
||||
}
|
||||
|
||||
if side_track is not None:
|
||||
limb_now = limb.normalized[body_indices]
|
||||
hand_now = hand.normalized[hand_start:hand_end]
|
||||
if side_track.limb_previous is not None:
|
||||
delta = limb_now - side_track.limb_previous
|
||||
side_track.limb_velocity = (
|
||||
delta
|
||||
if side_track.limb_velocity is None
|
||||
else 0.65 * side_track.limb_velocity + 0.35 * delta
|
||||
)
|
||||
if hand_valid:
|
||||
if side_track.hand_previous is not None:
|
||||
delta = hand_now - side_track.hand_previous
|
||||
side_track.hand_velocity = (
|
||||
delta
|
||||
if side_track.hand_velocity is None
|
||||
else 0.65 * side_track.hand_velocity + 0.35 * delta
|
||||
)
|
||||
side_track.hand_previous = hand_now
|
||||
side_track.hand_source = hand.source
|
||||
side_track.limb_previous = limb_now
|
||||
side_track.limb_source = limb.source
|
||||
side_track.previous_hand_eligible = distance_eligible
|
||||
|
||||
suppress_duplicate_hands(fused, metadata, bbox)
|
||||
return fused, metadata
|
||||
|
||||
|
||||
def adaptive_smooth(
|
||||
keypoints: np.ndarray,
|
||||
bbox: list[float],
|
||||
track: PersonTrack,
|
||||
) -> np.ndarray:
|
||||
normalized = normalized_keypoints(keypoints, bbox)
|
||||
if track.smoothed is None:
|
||||
track.smoothed = normalized
|
||||
else:
|
||||
displacement = np.linalg.norm(normalized - track.smoothed, axis=1)
|
||||
alpha = 0.32 + 0.58 * np.clip(displacement / 0.08, 0.0, 1.0)
|
||||
track.smoothed = alpha[:, None] * normalized + (1.0 - alpha[:, None]) * track.smoothed
|
||||
|
||||
x1, y1, x2, y2 = bbox
|
||||
result = track.smoothed.copy()
|
||||
result[:, 0] = result[:, 0] * max(x2 - x1, 1.0) + x1
|
||||
result[:, 1] = result[:, 1] * max(y2 - y1, 1.0) + y1
|
||||
return result
|
||||
|
||||
|
||||
def build_global_hands(
|
||||
keypoints: list[np.ndarray],
|
||||
metadata: list[dict],
|
||||
) -> list[dict]:
|
||||
candidates = []
|
||||
for person_idx, (person_keypoints, person_meta) in enumerate(
|
||||
zip(keypoints, metadata), start=1
|
||||
):
|
||||
for side, (start, end) in HAND_SLICES.items():
|
||||
meta = person_meta[side]
|
||||
if not meta["alarm_eligible"]:
|
||||
continue
|
||||
hand = person_keypoints[start:end]
|
||||
span_px = math.hypot(float(np.ptp(hand[:, 0])), float(np.ptp(hand[:, 1])))
|
||||
candidates.append(
|
||||
{
|
||||
"person_id": person_idx,
|
||||
"side": side,
|
||||
"source": meta["hand_source"],
|
||||
"state": meta["hand_state"],
|
||||
"distance_eligible": meta["distance_eligible"],
|
||||
"alarm_eligible": meta["alarm_eligible"],
|
||||
"score": meta["selected_hand_score"],
|
||||
"center": np.mean(hand, axis=0),
|
||||
"span_px": span_px,
|
||||
"keypoints": hand,
|
||||
"aliases": [],
|
||||
}
|
||||
)
|
||||
|
||||
kept: list[dict] = []
|
||||
for candidate in sorted(candidates, key=lambda item: item["score"], reverse=True):
|
||||
duplicate = None
|
||||
for existing in kept:
|
||||
threshold = float(np.clip(0.45 * max(candidate["span_px"], existing["span_px"]), 18, 70))
|
||||
if float(np.linalg.norm(candidate["center"] - existing["center"])) <= threshold:
|
||||
duplicate = existing
|
||||
break
|
||||
if duplicate is None:
|
||||
kept.append(candidate)
|
||||
else:
|
||||
duplicate["aliases"].append(
|
||||
{"person_id": candidate["person_id"], "side": candidate["side"]}
|
||||
)
|
||||
|
||||
return [
|
||||
{
|
||||
"person_id": item["person_id"],
|
||||
"side": item["side"],
|
||||
"source": item["source"],
|
||||
"state": item["state"],
|
||||
"distance_eligible": item["distance_eligible"],
|
||||
"alarm_eligible": item["alarm_eligible"],
|
||||
"score": item["score"],
|
||||
"center": np.round(item["center"], 3).tolist(),
|
||||
"span_px": round(item["span_px"], 3),
|
||||
"keypoints": np.round(item["keypoints"], 3).tolist(),
|
||||
"aliases": item["aliases"],
|
||||
}
|
||||
for item in kept
|
||||
]
|
||||
|
||||
|
||||
def draw_dashed_line(
|
||||
image: np.ndarray,
|
||||
start: np.ndarray,
|
||||
end: np.ndarray,
|
||||
color: tuple[int, int, int],
|
||||
thickness: int = 2,
|
||||
) -> None:
|
||||
a = np.asarray(start, dtype=np.float64)
|
||||
b = np.asarray(end, dtype=np.float64)
|
||||
length = float(np.linalg.norm(b - a))
|
||||
if length < 1.0:
|
||||
return
|
||||
direction = (b - a) / length
|
||||
for offset in np.arange(0.0, length, 14.0):
|
||||
segment_start = a + direction * offset
|
||||
segment_end = a + direction * min(offset + 8.0, length)
|
||||
cv2.line(
|
||||
image,
|
||||
tuple(np.round(segment_start).astype(int)),
|
||||
tuple(np.round(segment_end).astype(int)),
|
||||
color,
|
||||
thickness,
|
||||
cv2.LINE_AA,
|
||||
)
|
||||
|
||||
|
||||
def draw_fused(
|
||||
image: np.ndarray,
|
||||
boxes: list[list[float]],
|
||||
keypoints: list[np.ndarray],
|
||||
metadata: list[dict],
|
||||
title: str,
|
||||
) -> np.ndarray:
|
||||
canvas = image.copy()
|
||||
for person_idx, (bbox, kpts, person_meta) in enumerate(
|
||||
zip(boxes, keypoints, metadata), start=1
|
||||
):
|
||||
x1, y1, x2, y2 = (int(round(value)) for value in bbox)
|
||||
cv2.rectangle(canvas, (x1, y1), (x2, y2), (0, 220, 255), 2)
|
||||
cv2.putText(
|
||||
canvas,
|
||||
f"P{person_idx}",
|
||||
(x1 + 4, max(22, y1 + 24)),
|
||||
cv2.FONT_HERSHEY_SIMPLEX,
|
||||
0.72,
|
||||
(0, 220, 255),
|
||||
2,
|
||||
cv2.LINE_AA,
|
||||
)
|
||||
for side in ("left", "right"):
|
||||
shoulder, elbow, wrist = BODY_SIDE_INDICES[side]
|
||||
side_meta = person_meta[side]
|
||||
color = SIDE_COLORS[side]
|
||||
for a, b in ((shoulder, elbow), (elbow, wrist)):
|
||||
if side_meta["limb_state"] == "consensus":
|
||||
cv2.line(
|
||||
canvas,
|
||||
tuple(np.round(kpts[a]).astype(int)),
|
||||
tuple(np.round(kpts[b]).astype(int)),
|
||||
color,
|
||||
3,
|
||||
cv2.LINE_AA,
|
||||
)
|
||||
else:
|
||||
draw_dashed_line(canvas, kpts[a], kpts[b], (0, 140, 255), 3)
|
||||
|
||||
start, end = HAND_SLICES[side]
|
||||
hand = kpts[start:end]
|
||||
if side_meta["hand_state"] != "rejected":
|
||||
hand_color = HAND_COLORS[side]
|
||||
if not side_meta["distance_eligible"]:
|
||||
hand_color = (0, 140, 255)
|
||||
for a, b in HAND_EDGES_LOCAL:
|
||||
if side_meta["distance_eligible"]:
|
||||
cv2.line(
|
||||
canvas,
|
||||
tuple(np.round(hand[a]).astype(int)),
|
||||
tuple(np.round(hand[b]).astype(int)),
|
||||
hand_color,
|
||||
2,
|
||||
cv2.LINE_AA,
|
||||
)
|
||||
else:
|
||||
draw_dashed_line(canvas, hand[a], hand[b], hand_color, 2)
|
||||
label = (
|
||||
f"{side[0].upper()}:{side_meta['hand_state'][:4]}"
|
||||
f"/{side_meta['hand_source'][0]}"
|
||||
)
|
||||
label_y = max(45, y1 + (48 if side == "left" else 70))
|
||||
cv2.putText(
|
||||
canvas,
|
||||
label,
|
||||
(x1 + 4, label_y),
|
||||
cv2.FONT_HERSHEY_SIMPLEX,
|
||||
0.5,
|
||||
HAND_COLORS[side] if side_meta["distance_eligible"] else (0, 140, 255),
|
||||
1,
|
||||
cv2.LINE_AA,
|
||||
)
|
||||
|
||||
cv2.rectangle(canvas, (0, 0), (650, 45), (0, 0, 0), -1)
|
||||
cv2.putText(
|
||||
canvas,
|
||||
title,
|
||||
(10, 31),
|
||||
cv2.FONT_HERSHEY_SIMPLEX,
|
||||
0.8,
|
||||
(255, 255, 255),
|
||||
2,
|
||||
cv2.LINE_AA,
|
||||
)
|
||||
return canvas
|
||||
Reference in New Issue
Block a user