856 lines
31 KiB
Python
856 lines
31 KiB
Python
"""Fuse raw and person-isolated RTMW-L poses with temporal consistency.
|
|
|
|
The fusion keeps the two inference views as hypotheses. It selects body limbs
|
|
and hands separately, rejects implausible hand geometry, adds source hysteresis
|
|
and constant-velocity consistency, then applies a motion-adaptive EMA in
|
|
person-box coordinates. Occluded/conflicting limbs remain labelled as such so
|
|
they are not accidentally consumed as precise anti-pinch distance evidence.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import math
|
|
from collections import Counter, defaultdict
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
from compare_rtmw_ml import (
|
|
BODY_EDGES,
|
|
BODY_WRISTS,
|
|
HAND_EDGES_LOCAL,
|
|
HAND_SLICES,
|
|
draw_pose,
|
|
normalized_keypoints,
|
|
summarize,
|
|
temporal_metrics,
|
|
)
|
|
|
|
|
|
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 parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--video", type=Path, required=True)
|
|
parser.add_argument("--raw-jsonl", type=Path, required=True)
|
|
parser.add_argument("--isolated-jsonl", type=Path, required=True)
|
|
parser.add_argument("--output", type=Path, required=True)
|
|
return parser.parse_args()
|
|
|
|
|
|
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
|
|
|
|
|
|
def make_montage(
|
|
video_path: Path,
|
|
records: list[dict],
|
|
selected_frames: list[int],
|
|
output_path: Path,
|
|
) -> None:
|
|
selected = set(selected_frames)
|
|
by_frame = {record["frame_index"]: record for record in records}
|
|
cap = cv2.VideoCapture(str(video_path))
|
|
rows = []
|
|
frame_idx = 0
|
|
while True:
|
|
ok, frame = cap.read()
|
|
if not ok:
|
|
break
|
|
if frame_idx in selected:
|
|
record = by_frame[frame_idx]
|
|
boxes = record["boxes"]
|
|
raw = draw_pose(
|
|
frame,
|
|
boxes,
|
|
np.asarray(record["raw_keypoints"]),
|
|
f"RTMW-L raw | frame {frame_idx}",
|
|
)
|
|
spatial = draw_fused(
|
|
frame,
|
|
boxes,
|
|
[np.asarray(points) for points in record["spatial_keypoints"]],
|
|
record["spatial_metadata"],
|
|
f"spatial fusion | frame {frame_idx}",
|
|
)
|
|
temporal = draw_fused(
|
|
frame,
|
|
boxes,
|
|
[np.asarray(points) for points in record["temporal_keypoints"]],
|
|
record["temporal_metadata"],
|
|
f"temporal fusion | frame {frame_idx}",
|
|
)
|
|
panels = [
|
|
cv2.resize(panel, (480, 270), interpolation=cv2.INTER_AREA)
|
|
for panel in (raw, spatial, temporal)
|
|
]
|
|
rows.append(np.hstack(panels))
|
|
frame_idx += 1
|
|
cap.release()
|
|
if rows:
|
|
cv2.imwrite(str(output_path), np.vstack(rows))
|
|
|
|
|
|
def active_hand_temporal_metrics(records: list[dict], variant: str) -> dict:
|
|
values: dict[str, list[float]] = defaultdict(list)
|
|
previous: dict[tuple[int, str], tuple[int, np.ndarray]] = {}
|
|
for record in records:
|
|
for person_idx, bbox in enumerate(record["boxes"]):
|
|
kpts = normalized_keypoints(
|
|
np.asarray(record[f"{variant}_keypoints"][person_idx]), bbox
|
|
)
|
|
meta = record[f"{variant}_metadata"][person_idx]
|
|
for side, (start, end) in HAND_SLICES.items():
|
|
key = (person_idx, side)
|
|
if not meta[side]["distance_eligible"]:
|
|
previous.pop(key, None)
|
|
continue
|
|
prior = previous.get(key)
|
|
if prior is not None and prior[0] == record["frame_index"] - 1:
|
|
step = np.linalg.norm(kpts[start:end] - prior[1], axis=1)
|
|
values[f"P{person_idx + 1}_{side}"].extend(float(v) for v in step)
|
|
previous[key] = (record["frame_index"], kpts[start:end])
|
|
return {name: summarize(items) for name, items in values.items()}
|
|
|
|
|
|
def metadata_summary(records: list[dict], variant: str) -> dict:
|
|
state_counts: Counter = Counter()
|
|
source_counts: Counter = Counter()
|
|
eligible_counts: Counter = Counter()
|
|
alarm_counts: Counter = Counter()
|
|
switches: Counter = Counter()
|
|
previous_source: dict[tuple[int, str, str], str] = {}
|
|
for record in records:
|
|
for person_idx, person_meta in enumerate(record[f"{variant}_metadata"]):
|
|
for side in ("left", "right"):
|
|
meta = person_meta[side]
|
|
key_prefix = f"P{person_idx + 1}_{side}"
|
|
state_counts[f"{key_prefix}_{meta['hand_state']}"] += 1
|
|
source_counts[f"{key_prefix}_{meta['hand_source']}"] += 1
|
|
if meta["distance_eligible"]:
|
|
eligible_counts[key_prefix] += 1
|
|
if meta["alarm_eligible"]:
|
|
alarm_counts[key_prefix] += 1
|
|
for kind in ("limb", "hand"):
|
|
source = meta[f"{kind}_source"]
|
|
key = (person_idx, side, kind)
|
|
if key in previous_source and previous_source[key] != source:
|
|
switches[f"{key_prefix}_{kind}"] += 1
|
|
previous_source[key] = source
|
|
return {
|
|
"hand_state_counts": dict(sorted(state_counts.items())),
|
|
"source_counts": dict(sorted(source_counts.items())),
|
|
"distance_eligible_frames": dict(sorted(eligible_counts.items())),
|
|
"alarm_eligible_frames": dict(sorted(alarm_counts.items())),
|
|
"source_switches": dict(sorted(switches.items())),
|
|
}
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
args.output.mkdir(parents=True, exist_ok=True)
|
|
raw_records = [
|
|
json.loads(line) for line in args.raw_jsonl.read_text(encoding="utf-8").splitlines()
|
|
]
|
|
isolated_records = [
|
|
json.loads(line)
|
|
for line in args.isolated_jsonl.read_text(encoding="utf-8").splitlines()
|
|
]
|
|
|
|
cap = cv2.VideoCapture(str(args.video))
|
|
fps = float(cap.get(cv2.CAP_PROP_FPS))
|
|
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
|
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
|
writer = cv2.VideoWriter(
|
|
str(args.output / "rtmw_l_raw_spatial_temporal.mp4"),
|
|
cv2.VideoWriter_fourcc(*"mp4v"),
|
|
fps,
|
|
(1920, 360),
|
|
)
|
|
|
|
tracks = [PersonTrack(), PersonTrack()]
|
|
output_records: list[dict] = []
|
|
frame_idx = 0
|
|
while True:
|
|
ok, frame = cap.read()
|
|
if not ok:
|
|
break
|
|
raw_record = raw_records[frame_idx]
|
|
isolated_record = isolated_records[frame_idx]
|
|
boxes = raw_record["boxes"][:2]
|
|
raw_keypoints = np.asarray(raw_record["L"]["keypoints"][:2], dtype=np.float64)
|
|
raw_scores = np.asarray(raw_record["L"]["raw_scores"][:2], dtype=np.float64)
|
|
isolated_keypoints = np.asarray(
|
|
isolated_record["L"]["keypoints"][:2], dtype=np.float64
|
|
)
|
|
isolated_scores = np.asarray(
|
|
isolated_record["masked_raw_scores"][:2], dtype=np.float64
|
|
)
|
|
|
|
spatial_keypoints: list[np.ndarray] = []
|
|
spatial_metadata: list[dict] = []
|
|
temporal_keypoints: list[np.ndarray] = []
|
|
temporal_selected_keypoints: list[np.ndarray] = []
|
|
temporal_metadata: list[dict] = []
|
|
for person_idx in range(2):
|
|
spatial, spatial_meta = fuse_person(
|
|
raw_keypoints[person_idx],
|
|
raw_scores[person_idx],
|
|
isolated_keypoints[person_idx],
|
|
isolated_scores[person_idx],
|
|
boxes[person_idx],
|
|
track=None,
|
|
)
|
|
temporal_unfiltered, temporal_meta = fuse_person(
|
|
raw_keypoints[person_idx],
|
|
raw_scores[person_idx],
|
|
isolated_keypoints[person_idx],
|
|
isolated_scores[person_idx],
|
|
boxes[person_idx],
|
|
track=tracks[person_idx],
|
|
)
|
|
temporal = adaptive_smooth(
|
|
temporal_unfiltered, boxes[person_idx], tracks[person_idx]
|
|
)
|
|
spatial_keypoints.append(spatial)
|
|
spatial_metadata.append(spatial_meta)
|
|
temporal_keypoints.append(temporal)
|
|
temporal_selected_keypoints.append(temporal_unfiltered)
|
|
temporal_metadata.append(temporal_meta)
|
|
|
|
global_hands = build_global_hands(
|
|
temporal_selected_keypoints, temporal_metadata
|
|
)
|
|
|
|
output_record = {
|
|
"frame_index": frame_idx,
|
|
"boxes": boxes,
|
|
"raw_keypoints": np.round(raw_keypoints, 3).tolist(),
|
|
"spatial_keypoints": np.round(np.asarray(spatial_keypoints), 3).tolist(),
|
|
"spatial_metadata": spatial_metadata,
|
|
"temporal_keypoints": np.round(np.asarray(temporal_keypoints), 3).tolist(),
|
|
"temporal_selected_keypoints": np.round(
|
|
np.asarray(temporal_selected_keypoints), 3
|
|
).tolist(),
|
|
"temporal_metadata": temporal_metadata,
|
|
"global_hands": global_hands,
|
|
}
|
|
output_records.append(output_record)
|
|
|
|
raw_panel = draw_pose(frame, boxes, raw_keypoints, f"RTMW-L raw | frame {frame_idx}")
|
|
spatial_panel = draw_fused(
|
|
frame,
|
|
boxes,
|
|
spatial_keypoints,
|
|
spatial_metadata,
|
|
f"spatial fusion | frame {frame_idx}",
|
|
)
|
|
temporal_panel = draw_fused(
|
|
frame,
|
|
boxes,
|
|
temporal_keypoints,
|
|
temporal_metadata,
|
|
f"temporal fusion | frame {frame_idx}",
|
|
)
|
|
panels = [
|
|
cv2.resize(panel, (640, 360), interpolation=cv2.INTER_AREA)
|
|
for panel in (raw_panel, spatial_panel, temporal_panel)
|
|
]
|
|
writer.write(np.hstack(panels))
|
|
frame_idx += 1
|
|
|
|
cap.release()
|
|
writer.release()
|
|
|
|
with (args.output / "fusion.jsonl").open("w", encoding="utf-8") as handle:
|
|
for record in output_records:
|
|
handle.write(json.dumps(record, ensure_ascii=False) + "\n")
|
|
|
|
def comparison_records(variant: str) -> list[dict]:
|
|
return [
|
|
{
|
|
"frame_index": record["frame_index"],
|
|
"boxes": record["boxes"],
|
|
"M": {"keypoints": record["raw_keypoints"]},
|
|
"L": {"keypoints": record[f"{variant}_keypoints"]},
|
|
}
|
|
for record in output_records
|
|
]
|
|
|
|
summary = {
|
|
"frames": len(output_records),
|
|
"fusion_policy": {
|
|
"candidates": ["RTMW-L raw", "RTMW-L smaller-person isolated"],
|
|
"selection": [
|
|
"limb and hand selected separately",
|
|
"geometry validity",
|
|
"constant-velocity temporal consistency",
|
|
"source-switch hysteresis",
|
|
"motion-adaptive EMA in person-box coordinates",
|
|
],
|
|
"distance_rule": (
|
|
"Only consensus or temporally stable raw visible endpoints are eligible. "
|
|
"Inferred/conflict/rejected hands are excluded from precise distance."
|
|
),
|
|
"alarm_rule": (
|
|
"Alarm attack is asymmetric: a geometrically valid raw endpoint whose score "
|
|
"is clearly stronger than the isolated view is admitted immediately. Alarm "
|
|
"release should still require temporal confirmation and hysteresis."
|
|
),
|
|
"risk_distance_note": (
|
|
"Use global_hands and temporal_selected_keypoints (unfiltered) for alarm entry "
|
|
"to avoid smoothing delay. Use distance_eligible hands and temporal_keypoints "
|
|
"for stable reporting. When the door edge is available, a conservative "
|
|
"implementation can use the smaller of the unfiltered and smoothed distances."
|
|
),
|
|
},
|
|
"spatial": metadata_summary(output_records, "spatial"),
|
|
"temporal": metadata_summary(output_records, "temporal"),
|
|
"active_hand_frame_step": {
|
|
"spatial": active_hand_temporal_metrics(output_records, "spatial"),
|
|
"temporal": active_hand_temporal_metrics(output_records, "temporal"),
|
|
},
|
|
"raw_vs_spatial_temporal_metrics": temporal_metrics(comparison_records("spatial")),
|
|
"raw_vs_temporal_temporal_metrics": temporal_metrics(comparison_records("temporal")),
|
|
"global_hand_count_frames": dict(
|
|
sorted(Counter(len(record["global_hands"]) for record in output_records).items())
|
|
),
|
|
"cross_person_hand_aliases_total": sum(
|
|
len(hand["aliases"])
|
|
for record in output_records
|
|
for hand in record["global_hands"]
|
|
),
|
|
}
|
|
with (args.output / "summary.json").open("w", encoding="utf-8") as handle:
|
|
json.dump(summary, handle, ensure_ascii=False, indent=2)
|
|
|
|
key_frames = [0, 45, 50, 55, 120, 132, 136, 180, 240, 252, 295, 300, 304, 360, 411]
|
|
make_montage(
|
|
args.video,
|
|
output_records,
|
|
key_frames,
|
|
args.output / "rtmw_l_raw_spatial_temporal_montage.jpg",
|
|
)
|
|
print(json.dumps(summary, ensure_ascii=False, indent=2))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|