"""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()