1131 lines
44 KiB
Python
1131 lines
44 KiB
Python
"""Occlusion-aware temporal pose reconstruction on top of RTMW-L fusion.
|
|
|
|
Visible endpoints are filtered and optionally connected with a two-bone IK
|
|
solution. Short gaps are interpolated/predicted, medium gaps use a local motion
|
|
corridor, and long occlusions retain an undrawn shoulder sector for downstream
|
|
door-band intersection instead of emitting a fake precise wrist.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import math
|
|
from collections import Counter, defaultdict
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
from compare_rtmw_ml import HAND_EDGES_LOCAL, HAND_SLICES, normalized_keypoints, summarize
|
|
from fuse_rtmw_l_temporal import BODY_SIDE_INDICES, draw_dashed_line, draw_fused
|
|
|
|
|
|
STATE_COLORS = {
|
|
"observed": (40, 220, 40),
|
|
"provisional": (255, 190, 40),
|
|
"recovered": (220, 180, 40),
|
|
"reassigned": (220, 120, 40),
|
|
"interpolated": (0, 220, 255),
|
|
"predicted_short": (0, 140, 255),
|
|
"reachable_medium": (0, 120, 255),
|
|
"reachable_long": (30, 80, 255),
|
|
}
|
|
|
|
|
|
@dataclass
|
|
class Measurement:
|
|
keypoints: np.ndarray
|
|
precise: bool
|
|
provisional: bool
|
|
source: str = "model"
|
|
quality: float = 0.0
|
|
|
|
|
|
@dataclass
|
|
class EndpointTrack:
|
|
wrist: np.ndarray | None = None
|
|
velocity: np.ndarray | None = None
|
|
hand_shape: np.ndarray | None = None
|
|
last_observed_frame: int | None = None
|
|
missed_frames: int = 0
|
|
last_elbow: np.ndarray | None = None
|
|
last_direction: np.ndarray | None = None
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--video", type=Path, required=True)
|
|
parser.add_argument("--fusion-jsonl", type=Path, required=True)
|
|
parser.add_argument("--output", type=Path, required=True)
|
|
parser.add_argument("--short-occlusion-frames", type=int, default=12)
|
|
parser.add_argument("--medium-occlusion-frames", type=int, default=30)
|
|
parser.add_argument("--lookahead-frames", type=int, default=3)
|
|
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 point_distance_ratio(a: np.ndarray, b: np.ndarray, bbox: list[float]) -> float:
|
|
return float(np.linalg.norm(a - b)) / bbox_diagonal(bbox)
|
|
|
|
|
|
def collect_measurements(
|
|
records: list[dict],
|
|
person_idx: int,
|
|
side: str,
|
|
lookahead_frames: int,
|
|
) -> list[Measurement | None]:
|
|
start, end = HAND_SLICES[side]
|
|
measurements: list[Measurement | None] = [None] * len(records)
|
|
observed_frames = []
|
|
for frame_idx, record in enumerate(records):
|
|
meta = record["temporal_metadata"][person_idx][side]
|
|
if not meta["alarm_eligible"]:
|
|
continue
|
|
keypoints = np.asarray(
|
|
record["temporal_selected_keypoints"][person_idx][start:end],
|
|
dtype=np.float64,
|
|
)
|
|
measurements[frame_idx] = Measurement(
|
|
keypoints=keypoints,
|
|
precise=bool(meta["distance_eligible"]),
|
|
provisional=not bool(meta["distance_eligible"]),
|
|
quality=float(meta.get("selected_hand_score", 0.0)),
|
|
)
|
|
observed_frames.append(frame_idx)
|
|
|
|
# A small future window is allowed for stable reporting. The alarm fast
|
|
# path still uses the unmodified current-frame candidates from V1.
|
|
for left, right in zip(observed_frames, observed_frames[1:]):
|
|
gap = right - left - 1
|
|
if gap <= 0 or gap > lookahead_frames:
|
|
continue
|
|
left_points = measurements[left].keypoints
|
|
right_points = measurements[right].keypoints
|
|
for offset in range(1, gap + 1):
|
|
ratio = offset / (gap + 1)
|
|
points = (1.0 - ratio) * left_points + ratio * right_points
|
|
measurements[left + offset] = Measurement(
|
|
keypoints=points,
|
|
precise=False,
|
|
provisional=False,
|
|
source="fixed_lag_interpolation",
|
|
quality=min(measurements[left].quality, measurements[right].quality),
|
|
)
|
|
return measurements
|
|
|
|
|
|
def estimate_bone_ratios(records: list[dict]) -> dict[tuple[int, str], tuple[float, float]]:
|
|
values: dict[tuple[int, str], list[tuple[float, float]]] = defaultdict(list)
|
|
for record in records:
|
|
for person_idx, bbox in enumerate(record["boxes"]):
|
|
keypoints = np.asarray(
|
|
record["temporal_selected_keypoints"][person_idx], dtype=np.float64
|
|
)
|
|
diagonal = bbox_diagonal(bbox)
|
|
for side, (shoulder, elbow, wrist) in BODY_SIDE_INDICES.items():
|
|
meta = record["temporal_metadata"][person_idx][side]
|
|
if meta["limb_state"] != "consensus":
|
|
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:
|
|
values[(person_idx, side)].append((upper, forearm))
|
|
|
|
estimates: dict[tuple[int, str], tuple[float, float]] = {}
|
|
person_count = len(records[0]["boxes"])
|
|
for person_idx in range(person_count):
|
|
for side in ("left", "right"):
|
|
samples = values[(person_idx, side)]
|
|
if len(samples) < 20:
|
|
opposite = "right" if side == "left" else "left"
|
|
samples = values[(person_idx, opposite)]
|
|
if samples:
|
|
array = np.asarray(samples)
|
|
estimates[(person_idx, side)] = (
|
|
float(np.median(array[:, 0])),
|
|
float(np.median(array[:, 1])),
|
|
)
|
|
else:
|
|
estimates[(person_idx, side)] = (0.24, 0.22)
|
|
return estimates
|
|
|
|
|
|
def solve_two_bone_elbow(
|
|
shoulder: np.ndarray,
|
|
wrist: np.ndarray,
|
|
upper_length: float,
|
|
forearm_length: float,
|
|
reference_elbow: np.ndarray,
|
|
) -> np.ndarray:
|
|
direction = wrist - shoulder
|
|
distance = float(np.linalg.norm(direction))
|
|
if distance < 1e-6:
|
|
return reference_elbow.copy()
|
|
unit = direction / distance
|
|
|
|
if distance >= upper_length + forearm_length:
|
|
return shoulder + unit * upper_length
|
|
if distance <= abs(upper_length - forearm_length):
|
|
sign = 1.0 if upper_length >= forearm_length else -1.0
|
|
return shoulder + unit * upper_length * sign
|
|
|
|
along = (
|
|
upper_length * upper_length
|
|
- forearm_length * forearm_length
|
|
+ distance * distance
|
|
) / (2.0 * distance)
|
|
height = math.sqrt(max(upper_length * upper_length - along * along, 0.0))
|
|
base = shoulder + unit * along
|
|
perpendicular = np.asarray([-unit[1], unit[0]])
|
|
candidates = (base + perpendicular * height, base - perpendicular * height)
|
|
return min(candidates, key=lambda candidate: float(np.linalg.norm(candidate - reference_elbow)))
|
|
|
|
|
|
def update_endpoint_track(
|
|
track: EndpointTrack,
|
|
measurement: Measurement | None,
|
|
frame_idx: int,
|
|
short_occlusion_frames: int,
|
|
medium_occlusion_frames: int,
|
|
) -> tuple[np.ndarray | None, np.ndarray | None, str]:
|
|
if measurement is not None:
|
|
measured_wrist = measurement.keypoints[0]
|
|
measured_shape = measurement.keypoints - measured_wrist
|
|
if track.wrist is None or track.missed_frames > short_occlusion_frames:
|
|
track.wrist = measured_wrist.copy()
|
|
track.velocity = np.zeros(2, dtype=np.float64)
|
|
track.hand_shape = measured_shape.copy()
|
|
else:
|
|
prediction = track.wrist + track.velocity
|
|
residual = measured_wrist - prediction
|
|
alpha = 0.78 if measurement.precise else 0.88
|
|
beta = 0.18 if measurement.precise else 0.24
|
|
track.wrist = prediction + alpha * residual
|
|
track.velocity = 0.82 * track.velocity + beta * residual
|
|
shape_alpha = 0.55 if measurement.precise else 0.7
|
|
track.hand_shape = (
|
|
shape_alpha * measured_shape + (1.0 - shape_alpha) * track.hand_shape
|
|
)
|
|
track.last_observed_frame = frame_idx
|
|
track.missed_frames = 0
|
|
if measurement.source == "global_reassignment":
|
|
state = "reassigned"
|
|
elif measurement.source == "temporal_recovery":
|
|
state = "recovered"
|
|
elif measurement.precise:
|
|
state = "observed"
|
|
elif measurement.provisional:
|
|
state = "provisional"
|
|
else:
|
|
state = "interpolated"
|
|
return track.wrist.copy(), track.wrist + track.hand_shape, state
|
|
|
|
track.missed_frames += 1
|
|
if track.wrist is None:
|
|
return None, None, "reachable_long"
|
|
if track.missed_frames <= short_occlusion_frames:
|
|
track.wrist = track.wrist + track.velocity
|
|
track.velocity *= 0.86
|
|
return track.wrist.copy(), track.wrist + track.hand_shape, "predicted_short"
|
|
if track.missed_frames <= medium_occlusion_frames:
|
|
return None, None, "reachable_medium"
|
|
return None, None, "reachable_long"
|
|
|
|
|
|
def temporal_recovery_measurement(
|
|
keypoints: np.ndarray,
|
|
metadata: dict,
|
|
side: str,
|
|
bbox: list[float],
|
|
track: EndpointTrack,
|
|
medium_occlusion_frames: int,
|
|
) -> Measurement | None:
|
|
"""Recover a weak current-frame hand only when history strongly supports it."""
|
|
if track.wrist is None or track.missed_frames > medium_occlusion_frames:
|
|
return None
|
|
if metadata.get("hand_state") not in {"conflict", "inferred"}:
|
|
return None
|
|
if not metadata.get("selected_hand_valid", False):
|
|
return None
|
|
if float(metadata.get("selected_hand_score", 0.0)) < 2.0:
|
|
return None
|
|
link = float(metadata.get("hand_selection", {}).get("link_to_selected_limb", 1.0))
|
|
if link > 0.08:
|
|
return None
|
|
|
|
start, end = HAND_SLICES[side]
|
|
hand = np.asarray(keypoints[start:end], dtype=np.float64)
|
|
if hand.shape != (end - start, 2) or not np.isfinite(hand).all():
|
|
return None
|
|
prediction = track.wrist + track.velocity
|
|
diagonal = bbox_diagonal(bbox)
|
|
gate = min(0.20, 0.05 + 0.008 * max(track.missed_frames, 1)) * diagonal
|
|
if float(np.linalg.norm(hand[0] - prediction)) > gate:
|
|
return None
|
|
return Measurement(
|
|
keypoints=hand,
|
|
precise=False,
|
|
provisional=True,
|
|
source="temporal_recovery",
|
|
quality=float(metadata.get("selected_hand_score", 0.0)),
|
|
)
|
|
|
|
|
|
def hand_span_px(hand: np.ndarray) -> float:
|
|
return math.hypot(float(np.ptp(hand[:, 0])), float(np.ptp(hand[:, 1])))
|
|
|
|
|
|
def bbox_area(bbox: list[float]) -> float:
|
|
return max((bbox[2] - bbox[0]) * (bbox[3] - bbox[1]), 1.0)
|
|
|
|
|
|
def reset_endpoint_track(track: EndpointTrack) -> None:
|
|
track.wrist = None
|
|
track.velocity = None
|
|
track.hand_shape = None
|
|
track.last_observed_frame = None
|
|
track.missed_frames = 0
|
|
track.last_elbow = None
|
|
track.last_direction = None
|
|
|
|
|
|
def resolve_global_hand_assignments(
|
|
frame_measurements: dict[tuple[int, str], Measurement | None],
|
|
tracks: dict[tuple[int, str], EndpointTrack],
|
|
boxes: list[list[float]],
|
|
medium_occlusion_frames: int,
|
|
) -> tuple[dict[tuple[int, str], Measurement | None], dict[tuple[int, str], dict]]:
|
|
"""Resolve overlapping-person hands before binding them to body sides.
|
|
|
|
Large foreground top-down crops may include a smaller background person.
|
|
The same physical hand can therefore be emitted for both people, or only by
|
|
the smaller crop even though it belongs to the foreground person. Direct
|
|
duplicates are kept on the foreground track when their apparent scales
|
|
agree. A missing foreground side can then reclaim a nearby candidate using
|
|
its recent hand trajectory and historical hand scale.
|
|
"""
|
|
resolved = dict(frame_measurements)
|
|
events: dict[tuple[int, str], dict] = {}
|
|
keys = list(resolved)
|
|
|
|
# First collapse direct cross-person duplicates. Keep this deliberately
|
|
# asymmetric: a substantially larger foreground crop must also carry the
|
|
# substantially larger historical hand shape.
|
|
for index, first_key in enumerate(keys):
|
|
first = resolved[first_key]
|
|
if first is None:
|
|
continue
|
|
for second_key in keys[index + 1 :]:
|
|
second = resolved[second_key]
|
|
if second is None or first_key[0] == second_key[0]:
|
|
continue
|
|
first_hand = first.keypoints
|
|
second_hand = second.keypoints
|
|
first_span = hand_span_px(first_hand)
|
|
second_span = hand_span_px(second_hand)
|
|
center_distance = float(
|
|
np.linalg.norm(np.mean(first_hand, axis=0) - np.mean(second_hand, axis=0))
|
|
)
|
|
if center_distance > max(70.0, 0.5 * max(first_span, second_span)):
|
|
continue
|
|
|
|
first_area = bbox_area(boxes[first_key[0]])
|
|
second_area = bbox_area(boxes[second_key[0]])
|
|
if first_area >= 2.5 * second_area and (
|
|
first_span >= 1.45 * second_span
|
|
or first.quality >= second.quality + 0.6
|
|
):
|
|
owner_key, released_key = first_key, second_key
|
|
elif second_area >= 2.5 * first_area and (
|
|
second_span >= 1.45 * first_span
|
|
or second.quality >= first.quality + 0.6
|
|
):
|
|
owner_key, released_key = second_key, first_key
|
|
else:
|
|
continue
|
|
resolved[released_key] = None
|
|
reset_endpoint_track(tracks[released_key])
|
|
events[owner_key] = {
|
|
"method": "foreground_duplicate_owner",
|
|
"released_person_id": released_key[0] + 1,
|
|
"released_side": released_key[1],
|
|
}
|
|
events[released_key] = {
|
|
"method": "reassigned_away",
|
|
"assigned_person_id": owner_key[0] + 1,
|
|
"assigned_side": owner_key[1],
|
|
}
|
|
|
|
# Then recover missing foreground hands from another crop. The candidate
|
|
# must fall inside the target's recent motion gate and be much smaller than
|
|
# the target's historical foreground hand shape.
|
|
available_sources = [key for key, value in resolved.items() if value is not None]
|
|
target_candidates = []
|
|
for target_key, target_measurement in resolved.items():
|
|
if target_measurement is not None:
|
|
continue
|
|
target_track = tracks[target_key]
|
|
if (
|
|
target_track.wrist is None
|
|
or target_track.hand_shape is None
|
|
or target_track.missed_frames > medium_occlusion_frames
|
|
):
|
|
continue
|
|
target_area = bbox_area(boxes[target_key[0]])
|
|
predicted_center = target_track.wrist + np.mean(target_track.hand_shape, axis=0)
|
|
target_span = hand_span_px(target_track.hand_shape)
|
|
gate = 0.20 * bbox_diagonal(boxes[target_key[0]])
|
|
for source_key in available_sources:
|
|
if source_key[0] == target_key[0]:
|
|
continue
|
|
source = resolved[source_key]
|
|
if source is None or source.source != "model":
|
|
continue
|
|
source_area = bbox_area(boxes[source_key[0]])
|
|
if target_area < 2.5 * source_area:
|
|
continue
|
|
source_center = np.mean(source.keypoints, axis=0)
|
|
source_span = hand_span_px(source.keypoints)
|
|
distance = float(np.linalg.norm(source_center - predicted_center))
|
|
if distance > gate or target_span < 1.45 * max(source_span, 1.0):
|
|
continue
|
|
target_candidates.append(
|
|
(distance / gate, target_key, source_key, source_center)
|
|
)
|
|
|
|
used_targets = set()
|
|
used_sources = set()
|
|
for _, target_key, source_key, source_center in sorted(target_candidates):
|
|
if target_key in used_targets or source_key in used_sources:
|
|
continue
|
|
target_track = tracks[target_key]
|
|
shape_center = np.mean(target_track.hand_shape, axis=0)
|
|
reassigned_wrist = source_center - shape_center
|
|
reassigned_hand = reassigned_wrist + target_track.hand_shape
|
|
resolved[target_key] = Measurement(
|
|
keypoints=reassigned_hand,
|
|
precise=False,
|
|
provisional=True,
|
|
source="global_reassignment",
|
|
quality=resolved[source_key].quality,
|
|
)
|
|
resolved[source_key] = None
|
|
reset_endpoint_track(tracks[source_key])
|
|
events[target_key] = {
|
|
"method": "foreground_temporal_reassignment",
|
|
"source_person_id": source_key[0] + 1,
|
|
"source_side": source_key[1],
|
|
"observed_center": np.round(source_center, 3).tolist(),
|
|
}
|
|
events[source_key] = {
|
|
"method": "reassigned_away",
|
|
"assigned_person_id": target_key[0] + 1,
|
|
"assigned_side": target_key[1],
|
|
}
|
|
used_targets.add(target_key)
|
|
used_sources.add(source_key)
|
|
return resolved, events
|
|
|
|
|
|
def track_direction(track: EndpointTrack) -> np.ndarray:
|
|
if track.last_direction is not None:
|
|
return track.last_direction / max(float(np.linalg.norm(track.last_direction)), 1e-6)
|
|
if track.velocity is not None and float(np.linalg.norm(track.velocity)) > 1e-6:
|
|
return track.velocity / np.linalg.norm(track.velocity)
|
|
return np.asarray([1.0, 0.0], dtype=np.float64)
|
|
|
|
|
|
def ellipse_polygon(
|
|
center: np.ndarray,
|
|
major_axis: float,
|
|
minor_axis: float,
|
|
direction: np.ndarray,
|
|
point_count: int = 64,
|
|
) -> np.ndarray:
|
|
perpendicular = np.asarray([-direction[1], direction[0]], dtype=np.float64)
|
|
angles = np.linspace(0.0, 2.0 * math.pi, point_count, endpoint=False)
|
|
return (
|
|
center
|
|
+ np.cos(angles)[:, None] * major_axis * direction
|
|
+ np.sin(angles)[:, None] * minor_axis * perpendicular
|
|
)
|
|
|
|
|
|
def clip_polygon_from_anchor_to_circle(
|
|
points: np.ndarray,
|
|
anchor: np.ndarray,
|
|
circle_center: np.ndarray,
|
|
radius: float,
|
|
) -> np.ndarray:
|
|
"""Clip radial polygon samples to a shoulder-centered anatomical limit."""
|
|
clipped = []
|
|
offset = anchor - circle_center
|
|
for point in points:
|
|
if float(np.linalg.norm(point - circle_center)) <= radius:
|
|
clipped.append(point)
|
|
continue
|
|
ray = point - anchor
|
|
a = float(np.dot(ray, ray))
|
|
b = 2.0 * float(np.dot(offset, ray))
|
|
c = float(np.dot(offset, offset) - radius * radius)
|
|
discriminant = max(b * b - 4.0 * a * c, 0.0)
|
|
roots = [
|
|
value
|
|
for value in (
|
|
(-b - math.sqrt(discriminant)) / max(2.0 * a, 1e-6),
|
|
(-b + math.sqrt(discriminant)) / max(2.0 * a, 1e-6),
|
|
)
|
|
if 0.0 <= value <= 1.0
|
|
]
|
|
clipped.append(anchor + (max(roots) if roots else 0.0) * ray)
|
|
return np.asarray(clipped, dtype=np.float64)
|
|
|
|
|
|
def medium_motion_region(
|
|
shoulder: np.ndarray,
|
|
upper_length: float,
|
|
forearm_length: float,
|
|
bbox: list[float],
|
|
track: EndpointTrack,
|
|
short_occlusion_frames: int,
|
|
) -> dict:
|
|
diagonal = bbox_diagonal(bbox)
|
|
direction = track_direction(track)
|
|
hand_shape = (
|
|
track.hand_shape
|
|
if track.hand_shape is not None
|
|
else np.zeros((21, 2), dtype=np.float64)
|
|
)
|
|
shape_center = np.mean(hand_shape, axis=0)
|
|
centered_shape = hand_shape - shape_center
|
|
perpendicular = np.asarray([-direction[1], direction[0]], dtype=np.float64)
|
|
hand_major = float(np.max(np.abs(centered_shape @ direction)))
|
|
hand_minor = float(np.max(np.abs(centered_shape @ perpendicular)))
|
|
|
|
extra_frames = max(track.missed_frames - short_occlusion_frames, 1)
|
|
motion_major = min(0.20, 0.04 + 0.007 * extra_frames) * diagonal
|
|
motion_minor = min(0.13, 0.03 + 0.0045 * extra_frames) * diagonal
|
|
major_axis = max(hand_major + motion_major, 0.08 * diagonal)
|
|
minor_axis = max(hand_minor + motion_minor, 0.06 * diagonal)
|
|
|
|
center = track.wrist + shape_center
|
|
tracked_hand_extent = float(np.max(np.linalg.norm(centered_shape, axis=1)))
|
|
hand_reach = max(tracked_hand_extent, 0.10 * diagonal)
|
|
anatomical_reach = upper_length + forearm_length + hand_reach + 0.02 * diagonal
|
|
shoulder_to_center = center - shoulder
|
|
center_distance = float(np.linalg.norm(shoulder_to_center))
|
|
if center_distance > 0.96 * anatomical_reach:
|
|
center = shoulder + shoulder_to_center / center_distance * 0.96 * anatomical_reach
|
|
|
|
polygon = ellipse_polygon(center, major_axis, minor_axis, direction)
|
|
polygon = clip_polygon_from_anchor_to_circle(
|
|
polygon,
|
|
center,
|
|
shoulder,
|
|
anatomical_reach,
|
|
)
|
|
direction_deg = math.degrees(math.atan2(float(direction[1]), float(direction[0])))
|
|
return {
|
|
"geometry": "motion_corridor",
|
|
"center": np.round(center, 3).tolist(),
|
|
"predicted_wrist": np.round(track.wrist, 3).tolist(),
|
|
"major_axis_px": round(major_axis, 3),
|
|
"minor_axis_px": round(minor_axis, 3),
|
|
"outer_radius_px": round(major_axis, 3),
|
|
"direction_deg": round(direction_deg, 3),
|
|
"max_reach_from_shoulder_px": round(anatomical_reach, 3),
|
|
"polygon": np.round(polygon, 3).tolist(),
|
|
"requires_door_intersection": True,
|
|
"render_full_region": True,
|
|
"distance_lower_bound_rule": "minimum distance from motion-corridor polygon to door_edge",
|
|
}
|
|
|
|
|
|
def long_reachable_region(
|
|
shoulder: np.ndarray,
|
|
upper_length: float,
|
|
forearm_length: float,
|
|
bbox: list[float],
|
|
track: EndpointTrack,
|
|
medium_occlusion_frames: int,
|
|
) -> dict:
|
|
diagonal = bbox_diagonal(bbox)
|
|
height = max(bbox[3] - bbox[1], 1.0)
|
|
arm_reach = upper_length + forearm_length
|
|
hand_reach = min(0.13 * diagonal, 0.17 * height)
|
|
safety_margin = 0.025 * diagonal
|
|
outer_radius = arm_reach + hand_reach + safety_margin
|
|
inner_radius = max(abs(upper_length - forearm_length) - 0.04 * diagonal, 0.0)
|
|
if track.last_direction is None:
|
|
direction_deg = 0.0
|
|
half_angle_deg = 180.0
|
|
else:
|
|
direction = track_direction(track)
|
|
direction_deg = math.degrees(math.atan2(float(direction[1]), float(direction[0])))
|
|
half_angle_deg = min(
|
|
180.0,
|
|
65.0 + 2.0 * max(track.missed_frames - medium_occlusion_frames, 0),
|
|
)
|
|
return {
|
|
"geometry": "shoulder_sector",
|
|
"center": np.round(shoulder, 3).tolist(),
|
|
"inner_radius_px": round(inner_radius, 3),
|
|
"outer_radius_px": round(outer_radius, 3),
|
|
"arm_reach_px": round(arm_reach, 3),
|
|
"hand_reach_px": round(hand_reach, 3),
|
|
"safety_margin_px": round(safety_margin, 3),
|
|
"direction_deg": round(direction_deg, 3),
|
|
"half_angle_deg": round(half_angle_deg, 3),
|
|
"requires_door_intersection": True,
|
|
"render_full_region": False,
|
|
"distance_lower_bound_rule": "intersect shoulder sector with door danger band before alarming",
|
|
}
|
|
|
|
|
|
def reachable_region(
|
|
shoulder: np.ndarray,
|
|
upper_length: float,
|
|
forearm_length: float,
|
|
bbox: list[float],
|
|
track: EndpointTrack,
|
|
state: str = "reachable_long",
|
|
short_occlusion_frames: int = 12,
|
|
medium_occlusion_frames: int = 30,
|
|
) -> dict:
|
|
if state == "reachable_medium" and track.wrist is not None:
|
|
return medium_motion_region(
|
|
shoulder,
|
|
upper_length,
|
|
forearm_length,
|
|
bbox,
|
|
track,
|
|
short_occlusion_frames,
|
|
)
|
|
return long_reachable_region(
|
|
shoulder,
|
|
upper_length,
|
|
forearm_length,
|
|
bbox,
|
|
track,
|
|
medium_occlusion_frames,
|
|
)
|
|
|
|
|
|
def calibrated_bone_lengths(
|
|
upper_ratio: float,
|
|
forearm_ratio: float,
|
|
bbox: list[float],
|
|
) -> tuple[float, float]:
|
|
diagonal = bbox_diagonal(bbox)
|
|
height = max(bbox[3] - bbox[1], 1.0)
|
|
upper = upper_ratio * diagonal
|
|
forearm = forearm_ratio * diagonal
|
|
total = upper + forearm
|
|
# Keep only a loose sanity cap. The previous 0.55 * bbox-height cap was
|
|
# unsafe for close, clipped CMS views: a valid projected arm can occupy
|
|
# almost the full detection height. The ratios are already robust medians,
|
|
# so this cap only limits catastrophic calibration outliers.
|
|
maximum_total = 0.95 * diagonal
|
|
if total > maximum_total:
|
|
scale = maximum_total / total
|
|
upper *= scale
|
|
forearm *= scale
|
|
return upper, forearm
|
|
|
|
|
|
def process_sequence(
|
|
records: list[dict],
|
|
short_occlusion_frames: int,
|
|
medium_occlusion_frames: int,
|
|
lookahead_frames: int,
|
|
) -> list[dict]:
|
|
person_count = len(records[0]["boxes"])
|
|
measurements = {
|
|
(person_idx, side): collect_measurements(
|
|
records, person_idx, side, lookahead_frames
|
|
)
|
|
for person_idx in range(person_count)
|
|
for side in ("left", "right")
|
|
}
|
|
bone_ratios = estimate_bone_ratios(records)
|
|
tracks = {
|
|
(person_idx, side): EndpointTrack()
|
|
for person_idx in range(person_count)
|
|
for side in ("left", "right")
|
|
}
|
|
|
|
output = []
|
|
for frame_idx, record in enumerate(records):
|
|
boxes = record["boxes"]
|
|
base_keypoints = np.asarray(record["temporal_keypoints"], dtype=np.float64)
|
|
unfiltered_keypoints = np.asarray(
|
|
record["temporal_selected_keypoints"], dtype=np.float64
|
|
)
|
|
v2_keypoints = base_keypoints.copy()
|
|
frame_measurements = {
|
|
(person_idx, side): measurements[(person_idx, side)][frame_idx]
|
|
for person_idx in range(person_count)
|
|
for side in ("left", "right")
|
|
}
|
|
frame_measurements, assignment_events = resolve_global_hand_assignments(
|
|
frame_measurements,
|
|
tracks,
|
|
boxes,
|
|
medium_occlusion_frames,
|
|
)
|
|
people_meta = []
|
|
for person_idx, bbox in enumerate(boxes):
|
|
diagonal = bbox_diagonal(bbox)
|
|
person_meta = {}
|
|
for side, (shoulder_idx, elbow_idx, wrist_idx) in BODY_SIDE_INDICES.items():
|
|
track = tracks[(person_idx, side)]
|
|
measurement = frame_measurements[(person_idx, side)]
|
|
if measurement is None:
|
|
measurement = temporal_recovery_measurement(
|
|
unfiltered_keypoints[person_idx],
|
|
record["temporal_metadata"][person_idx][side],
|
|
side,
|
|
bbox,
|
|
track,
|
|
medium_occlusion_frames,
|
|
)
|
|
returning_after_long_occlusion = (
|
|
measurement is not None
|
|
and track.missed_frames > short_occlusion_frames
|
|
)
|
|
wrist, hand, state = update_endpoint_track(
|
|
track,
|
|
measurement,
|
|
frame_idx,
|
|
short_occlusion_frames,
|
|
medium_occlusion_frames,
|
|
)
|
|
shoulder = base_keypoints[person_idx, shoulder_idx]
|
|
upper_ratio, forearm_ratio = bone_ratios[(person_idx, side)]
|
|
upper_length, forearm_length = calibrated_bone_lengths(
|
|
upper_ratio, forearm_ratio, bbox
|
|
)
|
|
reachable = None
|
|
|
|
if wrist is not None and hand is not None:
|
|
reference_elbow = (
|
|
unfiltered_keypoints[person_idx, elbow_idx]
|
|
if returning_after_long_occlusion or track.last_elbow is None
|
|
else track.last_elbow
|
|
)
|
|
use_ik = (
|
|
record["temporal_metadata"][person_idx][side]["limb_state"]
|
|
!= "consensus"
|
|
or state in (
|
|
"provisional",
|
|
"recovered",
|
|
"reassigned",
|
|
"interpolated",
|
|
"predicted_short",
|
|
)
|
|
)
|
|
if use_ik:
|
|
elbow = solve_two_bone_elbow(
|
|
shoulder,
|
|
wrist,
|
|
upper_length,
|
|
forearm_length,
|
|
reference_elbow,
|
|
)
|
|
else:
|
|
elbow = base_keypoints[person_idx, elbow_idx]
|
|
track.last_elbow = elbow.copy()
|
|
direction = wrist - shoulder
|
|
if float(np.linalg.norm(direction)) > 1e-6:
|
|
track.last_direction = direction / np.linalg.norm(direction)
|
|
|
|
v2_keypoints[person_idx, shoulder_idx] = shoulder
|
|
v2_keypoints[person_idx, elbow_idx] = elbow
|
|
v2_keypoints[person_idx, wrist_idx] = wrist
|
|
start, end = HAND_SLICES[side]
|
|
v2_keypoints[person_idx, start:end] = hand
|
|
uncertainty_radius = (
|
|
0.025 * diagonal
|
|
if state == "observed"
|
|
else (0.04 + 0.015 * track.missed_frames) * diagonal
|
|
)
|
|
precise_distance = state == "observed"
|
|
fast_alarm = state in (
|
|
"observed",
|
|
"provisional",
|
|
"reassigned",
|
|
"interpolated",
|
|
"predicted_short",
|
|
)
|
|
else:
|
|
reachable = reachable_region(
|
|
shoulder,
|
|
upper_length,
|
|
forearm_length,
|
|
bbox,
|
|
track,
|
|
state,
|
|
short_occlusion_frames,
|
|
medium_occlusion_frames,
|
|
)
|
|
uncertainty_radius = reachable["outer_radius_px"]
|
|
precise_distance = False
|
|
fast_alarm = True
|
|
|
|
person_meta[side] = {
|
|
"state": state,
|
|
"precise_distance": precise_distance,
|
|
"fast_alarm_candidate": fast_alarm,
|
|
"risk_mode": (
|
|
"precise_hand"
|
|
if state == "observed"
|
|
else (
|
|
"motion_corridor_intersection"
|
|
if state == "reachable_medium"
|
|
else (
|
|
"door_intersection_only"
|
|
if state == "reachable_long"
|
|
else (
|
|
"reassigned_hand"
|
|
if state == "reassigned"
|
|
else "temporal_hand"
|
|
)
|
|
)
|
|
)
|
|
),
|
|
"requires_door_intersection": state
|
|
in ("reachable_medium", "reachable_long"),
|
|
"missed_frames": track.missed_frames,
|
|
"uncertainty_radius_px": round(float(uncertainty_radius), 3),
|
|
"bone_lengths_px": {
|
|
"upper": round(upper_length, 3),
|
|
"forearm": round(forearm_length, 3),
|
|
},
|
|
"identity_assignment": assignment_events.get(
|
|
(person_idx, side)
|
|
),
|
|
"reachable_region": reachable,
|
|
}
|
|
people_meta.append(person_meta)
|
|
|
|
output.append(
|
|
{
|
|
"frame_index": frame_idx,
|
|
"boxes": boxes,
|
|
"v1_keypoints": record["temporal_keypoints"],
|
|
"v1_metadata": record["temporal_metadata"],
|
|
"v2_keypoints": np.round(v2_keypoints, 3).tolist(),
|
|
"v2_metadata": people_meta,
|
|
}
|
|
)
|
|
return output
|
|
|
|
|
|
def sector_points(region: dict) -> np.ndarray:
|
|
center = np.asarray(region["center"], dtype=np.float64)
|
|
direction = region["direction_deg"]
|
|
half_angle = region["half_angle_deg"]
|
|
angles = np.linspace(direction - half_angle, direction + half_angle, 48)
|
|
radius = region["outer_radius_px"]
|
|
arc = np.column_stack(
|
|
[
|
|
center[0] + radius * np.cos(np.deg2rad(angles)),
|
|
center[1] + radius * np.sin(np.deg2rad(angles)),
|
|
]
|
|
)
|
|
if half_angle < 179.0:
|
|
return np.vstack([center, arc, center]).astype(np.int32)
|
|
return arc.astype(np.int32)
|
|
|
|
|
|
def draw_v2(
|
|
image: np.ndarray,
|
|
boxes: list[list[float]],
|
|
keypoints: np.ndarray,
|
|
metadata: list[dict],
|
|
title: str,
|
|
show_reachable: bool = True,
|
|
) -> np.ndarray:
|
|
canvas = image.copy()
|
|
overlay = canvas.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, (shoulder, elbow, wrist) in BODY_SIDE_INDICES.items():
|
|
meta = person_meta[side]
|
|
state = meta["state"]
|
|
color = STATE_COLORS[state]
|
|
if state in ("reachable_medium", "reachable_long"):
|
|
region = meta["reachable_region"]
|
|
if state == "reachable_medium":
|
|
if show_reachable:
|
|
polygon = np.round(region["polygon"]).astype(np.int32)
|
|
cv2.fillPoly(overlay, [polygon], color)
|
|
cv2.polylines(canvas, [polygon], True, color, 3, cv2.LINE_AA)
|
|
center = tuple(np.round(region["center"]).astype(int))
|
|
wrist_point = tuple(
|
|
np.round(region["predicted_wrist"]).astype(int)
|
|
)
|
|
cv2.line(canvas, center, wrist_point, color, 2, cv2.LINE_AA)
|
|
label = f"{side[0].upper()}:corridor"
|
|
else:
|
|
center = tuple(np.round(region["center"]).astype(int))
|
|
cv2.circle(canvas, center, 8, color, 2, cv2.LINE_AA)
|
|
if region["half_angle_deg"] >= 179.0:
|
|
cv2.putText(
|
|
canvas,
|
|
"?",
|
|
(center[0] + 11, center[1] + 9),
|
|
cv2.FONT_HERSHEY_SIMPLEX,
|
|
0.65,
|
|
color,
|
|
2,
|
|
cv2.LINE_AA,
|
|
)
|
|
else:
|
|
direction = math.radians(region["direction_deg"])
|
|
direction_end = (
|
|
int(round(center[0] + 55 * math.cos(direction))),
|
|
int(round(center[1] + 55 * math.sin(direction))),
|
|
)
|
|
cv2.arrowedLine(
|
|
canvas,
|
|
center,
|
|
direction_end,
|
|
color,
|
|
2,
|
|
cv2.LINE_AA,
|
|
tipLength=0.25,
|
|
)
|
|
if show_reachable and region.get("render_full_region", False):
|
|
polygon = sector_points(region)
|
|
if region["half_angle_deg"] < 179.0:
|
|
cv2.fillPoly(overlay, [polygon], color)
|
|
cv2.polylines(
|
|
canvas, [polygon], True, color, 3, cv2.LINE_AA
|
|
)
|
|
else:
|
|
radius = int(round(region["outer_radius_px"]))
|
|
cv2.circle(canvas, center, radius, color, 3, cv2.LINE_AA)
|
|
label = f"{side[0].upper()}:unknown"
|
|
else:
|
|
thickness = 3
|
|
if state in ("reassigned", "recovered"):
|
|
draw_dashed_line(
|
|
canvas,
|
|
kpts[shoulder],
|
|
kpts[elbow],
|
|
color,
|
|
thickness,
|
|
)
|
|
draw_dashed_line(
|
|
canvas,
|
|
kpts[elbow],
|
|
kpts[wrist],
|
|
color,
|
|
thickness,
|
|
)
|
|
else:
|
|
cv2.line(
|
|
canvas,
|
|
tuple(np.round(kpts[shoulder]).astype(int)),
|
|
tuple(np.round(kpts[elbow]).astype(int)),
|
|
color,
|
|
thickness,
|
|
cv2.LINE_AA,
|
|
)
|
|
cv2.line(
|
|
canvas,
|
|
tuple(np.round(kpts[elbow]).astype(int)),
|
|
tuple(np.round(kpts[wrist]).astype(int)),
|
|
color,
|
|
thickness,
|
|
cv2.LINE_AA,
|
|
)
|
|
start, end = HAND_SLICES[side]
|
|
hand = kpts[start:end]
|
|
for a, b in HAND_EDGES_LOCAL:
|
|
cv2.line(
|
|
canvas,
|
|
tuple(np.round(hand[a]).astype(int)),
|
|
tuple(np.round(hand[b]).astype(int)),
|
|
color,
|
|
2,
|
|
cv2.LINE_AA,
|
|
)
|
|
label = f"{side[0].upper()}:{state[:5]}"
|
|
cv2.putText(
|
|
canvas,
|
|
label,
|
|
(x1 + 4, max(45, y1 + (48 if side == "left" else 70))),
|
|
cv2.FONT_HERSHEY_SIMPLEX,
|
|
0.48,
|
|
color,
|
|
1,
|
|
cv2.LINE_AA,
|
|
)
|
|
canvas = cv2.addWeighted(overlay, 0.06, canvas, 0.94, 0.0)
|
|
cv2.rectangle(canvas, (0, 0), (670, 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 endpoint_metrics(records: list[dict]) -> dict:
|
|
states: Counter = Counter()
|
|
assignments: Counter = Counter()
|
|
steps: 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 = np.asarray(record["v2_keypoints"][person_idx], dtype=np.float64)
|
|
for side, (_, _, wrist) in BODY_SIDE_INDICES.items():
|
|
meta = record["v2_metadata"][person_idx][side]
|
|
key_name = f"P{person_idx + 1}_{side}"
|
|
states[f"{key_name}_{meta['state']}"] += 1
|
|
identity_assignment = meta.get("identity_assignment")
|
|
if identity_assignment is not None:
|
|
assignments[identity_assignment["method"]] += 1
|
|
if meta["state"] in ("reachable_medium", "reachable_long"):
|
|
previous.pop((person_idx, side), None)
|
|
continue
|
|
normalized = normalized_keypoints(kpts, bbox)[wrist]
|
|
prior = previous.get((person_idx, side))
|
|
if prior is not None and prior[0] == record["frame_index"] - 1:
|
|
steps[key_name].append(float(np.linalg.norm(normalized - prior[1])))
|
|
previous[(person_idx, side)] = (record["frame_index"], normalized)
|
|
return {
|
|
"state_counts": dict(sorted(states.items())),
|
|
"identity_assignment_counts": dict(sorted(assignments.items())),
|
|
"wrist_frame_step": {name: summarize(items) for name, items in steps.items()},
|
|
}
|
|
|
|
|
|
def write_video_and_montage(
|
|
video_path: Path,
|
|
records: list[dict],
|
|
output_dir: Path,
|
|
) -> None:
|
|
cap = cv2.VideoCapture(str(video_path))
|
|
fps = float(cap.get(cv2.CAP_PROP_FPS))
|
|
writer = cv2.VideoWriter(
|
|
str(output_dir / "rtmw_l_temporal_v1_vs_occlusion_v4.mp4"),
|
|
cv2.VideoWriter_fourcc(*"mp4v"),
|
|
fps,
|
|
(1920, 540),
|
|
)
|
|
key_frames = {0, 50, 120, 126, 132, 138, 180, 211, 220, 240, 252,
|
|
292, 296, 300, 304, 320, 360, 411}
|
|
montage_rows = []
|
|
frame_idx = 0
|
|
while True:
|
|
ok, frame = cap.read()
|
|
if not ok:
|
|
break
|
|
record = records[frame_idx]
|
|
v1 = draw_fused(
|
|
frame,
|
|
record["boxes"],
|
|
[np.asarray(points) for points in record["v1_keypoints"]],
|
|
record["v1_metadata"],
|
|
f"temporal fusion V1 | frame {frame_idx}",
|
|
)
|
|
v2 = draw_v2(
|
|
frame,
|
|
record["boxes"],
|
|
np.asarray(record["v2_keypoints"]),
|
|
record["v2_metadata"],
|
|
f"occlusion V4 global assignment | frame {frame_idx}",
|
|
)
|
|
left = cv2.resize(v1, (960, 540), interpolation=cv2.INTER_AREA)
|
|
right = cv2.resize(v2, (960, 540), interpolation=cv2.INTER_AREA)
|
|
writer.write(np.hstack([left, right]))
|
|
if frame_idx in key_frames:
|
|
montage_rows.append(
|
|
np.hstack(
|
|
[
|
|
cv2.resize(v1, (640, 360), interpolation=cv2.INTER_AREA),
|
|
cv2.resize(v2, (640, 360), interpolation=cv2.INTER_AREA),
|
|
]
|
|
)
|
|
)
|
|
frame_idx += 1
|
|
cap.release()
|
|
writer.release()
|
|
cv2.imwrite(
|
|
str(output_dir / "rtmw_l_temporal_v1_vs_occlusion_v4_montage.jpg"),
|
|
np.vstack(montage_rows),
|
|
)
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
args.output.mkdir(parents=True, exist_ok=True)
|
|
records = [
|
|
json.loads(line)
|
|
for line in args.fusion_jsonl.read_text(encoding="utf-8").splitlines()
|
|
]
|
|
output = process_sequence(
|
|
records,
|
|
short_occlusion_frames=args.short_occlusion_frames,
|
|
medium_occlusion_frames=args.medium_occlusion_frames,
|
|
lookahead_frames=args.lookahead_frames,
|
|
)
|
|
with (args.output / "occlusion_v4.jsonl").open("w", encoding="utf-8") as handle:
|
|
for record in output:
|
|
handle.write(json.dumps(record, ensure_ascii=False) + "\n")
|
|
|
|
summary = {
|
|
"algorithm_version": "occlusion_v4_global_assignment",
|
|
"frames": len(output),
|
|
"short_occlusion_frames": args.short_occlusion_frames,
|
|
"medium_occlusion_frames": args.medium_occlusion_frames,
|
|
"lookahead_frames": args.lookahead_frames,
|
|
"policy": {
|
|
"observed": "filtered visible endpoint; precise distance allowed",
|
|
"provisional": "strong current-frame endpoint; fast alarm only",
|
|
"recovered": "weak current-frame endpoint accepted by temporal gating",
|
|
"reassigned": "global hand candidate transferred using foreground scale and trajectory",
|
|
"interpolated": "fixed-lag short-gap interpolation; fast alarm only",
|
|
"predicted_short": "damped-velocity prediction; uncertainty grows each frame",
|
|
"reachable_medium": "track-centered motion corridor; intersect with door band",
|
|
"reachable_long": "no exact wrist; keep shoulder sector data for door intersection only",
|
|
"ik": "two-bone shoulder/wrist reconstruction using robust visible bone lengths",
|
|
"ownership": "one-to-one global hand reassignment using foreground scale and trajectory",
|
|
},
|
|
"metrics": endpoint_metrics(output),
|
|
}
|
|
with (args.output / "summary.json").open("w", encoding="utf-8") as handle:
|
|
json.dump(summary, handle, ensure_ascii=False, indent=2)
|
|
write_video_and_montage(args.video, output, args.output)
|
|
print(json.dumps(summary, ensure_ascii=False, indent=2))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|