531 lines
18 KiB
Python
531 lines
18 KiB
Python
"""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
|