zmv/vision-anti-pinch
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Trim use-can to buildable runtime sources and standalone Orin package

Browse Source
This commit is contained in:
zmv
2026-09-23 19:17:17 +08:00
parent 9aadf9cbda
commit 0f609612a6
108 changed files with 2365 additions and 35388 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pipeline/compare_rtmw_ml.py
-384
View File
@@ -1,384 +0,0 @@
"""Compare RTMW-M 256x192 and RTMW-L 384x288 on cached YOLO boxes.
This is an experiment script. It intentionally keeps the candidate weights
outside the formal models/ directory and does not change deployment code.
"""
from __future__ import annotations
import argparse
import json
import math
import time
from collections import defaultdict
from pathlib import Path
import cv2
import numpy as np
BODY_EDGES = ((5, 7), (7, 9), (6, 8), (8, 10))
HAND_EDGES_LOCAL = (
(0, 1), (1, 2), (2, 3), (3, 4),
(0, 5), (5, 6), (6, 7), (7, 8),
(0, 9), (9, 10), (10, 11), (11, 12),
(0, 13), (13, 14), (14, 15), (15, 16),
(0, 17), (17, 18), (18, 19), (19, 20),
)
HAND_SLICES = {"left": (91, 112), "right": (112, 133)}
BODY_WRISTS = {"left": 9, "right": 10}
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser()
parser.add_argument("--video", type=Path, required=True)
parser.add_argument("--boxes", type=Path, required=True)
parser.add_argument("--model-m", type=Path, required=True)
parser.add_argument("--model-l", type=Path, required=True)
parser.add_argument("--output", type=Path, required=True)
parser.add_argument("--max-people", type=int, default=2)
return parser.parse_args()
def load_boxes(path: Path, max_people: int) -> dict[int, list[list[float]]]:
frames: dict[int, list[list[float]]] = {}
with path.open("r", encoding="utf-8") as handle:
for line in handle:
item = json.loads(line)
boxes = [entry["xyxy"] for entry in item["boxes"]]
frames[int(item["frame_index"])] = sorted(boxes, key=lambda b: b[0])[:max_people]
return frames
def point(pt: np.ndarray) -> tuple[int, int]:
return int(round(float(pt[0]))), int(round(float(pt[1])))
def draw_pose(
image: np.ndarray,
boxes: list[list[float]],
keypoints: np.ndarray,
title: str,
) -> np.ndarray:
canvas = image.copy()
for person_idx, (bbox, kpts) in enumerate(zip(boxes, keypoints), 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.75,
(0, 220, 255),
2,
cv2.LINE_AA,
)
for edge_idx, (a, b) in enumerate(BODY_EDGES):
color = (40, 40, 255) if edge_idx < 2 else (30, 230, 30)
cv2.line(canvas, point(kpts[a]), point(kpts[b]), color, 3, cv2.LINE_AA)
for idx in (5, 6, 7, 8, 9, 10):
color = (40, 40, 255) if idx in (5, 7, 9) else (30, 230, 30)
cv2.circle(canvas, point(kpts[idx]), 5, color, -1, cv2.LINE_AA)
for side, (start, end) in HAND_SLICES.items():
hand = kpts[start:end]
color = (255, 80, 220) if side == "left" else (255, 190, 40)
for a, b in HAND_EDGES_LOCAL:
cv2.line(canvas, point(hand[a]), point(hand[b]), color, 2, cv2.LINE_AA)
for hand_point in hand:
cv2.circle(canvas, point(hand_point), 2, color, -1, cv2.LINE_AA)
cv2.rectangle(canvas, (0, 0), (610, 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 normalized_keypoints(kpts: np.ndarray, bbox: list[float]) -> np.ndarray:
x1, y1, x2, y2 = bbox
width = max(x2 - x1, 1.0)
height = max(y2 - y1, 1.0)
result = kpts.copy().astype(np.float64)
result[:, 0] = (result[:, 0] - x1) / width
result[:, 1] = (result[:, 1] - y1) / height
return result
def hand_span(kpts: np.ndarray, bbox: list[float], side: str) -> float:
start, end = HAND_SLICES[side]
hand = kpts[start:end]
span = math.hypot(
float(np.ptp(hand[:, 0])),
float(np.ptp(hand[:, 1])),
)
x1, y1, x2, y2 = bbox
return span / max(math.hypot(x2 - x1, y2 - y1), 1.0)
def summarize(values: list[float]) -> dict[str, float | int | None]:
if not values:
return {"count": 0, "median": None, "p95": None, "max": None}
array = np.asarray(values, dtype=np.float64)
return {
"count": int(array.size),
"median": round(float(np.median(array)), 5),
"p95": round(float(np.percentile(array, 95)), 5),
"max": round(float(np.max(array)), 5),
}
def temporal_metrics(records: list[dict]) -> dict:
by_model_person: dict[tuple[str, int], list[dict]] = defaultdict(list)
for record in records:
for model_name in ("M", "L"):
for person_idx in range(len(record[model_name]["keypoints"])):
by_model_person[(model_name, person_idx)].append(
{
"frame": record["frame_index"],
"bbox": record["boxes"][person_idx],
"kpts": np.asarray(
record[model_name]["keypoints"][person_idx], dtype=np.float64
),
}
)
result: dict[str, dict] = {"M": {}, "L": {}}
for (model_name, person_idx), sequence in by_model_person.items():
upper_steps: list[float] = []
wrist_steps = {"left": [], "right": []}
hand_steps = {"left": [], "right": []}
hand_spans = {"left": [], "right": []}
bone_lengths = {"left_upper": [], "left_forearm": [],
"right_upper": [], "right_forearm": []}
previous: np.ndarray | None = None
for entry in sequence:
normalized = normalized_keypoints(entry["kpts"], entry["bbox"])
if previous is not None:
steps = np.linalg.norm(normalized - previous, axis=1)
upper_steps.extend(float(steps[idx]) for idx in (5, 6, 7, 8, 9, 10))
for side, idx in BODY_WRISTS.items():
wrist_steps[side].append(float(steps[idx]))
for side, (start, end) in HAND_SLICES.items():
hand_steps[side].extend(float(value) for value in steps[start:end])
previous = normalized
hand_spans["left"].append(hand_span(entry["kpts"], entry["bbox"], "left"))
hand_spans["right"].append(hand_span(entry["kpts"], entry["bbox"], "right"))
bone_lengths["left_upper"].append(float(np.linalg.norm(normalized[5] - normalized[7])))
bone_lengths["left_forearm"].append(float(np.linalg.norm(normalized[7] - normalized[9])))
bone_lengths["right_upper"].append(float(np.linalg.norm(normalized[6] - normalized[8])))
bone_lengths["right_forearm"].append(float(np.linalg.norm(normalized[8] - normalized[10])))
span_summary = {}
for side, values in hand_spans.items():
median = float(np.median(values))
gross_limit = max(0.2, median * 2.5)
span_summary[side] = {
**summarize(values),
"gross_limit": round(gross_limit, 5),
"gross_outlier_frames": int(sum(value > gross_limit for value in values)),
}
bone_summary = {}
for name, values in bone_lengths.items():
array = np.asarray(values, dtype=np.float64)
mean = float(np.mean(array))
bone_summary[name] = {
**summarize(values),
"coefficient_of_variation": round(float(np.std(array) / max(mean, 1e-9)), 5),
}
result[model_name][f"P{person_idx + 1}"] = {
"upper_body_frame_step": summarize(upper_steps),
"wrist_frame_step": {side: summarize(values) for side, values in wrist_steps.items()},
"hand_landmark_frame_step": {
side: summarize(values) for side, values in hand_steps.items()
},
"hand_span_relative_to_person": span_summary,
"bone_length_relative_to_box": bone_summary,
}
return result
def model_disagreement(records: list[dict]) -> tuple[dict, list[int]]:
values: dict[str, list[float]] = defaultdict(list)
per_frame: list[tuple[float, int]] = []
for record in records:
frame_max = 0.0
for person_idx, bbox in enumerate(record["boxes"]):
m = normalized_keypoints(
np.asarray(record["M"]["keypoints"][person_idx]), bbox
)
l = normalized_keypoints(
np.asarray(record["L"]["keypoints"][person_idx]), bbox
)
distance = np.linalg.norm(m - l, axis=1)
values[f"P{person_idx + 1}_upper"].extend(
float(distance[idx]) for idx in (5, 6, 7, 8, 9, 10)
)
for side, idx in BODY_WRISTS.items():
values[f"P{person_idx + 1}_{side}_wrist"].append(float(distance[idx]))
frame_max = max(frame_max, float(np.max(distance[[5, 6, 7, 8, 9, 10]])))
per_frame.append((frame_max, record["frame_index"]))
worst_frames = [frame for _, frame in sorted(per_frame, reverse=True)[:12]]
return {name: summarize(items) for name, items in values.items()}, worst_frames
def write_montage(
video_path: Path,
records: list[dict],
selected_frames: list[int],
output_path: Path,
titles: tuple[str, str] = ("RTMW-M 256x192", "RTMW-L 384x288"),
) -> None:
selected = set(selected_frames)
record_map = {record["frame_index"]: record for record in records}
cap = cv2.VideoCapture(str(video_path))
rows: list[np.ndarray] = []
frame_idx = 0
while True:
ok, frame = cap.read()
if not ok:
break
if frame_idx in selected:
record = record_map[frame_idx]
m = draw_pose(
frame,
record["boxes"],
np.asarray(record["M"]["keypoints"]),
f"{titles[0]} | frame {frame_idx}",
)
l = draw_pose(
frame,
record["boxes"],
np.asarray(record["L"]["keypoints"]),
f"{titles[1]} | frame {frame_idx}",
)
m = cv2.resize(m, (640, 360), interpolation=cv2.INTER_AREA)
l = cv2.resize(l, (640, 360), interpolation=cv2.INTER_AREA)
rows.append(np.hstack([m, l]))
frame_idx += 1
cap.release()
if rows:
cv2.imwrite(str(output_path), np.vstack(rows))
def main() -> None:
from rtmlib import RTMPose
args = parse_args()
args.output.mkdir(parents=True, exist_ok=True)
cached_boxes = load_boxes(args.boxes, args.max_people)
pose_m = RTMPose(
str(args.model_m), model_input_size=(192, 256), backend="onnxruntime", device="cpu"
)
pose_l = RTMPose(
str(args.model_l), model_input_size=(288, 384), backend="onnxruntime", device="cpu"
)
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))
video_output = args.output / "reference_video_rtmw_m_vs_l.mp4"
writer = cv2.VideoWriter(
str(video_output),
cv2.VideoWriter_fourcc(*"mp4v"),
fps,
(width, height // 2),
)
records: list[dict] = []
timings = {"M": [], "L": []}
frame_idx = 0
while True:
ok, frame = cap.read()
if not ok:
break
boxes = cached_boxes[frame_idx]
started = time.perf_counter()
kpts_m, scores_m = pose_m(frame, bboxes=boxes)
timings["M"].append((time.perf_counter() - started) * 1000.0)
started = time.perf_counter()
kpts_l, scores_l = pose_l(frame, bboxes=boxes)
timings["L"].append((time.perf_counter() - started) * 1000.0)
records.append(
{
"frame_index": frame_idx,
"boxes": boxes,
"M": {
"keypoints": np.round(kpts_m, 3).tolist(),
"raw_scores": np.round(scores_m, 4).tolist(),
},
"L": {
"keypoints": np.round(kpts_l, 3).tolist(),
"raw_scores": np.round(scores_l, 4).tolist(),
},
}
)
left = draw_pose(frame, boxes, kpts_m, f"RTMW-M 256x192 | frame {frame_idx}")
right = draw_pose(frame, boxes, kpts_l, f"RTMW-L 384x288 | frame {frame_idx}")
left = cv2.resize(left, (width // 2, height // 2), interpolation=cv2.INTER_AREA)
right = cv2.resize(right, (width // 2, height // 2), interpolation=cv2.INTER_AREA)
writer.write(np.hstack([left, right]))
if frame_idx % 50 == 0:
print(f"processed frame {frame_idx}", flush=True)
frame_idx += 1
cap.release()
writer.release()
jsonl_path = args.output / "reference_video_rtmw_m_vs_l.jsonl"
with jsonl_path.open("w", encoding="utf-8") as handle:
for record in records:
handle.write(json.dumps(record, ensure_ascii=False) + "\n")
disagreement, worst_frames = model_disagreement(records)
fixed_frames = [0, 60, 120, 180, 240, 252, 300, 360, 411]
montage_frames = sorted(set(fixed_frames + worst_frames))
write_montage(
args.video,
records,
montage_frames,
args.output / "rtmw_m_vs_l_full_montage.jpg",
)
summary = {
"frames": len(records),
"input_sizes_wh": {"M": [192, 256], "L": [288, 384]},
"raw_score_note": (
"M and L raw SimCC peak scores use different scales and must not be compared "
"or thresholded with the same numeric value."
),
"timing_ms_two_people_cpu": {
name: summarize(values[5:] if len(values) > 5 else values)
for name, values in timings.items()
},
"temporal_metrics_normalized_to_yolo_box": temporal_metrics(records),
"m_l_position_disagreement_normalized_to_yolo_box": disagreement,
"largest_upper_body_disagreement_frames": worst_frames,
"montage_frames": montage_frames,
}
with (args.output / "summary.json").open("w", encoding="utf-8") as handle:
json.dump(summary, handle, ensure_ascii=False, indent=2)
print(json.dumps(summary, ensure_ascii=False, indent=2))
if __name__ == "__main__":
main()
pipeline/native_backend.py
-9
View File
@@ -1,11 +1,8 @@
"""Native Jetson TensorRT adapter; temporal/risk policy stays in upstream Python."""
from __future__ import annotations
import ctypes as C
from collections import defaultdict, deque
from pathlib import Path
import numpy as np
from live_pose_engine import LivePoseEngine
from predict_optimized_pose_videos import BoxTracker
ROOT = Path(__file__).resolve().parents[1]
@@ -67,9 +64,3 @@ class NativeBackend:
def close(self):
if self.camera:self.lib.vap_camera_close(self.camera);self.camera=None
if self.engine:self.lib.vap_engine_close(self.engine);self.engine=None
class NativePoseEngine(LivePoseEngine):
def __init__(self,args,native):
self.args=args;self.native=native;self.pose=native.pose
self.box_tracker=BoxTracker(max_age=30);self.people={}
self.timings=defaultdict(lambda:deque(maxlen=3000));self.isolated_inferences=0
pipeline/occlusion_aware_pose_v2.py
-1130
View File
File diff suppressed because it is too large Load Diff
pipeline/live_pose_engine.py → pipeline/pose_engine.py
+14 -114
View File
@@ -1,113 +1,28 @@
"""Resident YOLO11s + RTMW-L pose stage extracted from the v9.9 pipeline."""
"""Resident native TensorRT inference with causal pose fusion and tracking."""
from __future__ import annotations
import argparse
import os
import time
from collections import defaultdict
from pathlib import Path
from collections import defaultdict, deque
import numpy as np
from pose_tracking import BoxTracker
SIDES = ("left", "right")
class LivePoseEngine:
"""Run the v9.9 YOLO/RTMW pose and temporal-fusion path without EdgeTAM."""
def __init__(self, args: argparse.Namespace):
self._trt_dll_directory = None
if args.pose_provider == "tensorrt" and os.name == "nt":
try:
import tensorrt_libs
except ImportError as exc:
raise RuntimeError(
"--pose-provider tensorrt on Windows requires the "
"tensorrt-cu12-libs package in the active Python environment."
) from exc
trt_lib_dir = Path(tensorrt_libs.__file__).resolve().parent
self._trt_dll_directory = os.add_dll_directory(str(trt_lib_dir))
os.environ["PATH"] = (
str(trt_lib_dir) + os.pathsep + os.environ.get("PATH", "")
)
if args.pose_device.lower().startswith("cuda"):
import onnxruntime as ort
if "CUDAExecutionProvider" not in ort.get_available_providers():
raise RuntimeError(
"--pose-device cuda requires onnxruntime-gpu. Uninstall "
"onnxruntime and install onnxruntime-gpu==1.22.0."
)
from rtmlib import RTMPose
from ultralytics import YOLO
from predict_optimized_pose_videos import BoxTracker
class NativePoseEngine:
def __init__(self, args, native):
self.args = args
self.detector = YOLO(str(args.det_model))
self.pose = RTMPose(
str(args.pose_model),
model_input_size=(288, 384),
backend="onnxruntime",
device=args.pose_device,
)
if args.pose_provider == "tensorrt":
import onnxruntime as ort
available = ort.get_available_providers()
if "TensorrtExecutionProvider" not in available:
raise RuntimeError(
"--pose-provider tensorrt requires ONNX Runtime with "
"TensorrtExecutionProvider. "
f"Available providers: {available}"
)
cache_dir = args.pose_trt_cache.expanduser().resolve()
cache_dir.mkdir(parents=True, exist_ok=True)
device_id = (
int(str(args.pose_device).split(":")[-1])
if ":" in str(args.pose_device)
else 0
)
trt_options = {
"device_id": device_id,
"trt_fp16_enable": "True",
"trt_engine_cache_enable": "True",
"trt_engine_cache_path": str(cache_dir),
"trt_timing_cache_enable": "True",
"trt_timing_cache_path": str(cache_dir),
}
self.pose.session = ort.InferenceSession(
str(args.pose_model.resolve()),
providers=[
("TensorrtExecutionProvider", trt_options),
("CUDAExecutionProvider", {"device_id": device_id}),
],
)
active_providers = self.pose.session.get_providers()
if not active_providers or active_providers[0] != "TensorrtExecutionProvider":
raise RuntimeError(
"RTMW TensorRT initialization failed instead of using the "
"requested provider. Active providers: "
f"{active_providers}"
)
print(f"RTMW ONNX Runtime providers: {active_providers}")
self.native = native
self.pose = native.pose
self.box_tracker = BoxTracker(max_age=30)
self.people = {}
self.timings = defaultdict(list)
self.timings = defaultdict(lambda: deque(maxlen=3000))
self.isolated_inferences = 0
def infer(self, frame: np.ndarray, frame_idx: int) -> dict:
from compare_rtmw_ml import HAND_SLICES
from fuse_rtmw_l_temporal import adaptive_smooth, build_global_hands, fuse_person
from occlusion_aware_pose_v2 import Measurement, resolve_global_hand_assignments
from predict_optimized_pose_videos import (
from pose_geometry import HAND_SLICES
from pose_fusion import adaptive_smooth, build_global_hands, fuse_person
from pose_occlusion import Measurement, resolve_global_hand_assignments
from pose_tracking import (
LivePersonState,
filter_person_boxes,
mask_other_people,
update_live_v2,
)
@@ -118,23 +33,8 @@ class LivePoseEngine:
detection_height = max(
1, min(height, round(height * self.args.det_roi_y_max_ratio))
)
if hasattr(self, "native"):
boxes, confidences = self.native.detect(frame[:detection_height])
self.timings["detector_ms"].append((time.perf_counter() - started) * 1000)
else:
detection = self.detector.predict(
frame[:detection_height],
imgsz=self.args.det_imgsz,
conf=self.args.det_conf,
iou=0.30,
classes=[0],
device=self.args.det_device,
verbose=False,
)[0]
self.timings["detector_ms"].append((time.perf_counter() - started) * 1000)
boxes, confidences, _ = filter_person_boxes(
detection, width, height, self.args.max_people
)
boxes, confidences = self.native.detect(frame[:detection_height])
self.timings["detector_ms"].append((time.perf_counter() - started) * 1000)
track_ids = self.box_tracker.update(boxes, frame_idx, (height, width))
# Retire disappeared people in long-running camera sessions.
pipeline/fuse_rtmw_l_temporal.py → pipeline/pose_fusion.py
+6 -331
View File
@@ -1,38 +1,18 @@
"""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.
"""
"""Runtime pose fusion helpers."""
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,
)
from pose_geometry import BODY_WRISTS, HAND_EDGES_LOCAL, HAND_SLICES, normalized_keypoints
BODY_SIDE_INDICES = {"left": (5, 7, 9), "right": (6, 8, 10)}
SIDE_COLORS = {"left": (40, 40, 255), "right": (30, 230, 30)}
HAND_COLORS = {"left": (255, 80, 220), "right": (255, 190, 40)}
@@ -71,15 +51,6 @@ class PersonTrack:
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)
@@ -557,299 +528,3 @@ def draw_fused(
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()
pipeline/pose_geometry.py
+28
View File
@@ -0,0 +1,28 @@
"""Runtime pose geometry helpers."""
from __future__ import annotations
import numpy as np
HAND_EDGES_LOCAL = (
(0, 1), (1, 2), (2, 3), (3, 4),
(0, 5), (5, 6), (6, 7), (7, 8),
(0, 9), (9, 10), (10, 11), (11, 12),
(0, 13), (13, 14), (14, 15), (15, 16),
(0, 17), (17, 18), (18, 19), (19, 20),
)
HAND_SLICES = {"left": (91, 112), "right": (112, 133)}
BODY_WRISTS = {"left": 9, "right": 10}
def normalized_keypoints(kpts: np.ndarray, bbox: list[float]) -> np.ndarray:
x1, y1, x2, y2 = bbox
width = max(x2 - x1, 1.0)
height = max(y2 - y1, 1.0)
result = kpts.copy().astype(np.float64)
result[:, 0] = (result[:, 0] - x1) / width
result[:, 1] = (result[:, 1] - y1) / height
return result
pipeline/pose_occlusion.py
+513
View File
@@ -0,0 +1,513 @@
"""Runtime pose occlusion helpers."""
from __future__ import annotations
import math
from dataclasses import dataclass
import numpy as np
from pose_geometry import HAND_SLICES
@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 bbox_diagonal(bbox: list[float]) -> float:
return max(math.hypot(bbox[2] - bbox[0], bbox[3] - bbox[1]), 1.0)
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
pipeline/pose_tracking.py
+304
View File
@@ -0,0 +1,304 @@
"""Runtime pose tracking helpers."""
from __future__ import annotations
import math
from collections import deque
from dataclasses import dataclass, field
import numpy as np
from pose_geometry import HAND_SLICES
from pose_fusion import BODY_SIDE_INDICES, PersonTrack
from pose_occlusion import EndpointTrack, Measurement, calibrated_bone_lengths, reachable_region, solve_two_bone_elbow, temporal_recovery_measurement, update_endpoint_track
def bbox_iou(a: np.ndarray, b: np.ndarray) -> float:
x1 = max(float(a[0]), float(b[0]))
y1 = max(float(a[1]), float(b[1]))
x2 = min(float(a[2]), float(b[2]))
y2 = min(float(a[3]), float(b[3]))
intersection = max(x2 - x1, 0.0) * max(y2 - y1, 0.0)
union = (
max(float(a[2] - a[0]), 0.0) * max(float(a[3] - a[1]), 0.0)
+ max(float(b[2] - b[0]), 0.0) * max(float(b[3] - b[1]), 0.0)
- intersection
)
return intersection / max(union, 1.0)
class BoxTracker:
def __init__(self, max_age: int = 30):
self.max_age = max_age
self.next_id = 1
self.tracks: dict[int, dict] = {}
def update(
self,
boxes: list[list[float]],
frame_idx: int,
image_shape: tuple[int, int],
) -> list[int]:
height, width = image_shape
image_diagonal = math.hypot(width, height)
active_ids = [
track_id
for track_id, track in self.tracks.items()
if frame_idx - track["last_frame"] <= self.max_age
]
self.tracks = {track_id: self.tracks[track_id] for track_id in active_ids}
pairs = []
for detection_idx, box in enumerate(boxes):
box_array = np.asarray(box, dtype=np.float64)
center = (box_array[:2] + box_array[2:]) * 0.5
for track_id in active_ids:
prior = np.asarray(self.tracks[track_id]["box"], dtype=np.float64)
prior_center = (prior[:2] + prior[2:]) * 0.5
iou = bbox_iou(box_array, prior)
center_distance = float(np.linalg.norm(center - prior_center)) / image_diagonal
if iou >= 0.05 or center_distance <= 0.18:
score = 2.0 * iou + max(0.0, 1.0 - center_distance / 0.18)
pairs.append((score, detection_idx, track_id))
assigned_detections = set()
assigned_tracks = set()
result: list[int | None] = [None] * len(boxes)
for _, detection_idx, track_id in sorted(pairs, reverse=True):
if detection_idx in assigned_detections or track_id in assigned_tracks:
continue
assigned_detections.add(detection_idx)
assigned_tracks.add(track_id)
result[detection_idx] = track_id
for detection_idx, box in enumerate(boxes):
track_id = result[detection_idx]
if track_id is None:
track_id = self.next_id
self.next_id += 1
result[detection_idx] = track_id
self.tracks[track_id] = {"box": box, "last_frame": frame_idx}
return [int(track_id) for track_id in result]
@dataclass
class LivePersonState:
fusion: PersonTrack = field(default_factory=PersonTrack)
endpoints: dict[str, EndpointTrack] = field(
default_factory=lambda: {"left": EndpointTrack(), "right": EndpointTrack()}
)
bone_samples: dict[str, deque] = field(
default_factory=lambda: {
"left": deque(maxlen=120),
"right": deque(maxlen=120),
}
)
last_frame: int = -1
def bbox_diagonal(bbox: list[float]) -> float:
return max(math.hypot(bbox[2] - bbox[0], bbox[3] - bbox[1]), 1.0)
def update_bone_samples(
state: LivePersonState,
keypoints: np.ndarray,
metadata: dict,
bbox: list[float],
) -> None:
diagonal = bbox_diagonal(bbox)
for side, (shoulder, elbow, wrist) in BODY_SIDE_INDICES.items():
meta = metadata[side]
if meta["limb_state"] != "consensus" or not meta["distance_eligible"]:
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:
state.bone_samples[side].append((upper, forearm))
def current_bone_ratios(state: LivePersonState, side: str) -> tuple[float, float]:
samples = state.bone_samples[side]
if len(samples) < 10:
opposite = "right" if side == "left" else "left"
if len(state.bone_samples[opposite]) >= 10:
samples = state.bone_samples[opposite]
if not samples:
return 0.24, 0.22
values = np.asarray(samples, dtype=np.float64)
return float(np.median(values[:, 0])), float(np.median(values[:, 1]))
def update_live_v2(
frame_idx: int,
bbox: list[float],
selected_keypoints: np.ndarray,
smoothed_keypoints: np.ndarray,
fusion_metadata: dict,
state: LivePersonState,
assigned_measurements: dict[str, Measurement | None] | None,
identity_assignments: dict[str, dict] | None,
short_occlusion_frames: int,
medium_occlusion_frames: int,
) -> tuple[np.ndarray, dict]:
frame_gap = 0 if state.last_frame < 0 else max(frame_idx - state.last_frame - 1, 0)
for endpoint in state.endpoints.values():
endpoint.missed_frames += frame_gap
state.last_frame = frame_idx
update_bone_samples(state, selected_keypoints, fusion_metadata, bbox)
output = smoothed_keypoints.copy()
output_metadata = {}
diagonal = bbox_diagonal(bbox)
for side, (shoulder_idx, elbow_idx, wrist_idx) in BODY_SIDE_INDICES.items():
endpoint = state.endpoints[side]
meta = fusion_metadata[side]
start, end = HAND_SLICES[side]
measurement = (
assigned_measurements.get(side)
if assigned_measurements is not None
else None
)
if assigned_measurements is None and meta["alarm_eligible"]:
measurement = Measurement(
keypoints=selected_keypoints[start:end],
precise=bool(meta["distance_eligible"]),
provisional=not bool(meta["distance_eligible"]),
quality=float(meta.get("selected_hand_score", 0.0)),
)
if measurement is None:
measurement = temporal_recovery_measurement(
selected_keypoints,
meta,
side,
bbox,
endpoint,
medium_occlusion_frames,
)
returning_after_long = (
measurement is not None
and endpoint.missed_frames > short_occlusion_frames
)
wrist, hand, endpoint_state = update_endpoint_track(
endpoint,
measurement,
frame_idx,
short_occlusion_frames,
medium_occlusion_frames,
)
shoulder = smoothed_keypoints[shoulder_idx]
upper_ratio, forearm_ratio = current_bone_ratios(state, side)
upper_length, forearm_length = calibrated_bone_lengths(
upper_ratio, forearm_ratio, bbox
)
region = None
if wrist is not None and hand is not None:
reference_elbow = (
selected_keypoints[elbow_idx]
if returning_after_long or endpoint.last_elbow is None
else endpoint.last_elbow
)
use_ik = (
meta["limb_state"] != "consensus"
or endpoint_state
in ("provisional", "recovered", "reassigned", "predicted_short")
)
elbow = (
solve_two_bone_elbow(
shoulder,
wrist,
upper_length,
forearm_length,
reference_elbow,
)
if use_ik
else smoothed_keypoints[elbow_idx]
)
endpoint.last_elbow = elbow.copy()
direction = wrist - shoulder
if float(np.linalg.norm(direction)) > 1e-6:
endpoint.last_direction = direction / np.linalg.norm(direction)
output[shoulder_idx] = shoulder
output[elbow_idx] = elbow
output[wrist_idx] = wrist
output[start:end] = hand
uncertainty = (
0.025 * diagonal
if endpoint_state == "observed"
else (0.04 + 0.015 * endpoint.missed_frames) * diagonal
)
precise = endpoint_state == "observed"
else:
region = reachable_region(
shoulder,
upper_length,
forearm_length,
bbox,
endpoint,
endpoint_state,
short_occlusion_frames,
medium_occlusion_frames,
)
uncertainty = region["outer_radius_px"]
precise = False
output_metadata[side] = {
"state": endpoint_state,
"precise_distance": precise,
"fast_alarm_candidate": True,
"risk_mode": (
"precise_hand"
if endpoint_state == "observed"
else (
"motion_corridor_intersection"
if endpoint_state == "reachable_medium"
else (
"door_intersection_only"
if endpoint_state == "reachable_long"
else (
"reassigned_hand"
if endpoint_state == "reassigned"
else "temporal_hand"
)
)
)
),
"requires_door_intersection": endpoint_state
in ("reachable_medium", "reachable_long"),
"identity_assignment": (
identity_assignments.get(side)
if identity_assignments is not None
else None
),
"missed_frames": endpoint.missed_frames,
"uncertainty_radius_px": round(float(uncertainty), 3),
"bone_lengths_px": {
"upper": round(upper_length, 3),
"forearm": round(forearm_length, 3),
},
"reachable_region": region,
}
return output, output_metadata
def mask_other_people(
frame: np.ndarray,
boxes: list[list[float]],
target_idx: int,
) -> np.ndarray:
masked = frame.copy()
height, width = frame.shape[:2]
mask = np.zeros((height, width), dtype=bool)
target = boxes[target_idx]
target_area = max((target[2] - target[0]) * (target[3] - target[1]), 1.0)
for idx, bbox in enumerate(boxes):
if idx == target_idx:
continue
other_area = max((bbox[2] - bbox[0]) * (bbox[3] - bbox[1]), 1.0)
# The failure in this video is the small background person being
# absorbed by the much larger foreground person's 3:4 pose crop.
# Mask only markedly smaller people. Masking a larger occluder while
# estimating a small person removes genuine target pixels at overlap.
if other_area >= 0.5 * target_area:
continue
x1, y1, x2, y2 = (int(round(value)) for value in bbox)
x1, x2 = max(0, x1), min(width - 1, x2)
y1, y2 = max(0, y1), min(height - 1, y2)
mask[y1:y2 + 1, x1:x2 + 1] = True
masked[mask] = (124, 116, 104)
return masked
pipeline/predict_optimized_pose_videos.py
-683
View File
@@ -1,683 +0,0 @@
"""End-to-end optimized RTMW-L prediction for CMS videos.
Pipeline:
YOLO11s person -> fixed-camera false-positive filters -> lightweight track IDs
-> RTMW-L raw/smaller-person-isolated views -> temporal fusion -> causal
occlusion-aware V4 global hand assignment. Reachable regions stay in JSON and are hidden in the main
comparison video.
"""
from __future__ import annotations
import argparse
import json
import math
import time
from collections import Counter, defaultdict, deque
from dataclasses import dataclass, field
from pathlib import Path
import cv2
import numpy as np
from compare_rtmw_ml import HAND_SLICES, draw_pose, summarize
from fuse_rtmw_l_temporal import (
BODY_SIDE_INDICES,
PersonTrack,
adaptive_smooth,
build_global_hands,
draw_fused,
fuse_person,
)
from occlusion_aware_pose_v2 import (
EndpointTrack,
Measurement,
calibrated_bone_lengths,
draw_v2,
reachable_region,
resolve_global_hand_assignments,
solve_two_bone_elbow,
temporal_recovery_measurement,
update_endpoint_track,
)
from rtmw_l_person_mask_ablation import mask_other_people
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser()
parser.add_argument("--videos", type=Path, nargs="+", required=True)
parser.add_argument("--det-model", type=Path, required=True)
parser.add_argument("--pose-model", type=Path, required=True)
parser.add_argument("--output", type=Path, required=True)
parser.add_argument(
"--det-device",
default="0",
help="Ultralytics YOLO device, for example 0, cuda:0, or cpu.",
)
parser.add_argument(
"--pose-device",
default="cpu",
help="RTMW ONNX Runtime device: cpu or cuda.",
)
parser.add_argument("--det-imgsz", type=int, default=960)
parser.add_argument("--det-conf", type=float, default=0.30)
parser.add_argument("--det-roi-y-max-ratio", type=float, default=1.0)
parser.add_argument("--short-occlusion-frames", type=int, default=12)
parser.add_argument("--medium-occlusion-frames", type=int, default=30)
parser.add_argument("--max-people", type=int, default=3)
return parser.parse_args()
def bbox_iou(a: np.ndarray, b: np.ndarray) -> float:
x1 = max(float(a[0]), float(b[0]))
y1 = max(float(a[1]), float(b[1]))
x2 = min(float(a[2]), float(b[2]))
y2 = min(float(a[3]), float(b[3]))
intersection = max(x2 - x1, 0.0) * max(y2 - y1, 0.0)
union = (
max(float(a[2] - a[0]), 0.0) * max(float(a[3] - a[1]), 0.0)
+ max(float(b[2] - b[0]), 0.0) * max(float(b[3] - b[1]), 0.0)
- intersection
)
return intersection / max(union, 1.0)
class BoxTracker:
def __init__(self, max_age: int = 30):
self.max_age = max_age
self.next_id = 1
self.tracks: dict[int, dict] = {}
def update(
self,
boxes: list[list[float]],
frame_idx: int,
image_shape: tuple[int, int],
) -> list[int]:
height, width = image_shape
image_diagonal = math.hypot(width, height)
active_ids = [
track_id
for track_id, track in self.tracks.items()
if frame_idx - track["last_frame"] <= self.max_age
]
self.tracks = {track_id: self.tracks[track_id] for track_id in active_ids}
pairs = []
for detection_idx, box in enumerate(boxes):
box_array = np.asarray(box, dtype=np.float64)
center = (box_array[:2] + box_array[2:]) * 0.5
for track_id in active_ids:
prior = np.asarray(self.tracks[track_id]["box"], dtype=np.float64)
prior_center = (prior[:2] + prior[2:]) * 0.5
iou = bbox_iou(box_array, prior)
center_distance = float(np.linalg.norm(center - prior_center)) / image_diagonal
if iou >= 0.05 or center_distance <= 0.18:
score = 2.0 * iou + max(0.0, 1.0 - center_distance / 0.18)
pairs.append((score, detection_idx, track_id))
assigned_detections = set()
assigned_tracks = set()
result: list[int | None] = [None] * len(boxes)
for _, detection_idx, track_id in sorted(pairs, reverse=True):
if detection_idx in assigned_detections or track_id in assigned_tracks:
continue
assigned_detections.add(detection_idx)
assigned_tracks.add(track_id)
result[detection_idx] = track_id
for detection_idx, box in enumerate(boxes):
track_id = result[detection_idx]
if track_id is None:
track_id = self.next_id
self.next_id += 1
result[detection_idx] = track_id
self.tracks[track_id] = {"box": box, "last_frame": frame_idx}
return [int(track_id) for track_id in result]
@dataclass
class LivePersonState:
fusion: PersonTrack = field(default_factory=PersonTrack)
endpoints: dict[str, EndpointTrack] = field(
default_factory=lambda: {"left": EndpointTrack(), "right": EndpointTrack()}
)
bone_samples: dict[str, deque] = field(
default_factory=lambda: {
"left": deque(maxlen=120),
"right": deque(maxlen=120),
}
)
last_frame: int = -1
def bbox_diagonal(bbox: list[float]) -> float:
return max(math.hypot(bbox[2] - bbox[0], bbox[3] - bbox[1]), 1.0)
def update_bone_samples(
state: LivePersonState,
keypoints: np.ndarray,
metadata: dict,
bbox: list[float],
) -> None:
diagonal = bbox_diagonal(bbox)
for side, (shoulder, elbow, wrist) in BODY_SIDE_INDICES.items():
meta = metadata[side]
if meta["limb_state"] != "consensus" or not meta["distance_eligible"]:
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:
state.bone_samples[side].append((upper, forearm))
def current_bone_ratios(state: LivePersonState, side: str) -> tuple[float, float]:
samples = state.bone_samples[side]
if len(samples) < 10:
opposite = "right" if side == "left" else "left"
if len(state.bone_samples[opposite]) >= 10:
samples = state.bone_samples[opposite]
if not samples:
return 0.24, 0.22
values = np.asarray(samples, dtype=np.float64)
return float(np.median(values[:, 0])), float(np.median(values[:, 1]))
def update_live_v2(
frame_idx: int,
bbox: list[float],
selected_keypoints: np.ndarray,
smoothed_keypoints: np.ndarray,
fusion_metadata: dict,
state: LivePersonState,
assigned_measurements: dict[str, Measurement | None] | None,
identity_assignments: dict[str, dict] | None,
short_occlusion_frames: int,
medium_occlusion_frames: int,
) -> tuple[np.ndarray, dict]:
frame_gap = 0 if state.last_frame < 0 else max(frame_idx - state.last_frame - 1, 0)
for endpoint in state.endpoints.values():
endpoint.missed_frames += frame_gap
state.last_frame = frame_idx
update_bone_samples(state, selected_keypoints, fusion_metadata, bbox)
output = smoothed_keypoints.copy()
output_metadata = {}
diagonal = bbox_diagonal(bbox)
for side, (shoulder_idx, elbow_idx, wrist_idx) in BODY_SIDE_INDICES.items():
endpoint = state.endpoints[side]
meta = fusion_metadata[side]
start, end = HAND_SLICES[side]
measurement = (
assigned_measurements.get(side)
if assigned_measurements is not None
else None
)
if assigned_measurements is None and meta["alarm_eligible"]:
measurement = Measurement(
keypoints=selected_keypoints[start:end],
precise=bool(meta["distance_eligible"]),
provisional=not bool(meta["distance_eligible"]),
quality=float(meta.get("selected_hand_score", 0.0)),
)
if measurement is None:
measurement = temporal_recovery_measurement(
selected_keypoints,
meta,
side,
bbox,
endpoint,
medium_occlusion_frames,
)
returning_after_long = (
measurement is not None
and endpoint.missed_frames > short_occlusion_frames
)
wrist, hand, endpoint_state = update_endpoint_track(
endpoint,
measurement,
frame_idx,
short_occlusion_frames,
medium_occlusion_frames,
)
shoulder = smoothed_keypoints[shoulder_idx]
upper_ratio, forearm_ratio = current_bone_ratios(state, side)
upper_length, forearm_length = calibrated_bone_lengths(
upper_ratio, forearm_ratio, bbox
)
region = None
if wrist is not None and hand is not None:
reference_elbow = (
selected_keypoints[elbow_idx]
if returning_after_long or endpoint.last_elbow is None
else endpoint.last_elbow
)
use_ik = (
meta["limb_state"] != "consensus"
or endpoint_state
in ("provisional", "recovered", "reassigned", "predicted_short")
)
elbow = (
solve_two_bone_elbow(
shoulder,
wrist,
upper_length,
forearm_length,
reference_elbow,
)
if use_ik
else smoothed_keypoints[elbow_idx]
)
endpoint.last_elbow = elbow.copy()
direction = wrist - shoulder
if float(np.linalg.norm(direction)) > 1e-6:
endpoint.last_direction = direction / np.linalg.norm(direction)
output[shoulder_idx] = shoulder
output[elbow_idx] = elbow
output[wrist_idx] = wrist
output[start:end] = hand
uncertainty = (
0.025 * diagonal
if endpoint_state == "observed"
else (0.04 + 0.015 * endpoint.missed_frames) * diagonal
)
precise = endpoint_state == "observed"
else:
region = reachable_region(
shoulder,
upper_length,
forearm_length,
bbox,
endpoint,
endpoint_state,
short_occlusion_frames,
medium_occlusion_frames,
)
uncertainty = region["outer_radius_px"]
precise = False
output_metadata[side] = {
"state": endpoint_state,
"precise_distance": precise,
"fast_alarm_candidate": True,
"risk_mode": (
"precise_hand"
if endpoint_state == "observed"
else (
"motion_corridor_intersection"
if endpoint_state == "reachable_medium"
else (
"door_intersection_only"
if endpoint_state == "reachable_long"
else (
"reassigned_hand"
if endpoint_state == "reassigned"
else "temporal_hand"
)
)
)
),
"requires_door_intersection": endpoint_state
in ("reachable_medium", "reachable_long"),
"identity_assignment": (
identity_assignments.get(side)
if identity_assignments is not None
else None
),
"missed_frames": endpoint.missed_frames,
"uncertainty_radius_px": round(float(uncertainty), 3),
"bone_lengths_px": {
"upper": round(upper_length, 3),
"forearm": round(forearm_length, 3),
},
"reachable_region": region,
}
return output, output_metadata
def filter_person_boxes(
result,
width: int,
height: int,
max_people: int,
) -> tuple[list[list[float]], list[float], Counter]:
boxes = result.boxes.xyxy.cpu().numpy()
scores = result.boxes.conf.cpu().numpy()
rejected = Counter()
kept = []
for box, score in zip(boxes, scores):
center_x = float(box[0] + box[2]) * 0.5
box_height = float(box[3] - box[1])
if center_x > 0.80 * width:
rejected["right_reflection_roi"] += 1
continue
if box_height < 0.20 * height:
rejected["too_short"] += 1
continue
kept.append((box.tolist(), float(score)))
kept.sort(key=lambda item: item[1], reverse=True)
if len(kept) > max_people:
rejected["over_max_people"] += len(kept) - max_people
kept = kept[:max_people]
kept.sort(key=lambda item: item[0][0])
return [item[0] for item in kept], [item[1] for item in kept], rejected
def process_video(
video_path: Path,
output_dir: Path,
detector: YOLO,
pose: RTMPose,
det_device: str,
det_imgsz: int,
det_conf: float,
det_roi_y_max_ratio: float,
short_occlusion_frames: int,
medium_occlusion_frames: int,
max_people: int,
video_label: str,
) -> dict:
output_dir.mkdir(parents=True, exist_ok=True)
cap = cv2.VideoCapture(str(video_path))
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(output_dir / "raw_temporal_occlusion_v4.mp4"),
cv2.VideoWriter_fourcc(*"mp4v"),
fps,
(1920, 360),
)
box_tracker = BoxTracker(max_age=30)
people: dict[int, LivePersonState] = {}
records = []
detection_counts = Counter()
rejection_counts = Counter()
state_counts = Counter()
timings = defaultdict(list)
isolated_inferences = 0
frame_idx = 0
while True:
ok, frame = cap.read()
if not ok:
break
started = time.perf_counter()
detection_height = max(1, min(height, round(height * det_roi_y_max_ratio)))
detection = detector.predict(
frame[:detection_height],
imgsz=det_imgsz,
conf=det_conf,
iou=0.30,
classes=[0],
device=det_device,
verbose=False,
)[0]
timings["detector_ms"].append((time.perf_counter() - started) * 1000.0)
boxes, confidences, rejected = filter_person_boxes(
detection, width, height, max_people
)
rejection_counts.update(rejected)
detection_counts[len(boxes)] += 1
track_ids = box_tracker.update(boxes, frame_idx, (height, width))
if boxes:
started = time.perf_counter()
raw_keypoints, raw_scores = pose(frame, bboxes=boxes)
timings["raw_pose_ms"].append((time.perf_counter() - started) * 1000.0)
isolated_keypoints = raw_keypoints.copy()
isolated_scores = raw_scores.copy()
started = time.perf_counter()
for person_idx, bbox in enumerate(boxes):
target_area = max((bbox[2] - bbox[0]) * (bbox[3] - bbox[1]), 1.0)
needs_isolation = any(
idx != person_idx
and max((other[2] - other[0]) * (other[3] - other[1]), 1.0)
< 0.5 * target_area
for idx, other in enumerate(boxes)
)
if not needs_isolation:
continue
masked = mask_other_people(frame, boxes, person_idx)
kpts, scores = pose(masked, bboxes=[bbox])
isolated_keypoints[person_idx] = kpts[0]
isolated_scores[person_idx] = scores[0]
isolated_inferences += 1
timings["isolated_extra_ms"].append((time.perf_counter() - started) * 1000.0)
else:
raw_keypoints = np.empty((0, 133, 2), dtype=np.float64)
raw_scores = np.empty((0, 133), dtype=np.float64)
isolated_keypoints = raw_keypoints.copy()
isolated_scores = raw_scores.copy()
selected_list = []
smoothed_list = []
fusion_meta_list = []
person_states = []
for person_idx, (bbox, track_id) in enumerate(zip(boxes, track_ids)):
state = people.setdefault(track_id, LivePersonState())
selected, fusion_meta = fuse_person(
raw_keypoints[person_idx],
raw_scores[person_idx],
isolated_keypoints[person_idx],
isolated_scores[person_idx],
bbox,
state.fusion,
)
smoothed = adaptive_smooth(selected, bbox, state.fusion)
selected_list.append(selected)
smoothed_list.append(smoothed)
fusion_meta_list.append(fusion_meta)
person_states.append(state)
frame_measurements = {}
endpoint_tracks = {}
for person_idx, (selected, fusion_meta, state) in enumerate(
zip(selected_list, fusion_meta_list, person_states)
):
for side, (start, end) in HAND_SLICES.items():
meta = fusion_meta[side]
frame_measurements[(person_idx, side)] = (
Measurement(
keypoints=selected[start:end],
precise=bool(meta["distance_eligible"]),
provisional=not bool(meta["distance_eligible"]),
quality=float(meta.get("selected_hand_score", 0.0)),
)
if meta["alarm_eligible"]
else None
)
endpoint_tracks[(person_idx, side)] = state.endpoints[side]
frame_measurements, assignment_events = resolve_global_hand_assignments(
frame_measurements,
endpoint_tracks,
boxes,
medium_occlusion_frames,
)
v2_list = []
v2_meta_list = []
for person_idx, (bbox, selected, smoothed, fusion_meta, state) in enumerate(
zip(
boxes,
selected_list,
smoothed_list,
fusion_meta_list,
person_states,
)
):
v2, v2_meta = update_live_v2(
frame_idx,
bbox,
selected,
smoothed,
fusion_meta,
state,
{
side: frame_measurements[(person_idx, side)]
for side in ("left", "right")
},
{
side: assignment_events.get((person_idx, side))
for side in ("left", "right")
},
short_occlusion_frames,
medium_occlusion_frames,
)
v2_list.append(v2)
v2_meta_list.append(v2_meta)
for side in ("left", "right"):
state_counts[f"{v2_meta[side]['state']}"] += 1
global_hands = build_global_hands(selected_list, fusion_meta_list)
records.append(
{
"frame_index": frame_idx,
"timestamp_ms": round(frame_idx / fps * 1000.0, 3),
"track_ids": track_ids,
"boxes": np.round(np.asarray(boxes), 3).tolist() if boxes else [],
"det_confidences": [round(value, 4) for value in confidences],
"raw_keypoints": np.round(raw_keypoints, 3).tolist(),
"temporal_selected_keypoints": np.round(np.asarray(selected_list), 3).tolist(),
"temporal_keypoints": np.round(np.asarray(smoothed_list), 3).tolist(),
"temporal_metadata": fusion_meta_list,
"v2_keypoints": np.round(np.asarray(v2_list), 3).tolist(),
"v2_metadata": v2_meta_list,
"global_hands": global_hands,
}
)
raw_panel = draw_pose(
frame,
boxes,
raw_keypoints,
f"{video_label} RTMW-L raw | frame {frame_idx}",
)
temporal_panel = draw_fused(
frame,
boxes,
smoothed_list,
fusion_meta_list,
f"{video_label} temporal fusion | frame {frame_idx}",
)
v2_panel = draw_v2(
frame,
boxes,
np.asarray(v2_list),
v2_meta_list,
f"{video_label} occlusion V4 | frame {frame_idx}",
show_reachable=False,
)
panels = [
cv2.resize(panel, (640, 360), interpolation=cv2.INTER_AREA)
for panel in (raw_panel, temporal_panel, v2_panel)
]
writer.write(np.hstack(panels))
if frame_idx == 0 and video_path.suffix.lower() in {
".jpg",
".jpeg",
".png",
".bmp",
".webp",
}:
cv2.imwrite(str(output_dir / "raw_pose.jpg"), raw_panel)
cv2.imwrite(str(output_dir / "spatial_fusion.jpg"), temporal_panel)
cv2.imwrite(str(output_dir / "occlusion_v2.jpg"), v2_panel)
panel_width = 960
panel_height = round(height * panel_width / width)
image_panels = [
cv2.resize(
panel,
(panel_width, panel_height),
interpolation=cv2.INTER_AREA,
)
for panel in (raw_panel, temporal_panel, v2_panel)
]
cv2.imwrite(
str(output_dir / "raw_spatial_occlusion_v2.jpg"),
np.hstack(image_panels),
)
if frame_idx % 100 == 0:
print(f"{video_label}: processed frame {frame_idx}", flush=True)
frame_idx += 1
cap.release()
writer.release()
with (output_dir / "predictions.jsonl").open("w", encoding="utf-8") as handle:
for record in records:
handle.write(json.dumps(record, ensure_ascii=False) + "\n")
summary = {
"occlusion_algorithm": "v4_global_assignment",
"video": str(video_path.resolve()),
"frames": len(records),
"fps": round(fps, 3),
"resolution": [width, height],
"person_filter": {
"center_x_max_ratio": 0.80,
"minimum_box_height_ratio": 0.20,
"max_people": max_people,
},
"detector_settings": {
"imgsz": det_imgsz,
"confidence": det_conf,
"roi_y_max_ratio": det_roi_y_max_ratio,
},
"filtered_person_count_frames": dict(sorted(detection_counts.items())),
"rejected_detection_counts": dict(sorted(rejection_counts.items())),
"track_ids_created": box_tracker.next_id - 1,
"isolated_pose_inferences": isolated_inferences,
"short_occlusion_frames": short_occlusion_frames,
"medium_occlusion_frames": medium_occlusion_frames,
"v2_state_counts": dict(sorted(state_counts.items())),
"timing_ms_cpu": {
name: summarize(values[5:] if len(values) > 5 else values)
for name, values in timings.items()
},
"output_note": (
"Main video hides reachable-region arcs. Reachable-region geometry remains in "
"predictions.jsonl for later door-distance lower-bound calculation."
),
}
with (output_dir / "summary.json").open("w", encoding="utf-8") as handle:
json.dump(summary, handle, ensure_ascii=False, indent=2)
return summary
def main() -> None:
from rtmlib import RTMPose
from ultralytics import YOLO
args = parse_args()
detector = YOLO(str(args.det_model))
pose = RTMPose(
str(args.pose_model),
model_input_size=(288, 384),
backend="onnxruntime",
device=args.pose_device,
)
all_summaries = {}
for video_idx, video_path in enumerate(args.videos, start=1):
label = f"video{video_idx}"
output_dir = args.output / label
all_summaries[label] = process_video(
video_path,
output_dir,
detector,
pose,
args.det_device,
args.det_imgsz,
args.det_conf,
args.det_roi_y_max_ratio,
args.short_occlusion_frames,
args.medium_occlusion_frames,
args.max_people,
label,
)
args.output.mkdir(parents=True, exist_ok=True)
with (args.output / "summary.json").open("w", encoding="utf-8") as handle:
json.dump(all_summaries, handle, ensure_ascii=False, indent=2)
print(json.dumps(all_summaries, ensure_ascii=False, indent=2))
if __name__ == "__main__":
main()
pipeline/rtmw_l_person_mask_ablation.py
-166
View File
@@ -1,166 +0,0 @@
"""Test whether other people inside a top-down crop cause RTMW-L mix-ups."""
from __future__ import annotations
import argparse
import json
import time
from pathlib import Path
import cv2
import numpy as np
from compare_rtmw_ml import (
draw_pose,
model_disagreement,
summarize,
temporal_metrics,
write_montage,
)
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser()
parser.add_argument("--video", type=Path, required=True)
parser.add_argument("--baseline-jsonl", type=Path, required=True)
parser.add_argument("--model-l", type=Path, required=True)
parser.add_argument("--output", type=Path, required=True)
return parser.parse_args()
def mask_other_people(
frame: np.ndarray,
boxes: list[list[float]],
target_idx: int,
) -> np.ndarray:
masked = frame.copy()
height, width = frame.shape[:2]
mask = np.zeros((height, width), dtype=bool)
target = boxes[target_idx]
target_area = max((target[2] - target[0]) * (target[3] - target[1]), 1.0)
for idx, bbox in enumerate(boxes):
if idx == target_idx:
continue
other_area = max((bbox[2] - bbox[0]) * (bbox[3] - bbox[1]), 1.0)
# The failure in this video is the small background person being
# absorbed by the much larger foreground person's 3:4 pose crop.
# Mask only markedly smaller people. Masking a larger occluder while
# estimating a small person removes genuine target pixels at overlap.
if other_area >= 0.5 * target_area:
continue
x1, y1, x2, y2 = (int(round(value)) for value in bbox)
x1, x2 = max(0, x1), min(width - 1, x2)
y1, y2 = max(0, y1), min(height - 1, y2)
mask[y1:y2 + 1, x1:x2 + 1] = True
masked[mask] = (124, 116, 104)
return masked
def main() -> None:
from rtmlib import RTMPose
args = parse_args()
args.output.mkdir(parents=True, exist_ok=True)
baseline = [
json.loads(line)
for line in args.baseline_jsonl.read_text(encoding="utf-8").splitlines()
]
pose_l = RTMPose(
str(args.model_l), model_input_size=(288, 384), backend="onnxruntime", device="cpu"
)
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_vs_person_masked.mp4"),
cv2.VideoWriter_fourcc(*"mp4v"),
fps,
(width, height // 2),
)
comparison: list[dict] = []
timings: list[float] = []
frame_idx = 0
while True:
ok, frame = cap.read()
if not ok:
break
baseline_record = baseline[frame_idx]
boxes = baseline_record["boxes"][:2]
masked_keypoints = []
masked_scores = []
started = time.perf_counter()
for person_idx, bbox in enumerate(boxes):
masked_frame = mask_other_people(frame, boxes, person_idx)
keypoints, scores = pose_l(masked_frame, bboxes=[bbox])
masked_keypoints.append(keypoints[0])
masked_scores.append(scores[0])
timings.append((time.perf_counter() - started) * 1000.0)
masked_keypoints_array = np.asarray(masked_keypoints)
raw_keypoints = np.asarray(baseline_record["L"]["keypoints"][:2])
comparison.append(
{
"frame_index": frame_idx,
"boxes": boxes,
"M": {"keypoints": raw_keypoints.tolist()},
"L": {"keypoints": np.round(masked_keypoints_array, 3).tolist()},
"masked_raw_scores": np.round(np.asarray(masked_scores), 4).tolist(),
}
)
left = draw_pose(frame, boxes, raw_keypoints, f"RTMW-L raw | frame {frame_idx}")
right = draw_pose(
frame,
boxes,
masked_keypoints_array,
f"RTMW-L mask other person | frame {frame_idx}",
)
left = cv2.resize(left, (width // 2, height // 2), interpolation=cv2.INTER_AREA)
right = cv2.resize(right, (width // 2, height // 2), interpolation=cv2.INTER_AREA)
writer.write(np.hstack([left, right]))
if frame_idx % 50 == 0:
print(f"processed frame {frame_idx}", flush=True)
frame_idx += 1
cap.release()
writer.release()
with (args.output / "rtmw_l_raw_vs_person_masked.jsonl").open(
"w", encoding="utf-8"
) as handle:
for record in comparison:
handle.write(json.dumps(record, ensure_ascii=False) + "\n")
disagreement, worst_frames = model_disagreement(comparison)
fixed_frames = [0, 60, 120, 132, 136, 180, 240, 252, 300, 360, 411]
montage_frames = sorted(set(fixed_frames + worst_frames[:8]))
write_montage(
args.video,
comparison,
montage_frames,
args.output / "rtmw_l_raw_vs_person_masked_montage.jpg",
titles=("RTMW-L raw", "RTMW-L smaller-person masked"),
)
summary = {
"frames": len(comparison),
"labels": {
"M": "RTMW-L raw crop",
"L": "RTMW-L with markedly smaller people masked",
},
"timing_ms_two_people_cpu": summarize(timings[5:]),
"temporal_metrics_normalized_to_yolo_box": temporal_metrics(comparison),
"position_change_after_masking": disagreement,
"largest_change_frames": worst_frames,
"montage_frames": montage_frames,
}
with (args.output / "summary.json").open("w", encoding="utf-8") as handle:
json.dump(summary, handle, ensure_ascii=False, indent=2)
print(json.dumps(summary, ensure_ascii=False, indent=2))
if __name__ == "__main__":
main()
pipeline/rtmw_risk_pipeline.py
-390
View File
@@ -1,390 +0,0 @@
"""Lightweight causal door-risk inference using only YOLO11s and RTMW-L."""
from __future__ import annotations
import argparse
import json
import sys
import time
from pathlib import Path
import cv2
import numpy as np
PACKAGE_DIR = Path(__file__).resolve().parents[1]
PIPELINE_DIR = Path(__file__).resolve().parent
for import_dir in (PIPELINE_DIR, PACKAGE_DIR):
if str(import_dir) not in sys.path:
sys.path.insert(0, str(import_dir))
from door_mask_risk import DoorMaskSequence, DoorRiskEvaluator # noqa: E402
from door_zone_risk import ( # noqa: E402
DoorZoneAssigner,
DoorZoneConfig,
build_door_risks,
)
from live_pose_engine import LivePoseEngine # noqa: E402
def parse_image_size(value: str) -> int | tuple[int, int]:
parts = value.lower().replace(",", "x").split("x")
if len(parts) == 1:
return int(parts[0])
if len(parts) == 2:
return int(parts[0]), int(parts[1])
raise argparse.ArgumentTypeError(
"image size must be one integer or HEIGHTxWIDTH, for example 960 or 544x960"
)
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--video", type=Path, required=True)
door_source = parser.add_mutually_exclusive_group(required=True)
door_source.add_argument("--door-mask-dir", type=Path)
door_source.add_argument(
"--constant-door-mask",
type=Path,
help="One X-AnyLabeling JSON or raster mask reused for every video frame.",
)
parser.add_argument("--output", type=Path, required=True)
parser.add_argument(
"--det-model", type=Path, default=PACKAGE_DIR / "models" / "yolo11s.onnx"
)
parser.add_argument(
"--pose-model", type=Path, default=PACKAGE_DIR / "models" / "rtmw-l-384.onnx"
)
parser.add_argument(
"--pose-jsonl",
type=Path,
help="Diagnostic replay input; bypasses live YOLO/RTMW model execution.",
)
parser.add_argument("--det-device", default="0")
parser.add_argument("--pose-device", default="cuda")
parser.add_argument(
"--pose-provider",
choices=("cuda", "tensorrt"),
default="cuda",
help="ONNX Runtime provider for RTMW.",
)
parser.add_argument(
"--pose-trt-cache",
type=Path,
default=PACKAGE_DIR / "models" / "rtmw_l_trt_cache",
)
parser.add_argument("--det-imgsz", type=parse_image_size, default=960)
parser.add_argument("--det-conf", type=float, default=0.30)
parser.add_argument("--det-roi-y-max-ratio", type=float, default=1.0)
parser.add_argument("--max-people", type=int, default=3)
parser.add_argument(
"--door-mask-format",
choices=("auto", "xanylabeling", "raster"),
default="auto",
)
parser.add_argument("--door-label", action="append")
parser.add_argument("--door-frame-offset", type=int, default=0)
parser.add_argument("--require-door-mask", action="store_true")
parser.add_argument("--risk-warning-relative", type=float, default=0.30)
parser.add_argument("--risk-release-relative", type=float, default=0.45)
parser.add_argument(
"--active-door",
choices=("front", "rear", "both", "none"),
default="both",
help="Door(s) whose CAN state is 0. Use front/rear for deployment simulation.",
)
parser.add_argument(
"--door-zone-config",
type=Path,
default=PACKAGE_DIR / "annotations" / "door_zone_multifeature.json",
help="Pose-primary fixed-camera front/rear calibration JSON.",
)
parser.add_argument("--max-frames", type=int, default=0)
parser.add_argument("--write-video", action="store_true")
parser.add_argument("--progress-every", type=int, default=100)
return parser.parse_args()
def load_pose_records(path: Path) -> dict[int, dict]:
records = {}
with path.open("r", encoding="utf-8") as handle:
for line in handle:
if not line.strip():
continue
record = json.loads(line)
records[int(record["frame_index"])] = record
return records
def percentile(values: list[float], p: float) -> float | None:
return round(float(np.percentile(values, p)), 3) if values else None
def timing_stats(values: list[float]) -> dict:
return {
"count": len(values),
"total": round(float(np.sum(values)), 3) if values else 0.0,
"mean": round(float(np.mean(values)), 3) if values else None,
"median": percentile(values, 50),
"p95": percentile(values, 95),
}
def visualization_skeletons(record: dict) -> list[dict]:
"""Return the exact fused keypoints and metadata consumed by draw_fused."""
boxes = record.get("boxes", [])
track_ids = [int(value) for value in record.get("track_ids", [])]
keypoints = record.get("temporal_selected_keypoints", [])
metadata = record.get("temporal_metadata", [])
if not (len(boxes) == len(keypoints) == len(metadata)):
return []
return [
{
"track_id": track_ids[position] if position < len(track_ids) else position + 1,
"person_index": position,
"layout": "coco_wholebody_133",
"keypoints": person_keypoints,
"metadata": metadata[position],
}
for position, person_keypoints in enumerate(keypoints)
]
def render_frame(
frame: np.ndarray,
record: dict,
frame_risk: dict,
door_mask: np.ndarray | None,
door_assignments: list[dict],
) -> np.ndarray:
canvas = frame.copy()
boxes = record.get("boxes", [])
keypoints = record.get("temporal_selected_keypoints", [])
metadata = record.get("temporal_metadata", [])
if len(boxes) == len(keypoints) == len(metadata) and boxes:
from fuse_rtmw_l_temporal import draw_fused
canvas = draw_fused(canvas, boxes, keypoints, metadata, "YOLO + RTMW risk")
else:
for position, box in enumerate(boxes):
x1, y1, x2, y2 = [int(round(float(value))) for value in box[:4]]
cv2.rectangle(canvas, (x1, y1), (x2, y2), (0, 220, 255), 2)
cv2.putText(
canvas,
f"P{position + 1}",
(x1 + 4, max(22, y1 + 24)),
cv2.FONT_HERSHEY_SIMPLEX,
0.7,
(0, 220, 255),
2,
cv2.LINE_AA,
)
zone_colors = {
"front": (0, 220, 255),
"rear": (255, 180, 0),
"transition": (0, 140, 255),
"unassigned": (150, 150, 150),
}
for assignment in door_assignments:
position = int(assignment["person_index"])
if not 0 <= position < len(boxes):
continue
box = boxes[position]
x1, y1 = int(round(float(box[0]))), int(round(float(box[1])))
zone = str(assignment["door_zone"])
cv2.putText(
canvas,
f"T{assignment['track_id']} {zone.upper()}",
(max(0, x1 + 4), max(24, y1 - 8)),
cv2.FONT_HERSHEY_SIMPLEX,
0.68,
zone_colors.get(zone, (160, 160, 160)),
2,
cv2.LINE_AA,
)
if door_mask is not None and np.any(door_mask):
overlay = np.zeros_like(canvas)
overlay[door_mask > 0] = (255, 180, 0)
canvas = cv2.addWeighted(canvas, 1.0, overlay, 0.22, 0)
score = frame_risk.get("risk_score")
score_text = "N/A" if score is None else f"{float(score):.3f}"
level = str(frame_risk.get("risk_level", "unavailable"))
color = (0, 0, 255) if level == "danger" else (0, 180, 255)
if level in ("clear", "unavailable", "inactive", "unknown"):
color = (40, 210, 80) if level == "clear" else (160, 160, 160)
cv2.rectangle(canvas, (12, 55), (490, 95), (0, 0, 0), -1)
cv2.putText(
canvas,
f"{frame_risk.get('risk_type', 'DOOR_RISK')} score={score_text} level={level}",
(24, 83),
cv2.FONT_HERSHEY_SIMPLEX,
0.68,
color,
2,
cv2.LINE_AA,
)
return canvas
def main() -> None:
args = parse_args()
args.output.mkdir(parents=True, exist_ok=True)
if args.pose_provider == "tensorrt" and not args.pose_device.lower().startswith("cuda"):
raise ValueError("--pose-provider tensorrt requires --pose-device cuda")
capture = cv2.VideoCapture(str(args.video.resolve()))
if not capture.isOpened():
raise RuntimeError(f"Cannot open video: {args.video}")
input_fps = float(capture.get(cv2.CAP_PROP_FPS)) or 30.0
width = int(capture.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT))
risk_evaluator = DoorRiskEvaluator(
DoorMaskSequence(
args.door_mask_dir,
mask_format=args.door_mask_format,
labels=tuple(args.door_label or ("door", "车门")),
frame_offset=args.door_frame_offset,
require_every_frame=args.require_door_mask,
constant_mask=args.constant_door_mask,
),
warning_relative=args.risk_warning_relative,
release_relative=args.risk_release_relative,
)
zone_assigner = DoorZoneAssigner(DoorZoneConfig.load(args.door_zone_config))
print(
f"Door masks: {risk_evaluator.sequence.source_format} "
+ (
f"(constant: {risk_evaluator.sequence.constant_path})"
if risk_evaluator.sequence.constant_path is not None
else f"({len(risk_evaluator.sequence.paths)} indexed frames)"
)
)
pose_records = load_pose_records(args.pose_jsonl) if args.pose_jsonl else None
load_started = time.perf_counter()
pose_engine = None if pose_records is not None else LivePoseEngine(args)
model_load_seconds = time.perf_counter() - load_started
writer = None
if args.write_video:
writer = cv2.VideoWriter(
str(args.output / "rtmw_risk_overlay.mp4"),
cv2.VideoWriter_fourcc(*"mp4v"),
input_fps,
(width, height),
)
frame_latencies = []
pose_latencies = []
risk_latencies = []
render_latencies = []
started_all = time.perf_counter()
frame_index = 0
output_path = args.output / "causal_results.jsonl"
with output_path.open("w", encoding="utf-8") as output_handle:
while True:
if args.max_frames > 0 and frame_index >= args.max_frames:
break
ok, frame = capture.read()
if not ok:
break
frame_started = time.perf_counter()
pose_started = time.perf_counter()
if pose_records is None:
record = pose_engine.infer(frame, frame_index)
else:
record = pose_records.get(frame_index)
if record is None:
raise RuntimeError(f"Pose JSONL has no record for frame {frame_index}")
pose_latencies.append((time.perf_counter() - pose_started) * 1000)
risk_started = time.perf_counter()
assignments = zone_assigner.update(
frame_index, record, (height, width)
)
raw_risk, door_mask = risk_evaluator.evaluate(
frame_index, record, (height, width)
)
door_risks, frame_risk = build_door_risks(
raw_risk, assignments, args.active_door
)
risk_latencies.append((time.perf_counter() - risk_started) * 1000)
output_record = {
"frame_index": frame_index,
"track_ids": record.get("track_ids", []),
"boxes": record.get("boxes", []),
"skeletons": visualization_skeletons(record),
"active_door": args.active_door,
"door_assignments": assignments,
"door_risks": door_risks,
"risk": frame_risk,
"ungated_risk": raw_risk,
}
output_handle.write(json.dumps(output_record, ensure_ascii=False) + "\n")
if writer is not None:
render_started = time.perf_counter()
writer.write(
render_frame(frame, record, frame_risk, door_mask, assignments)
)
render_latencies.append((time.perf_counter() - render_started) * 1000)
frame_latencies.append((time.perf_counter() - frame_started) * 1000)
if args.progress_every > 0 and frame_index % args.progress_every == 0:
recent = frame_latencies[-min(len(frame_latencies), args.progress_every) :]
print(
f"frame={frame_index} recent_ms={np.mean(recent):.1f} "
f"risk={frame_risk['risk_level']}",
flush=True,
)
frame_index += 1
capture.release()
if writer is not None:
writer.release()
elapsed = time.perf_counter() - started_all
summary = {
"mode": "yolo_rtmw_door_risk_multifeature_v9_10",
"video": str(args.video.resolve()),
"pose_source": "jsonl_diagnostic" if pose_records is not None else "resident_yolo_rtmw",
"models_loaded": [] if pose_records is not None else ["YOLO11s", "RTMW-L"],
"edgetam_loaded": False,
"risk_region_policy": (
"RTMW alarm-eligible hand hulls; YOLO body box only when a person "
"has no usable RTMW hand region"
),
"frames": frame_index,
"input_fps": round(input_fps, 3),
"elapsed_seconds": round(elapsed, 3),
"throughput_fps": round(frame_index / elapsed, 3) if elapsed else None,
"model_load_seconds": round(model_load_seconds, 3),
"latency_ms": {
"total": timing_stats(frame_latencies),
"pose": timing_stats(pose_latencies),
"risk": timing_stats(risk_latencies),
"render": timing_stats(render_latencies),
},
"write_video": args.write_video,
"active_door": args.active_door,
"door_zone": zone_assigner.summary(),
"door_risk": risk_evaluator.summary(),
}
if pose_engine is not None:
summary["pose"] = {
"provider": args.pose_provider,
"isolated_inferences": pose_engine.isolated_inferences,
"timing_ms": {
key: timing_stats(values) for key, values in pose_engine.timings.items()
},
}
(args.output / "summary.json").write_text(
json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8"
)
print(json.dumps(summary, ensure_ascii=False, indent=2))
if __name__ == "__main__":
main()
pipeline/visualization.py
+116
View File
@@ -0,0 +1,116 @@
"""Runtime visualization helpers."""
from __future__ import annotations
import cv2
import numpy as np
def percentile(values: list[float], p: float) -> float | None:
return round(float(np.percentile(values, p)), 3) if values else None
def timing_stats(values: list[float]) -> dict:
return {
"count": len(values),
"total": round(float(np.sum(values)), 3) if values else 0.0,
"mean": round(float(np.mean(values)), 3) if values else None,
"median": percentile(values, 50),
"p95": percentile(values, 95),
}
def visualization_skeletons(record: dict) -> list[dict]:
"""Return the exact fused keypoints and metadata consumed by draw_fused."""
boxes = record.get("boxes", [])
track_ids = [int(value) for value in record.get("track_ids", [])]
keypoints = record.get("temporal_selected_keypoints", [])
metadata = record.get("temporal_metadata", [])
if not (len(boxes) == len(keypoints) == len(metadata)):
return []
return [
{
"track_id": track_ids[position] if position < len(track_ids) else position + 1,
"person_index": position,
"layout": "coco_wholebody_133",
"keypoints": person_keypoints,
"metadata": metadata[position],
}
for position, person_keypoints in enumerate(keypoints)
]
def render_frame(
frame: np.ndarray,
record: dict,
frame_risk: dict,
door_mask: np.ndarray | None,
door_assignments: list[dict],
) -> np.ndarray:
canvas = frame.copy()
boxes = record.get("boxes", [])
keypoints = record.get("temporal_selected_keypoints", [])
metadata = record.get("temporal_metadata", [])
if len(boxes) == len(keypoints) == len(metadata) and boxes:
from pose_fusion import draw_fused
canvas = draw_fused(canvas, boxes, keypoints, metadata, "YOLO + RTMW risk")
else:
for position, box in enumerate(boxes):
x1, y1, x2, y2 = [int(round(float(value))) for value in box[:4]]
cv2.rectangle(canvas, (x1, y1), (x2, y2), (0, 220, 255), 2)
cv2.putText(
canvas,
f"P{position + 1}",
(x1 + 4, max(22, y1 + 24)),
cv2.FONT_HERSHEY_SIMPLEX,
0.7,
(0, 220, 255),
2,
cv2.LINE_AA,
)
zone_colors = {
"front": (0, 220, 255),
"rear": (255, 180, 0),
"transition": (0, 140, 255),
"unassigned": (150, 150, 150),
}
for assignment in door_assignments:
position = int(assignment["person_index"])
if not 0 <= position < len(boxes):
continue
box = boxes[position]
x1, y1 = int(round(float(box[0]))), int(round(float(box[1])))
zone = str(assignment["door_zone"])
cv2.putText(
canvas,
f"T{assignment['track_id']} {zone.upper()}",
(max(0, x1 + 4), max(24, y1 - 8)),
cv2.FONT_HERSHEY_SIMPLEX,
0.68,
zone_colors.get(zone, (160, 160, 160)),
2,
cv2.LINE_AA,
)
if door_mask is not None and np.any(door_mask):
overlay = np.zeros_like(canvas)
overlay[door_mask > 0] = (255, 180, 0)
canvas = cv2.addWeighted(canvas, 1.0, overlay, 0.22, 0)
score = frame_risk.get("risk_score")
score_text = "N/A" if score is None else f"{float(score):.3f}"
level = str(frame_risk.get("risk_level", "unavailable"))
color = (0, 0, 255) if level == "danger" else (0, 180, 255)
if level in ("clear", "unavailable", "inactive", "unknown"):
color = (40, 210, 80) if level == "clear" else (160, 160, 160)
cv2.rectangle(canvas, (12, 55), (490, 95), (0, 0, 0), -1)
cv2.putText(
canvas,
f"{frame_risk.get('risk_type', 'DOOR_RISK')} score={score_text} level={level}",
(24, 83),
cv2.FONT_HERSHEY_SIMPLEX,
0.68,
color,
2,
cv2.LINE_AA,
)
return canvas