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
orin/README.md
+17 -23
View File
@@ -1,39 +1,33 @@
# Orin 可直接运行包 — use-can
# Orin 直接运行包
这个目录可以单独复制到 Orin 的任意目录，不依赖兄弟目录 0920/0923，不需要重新编译。
包含板端编译的 TensorRT/CUDA 原生库、YOLO/RTMW engine、Python 风险代码、CAN USB 库、标定文件和原有 GMSL 驱动。
环境：Jetson AGX Orin，JetPack 5.1.2 / L4T R35.4.1，Linux 5.10.120-tegra，CUDA 11.4，TensorRT 8.5.x，系统 OpenCV 4.5.4 库；Python 3.8+，numpy、cv2。系统 CUDA/TensorRT/OpenCV 运行库由 JetPack 提供，未重复打包。
适用 Jetson AGX Orin / JetPack 5.1.2 (R35.4.1)、CUDA 11.4、TensorRT 8.5.x、OpenCV 4.5.4。
系统需要 Python 3.8+、numpy、cv2。无需重新编译，可单独复制整个目录。
```bash
cd orin
sha256sum -c SHA256SUMS
./scripts/initialize_camera.sh # 开机后无 /dev/video5 时执行；需要 sudo 权限
./run.sh # 默认 CAN 左侧、无可视化，结果保存在内存
./run.sh --visualize # 可选浏览器预览，默认端口 8765
./scripts/initialize_camera.sh # 开机后无 /dev/video5 时执行，需要 sudo
./run.sh # 默认 CAN 模式，结果保存在内存
./run.sh --visualize # 可选浏览器预览，http://设备IP:8765
```
Ctrl+C / SIGTERM 退出。仅有一个进程可以占用同一个摄像头。
只运行一个摄像头采集进程，Ctrl+C / SIGTERM 停止。
```bash
./run.sh --no-can --active-door both --duration 10
./run.sh --save-results --output ./outputs/session
./run.sh --visualize --save-results --save-snapshots --output ./outputs/session
```
- data[0..3]：左前、右前、左后、右后；`--can-side left/right` 选择摄像头侧。
- 未收到有效 CAN 或超时：CAN_UNKNOWN，score=null；不发送执行器控制报文。
- `--save-results` 保存 JSONL；`--save-snapshots` 保存快照；`--write-video` 保存叠加视频。默认均关闭。
- `--output PATH` 指定输出目录；`--duration 30` 限时运行。完整参数见 `./run.sh --help`。
- 标定参数仍需匹配实际机位，默认 calibration_verified=false。
CAN 0x391 的 data[0..3] 顺序为左前、右前、左后、右后。`--can-side right` 切换右侧。未收到有效 CAN 或报文过期时输出 CAN_UNKNOWN，不发送执行器控制报文。
标定文件仍为研发机位，默认 calibration_verified=false。
同进程下游接口：
同进程下游示例（先停止其他采集进程）：
```python
from run_camera import CameraPipeline, parse_args
with CameraPipeline(parse_args([])) as pipeline:
result = pipeline.wait_for_result(timeout=10)
result = pipeline.latest_result # 等价于 pipeline.latest()
result = pipeline.latest_result # 等价于 pipeline.latest()
if result is not None:
print(result['risk'], result['door_risks'])
print(result['risk'], result['door_risks'], result['skeletons'])
```
持续读取示例：`python3 examples/read_latest.py`（先停止其他采集进程）。`latest_result` 是单个只读快照；它不是跨进程共享变量。
TensorRT engine 绑定目标 GPU/TRT 环境，不能直接用于 Thor、其他 JetPack 版本或桌面 GPU。
变量是线程安全的最新只读快照，读取不会消费结果，不是跨进程共享变量。
系统 CUDA/TensorRT/OpenCV 运行库由 JetPack 提供。引擎和驱动不能直接用于其他 GPU 或不匹配的内核。
orin/SHA256SUMS
+11 -13
View File
@@ -1,27 +1,25 @@
4171567ecbea2f15b3c945b8d8363362d9f422b8c732d33ee1e15461b40014da README.md
b6295cf4b0705c5c3d83c9ba35e41b2a0b865ea67742c09cfd1690f263f556d0 README.md
9ba768d559fd3b64e833b50c62c70f2adfd6c169620aea5fba32cfb1b4c5b0e9 annotations/constant_door_mask.json
3544fb1c7b1b27b39db27d8f849d5e8652f5d09a950a59c5dd0c7703086b4f6b annotations/door_zone_multifeature.json
6d5e874fd27fcfd75d303d7e35d9672028af71041f52dcc82a1e83676faefa65 bin/drivers/max9295.ko
8a78d04ef8be44e9b62a5e4ee315d3d6d679e34aedff02a234bbf4d40e222933 bin/drivers/max9296.ko
3da628587477cc35d5e3f8aee86e718e2ddf16808e893e00ba6e164b2307e843 bin/drivers/sgx-yuv-gmsl2.ko
cd75632beaaf81df2f59ce3e83fb335c21cc745068ef2b721c74cc3b2b03db26 build/libvision_native.so
49cc1067c1fa9f82a3d57dbfa77a55a54e5e5062ceae7ae84144b77745db79cf examples/read_latest.py
5ba7e4f0371a411cc1a6dfd1275e3b824032a14d07a9ccd7ca3d5a6630b322ef build/libvision_native.so
8dae8c557deadbe6bd7241c0dcd7025cb007077ee39a79b115e953a2b22cd9c1 lib/libusb-1.0.so
ab8b7f7249d5fb73d1d33d14fa7532a522088ffb6c25f37ba334daf4c6b4d313 lib/libusbcanfd.so
ebcd09e8c2e35008adbb0c1d222e580b17e0df1c8093df3481ba3374cd9a8174 models/engines/rtmw_l_384.engine
f6a65b926ac1144484f4aa93cd26d7997d32d62310ab93fe867f2bd4412222e9 models/engines/yolo11s.engine
fd3ae92255abf95ef037a1060ee75897f9705b4040f9394c3130f6ab860667e1 pipeline/compare_rtmw_ml.py
6d6d14a972bc4e5e9f8256514169e73ba0392a301cacd1221eaa26146aad00d2 pipeline/door_mask_risk.py
46d70e1e90a1f36d544b7492ef1fbf08821377a1c460a6683f42d4d7006055c7 pipeline/door_zone_risk.py
bc74a400ce5c5608c2fa49a69ccd8b41ec8633cd6a8239d9d594b8d9e5c3732a pipeline/fuse_rtmw_l_temporal.py
bcf2cc99dd27c8749801a5b743656767d9d64ef03d568d227cb654c573d70b00 pipeline/latest_result.py
cc8ea36d543896cd90f7a00822f06fc67771a1811859a4087995ac372df30276 pipeline/live_pose_engine.py
da299c11bd25127a09ffa738876f1de6e0d53d27add2840aeb6092ab0fe7016b pipeline/native_backend.py
81088b8bc072162bec13f57c65eb0d396e6d1e1074a83f92bcf29e393e27e11f pipeline/occlusion_aware_pose_v2.py
059e2a67c17ce63354986d3242a5d361103ee4241c5391f18140eeb19e47c939 pipeline/predict_optimized_pose_videos.py
bf8bb754cd841a52b4410f958a651b005de36ac4d17cc0992fbe1c90ed30caca pipeline/rtmw_l_person_mask_ablation.py
1bb58d851ba5d684527b8d1e1940db6241d2b0b2862687fa3b151899a65320d2 pipeline/rtmw_risk_pipeline.py
294cdc11aad0e5475275fd8dece2057e229b9bdedba69c26afa62176fd51a4e9 pipeline/native_backend.py
8b8c5b8c9c76c9bd2e1efa6cba60be337a18c9366e8f748038465e9e2fba2672 pipeline/pose_engine.py
93d3323ea431ed1d3c8397cfed7d800d65ec6d7700690487335861945040a486 pipeline/pose_fusion.py
2ebf58114835ba7c4779d4b66aeff65b8a94107128989c0336de556a4cf95fda pipeline/pose_geometry.py
f32fa5bebf893a3dc6b9ccfaa09a8e3a153ad5b36dde5373fdc561f10f9e526b pipeline/pose_occlusion.py
1a3f1f55d7b8a3f0ac96e69bf77f02f82017e8970adef96181d2bfade26e5bd5 pipeline/pose_tracking.py
dcf3cce5562ea52400ffa9842fdddaadac1efd672dabb0ed12ce34f20f779f37 pipeline/visualization.py
8a10e035334ed97b7480f558aff7b0464c255ac6138dd99f55a08ade68853e0e run.sh
4393f367b4ac7e5a00908043ffd21d5573cc12ed19b956fd101897ada22afd45 run_camera.py
23e9af5e82999889fe77c762a77c77036d1b1b699525ff2bd5ca8a418183ba50 run_camera.py
11b5fbc3c17c0610ec0514a0a26cf5d92588d15c0e50e7086d9db5a5c0306e99 scripts/initialize_camera.sh
4cc40e244495e1de60a0429a8fc58cdf383ccda471955e53360a1c36a278c822 scripts/initialize_camera_impl.sh
bce5a22d49633454ed2c6bc92b608edf3b718af2424c283ed8f56ba93987200a scripts/initialize_camera_impl.sh
orin/build/libvision_native.so
BIN
View File
Binary file not shown.
orin/examples/read_latest.py
-36
View File
@@ -1,36 +0,0 @@
#!/usr/bin/env python3
"""Same-process downstream example. Run on Orin, which owns the camera/GPU."""
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from run_camera import CameraPipeline, parse_args
def main():
# No visualization or disk output by default. CLI flags can enable either.
pipeline = CameraPipeline(parse_args())
last_frame = -1
try:
pipeline.start()
while True:
# Reading the variable is enough for a consumer with its own clock:
# result = pipeline.latest_result # equivalent to pipeline.latest()
# This sample waits to avoid busy polling; still only the latest
# result is returned, without queuing/acknowledging individual frames.
result = pipeline.wait_for_result(last_frame, timeout=1.0)
if result is not None:
last_frame = result['frame_index']
print(last_frame, result['risk']['risk_type'],
result['risk']['risk_score'], flush=True)
# Downstream code uses result['door_risks'], ['skeletons'], ['can'].
elif not pipeline.running():
break
except KeyboardInterrupt:
pass
finally:
pipeline.request_stop()
pipeline.join() # Also propagates worker errors to the caller.
if __name__ == '__main__':
main()
orin/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()
orin/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
orin/pipeline/occlusion_aware_pose_v2.py
-1130
View File
File diff suppressed because it is too large Load Diff
orin/pipeline/live_pose_engine.py → orin/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.
orin/pipeline/fuse_rtmw_l_temporal.py → orin/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()
orin/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
orin/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
orin/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
orin/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()
orin/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()
orin/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()
orin/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
orin/run_camera.py
+4 -3
View File
@@ -16,11 +16,12 @@ import cv2
import numpy as np
ROOT=Path(__file__).resolve().parent
sys.path.insert(0,str(ROOT/'pipeline'))
from native_backend import NativeBackend, NativePoseEngine
from native_backend import NativeBackend
from pose_engine import NativePoseEngine
from latest_result import LatestResult, result_to_dict
from door_mask_risk import DoorMaskSequence, DoorRiskEvaluator
from door_zone_risk import DoorZoneConfig, DoorZoneAssigner, build_door_risks
from rtmw_risk_pipeline import render_frame, visualization_skeletons, timing_stats
from visualization import render_frame, visualization_skeletons, timing_stats
def select_active_door(mode, snapshot, side='left', timeout_ms=1000):
@@ -121,7 +122,7 @@ class Preview:
def parse_args(argv=None):
p=argparse.ArgumentParser(description=__doc__)
p.add_argument('--video',type=Path,help='Diagnostic video replay through the same TensorRT/risk chain');p.add_argument('--camera',default='/dev/video5');p.add_argument('--width',type=int,default=1920);p.add_argument('--height',type=int,default=1536);p.add_argument('--fps',type=int,default=30)
p.add_argument('--engines',type=Path,default=ROOT/'models/engines')
p.add_argument('--engines',type=Path,default=ROOT/'orin/models/engines' if (ROOT/'orin').is_dir() else ROOT/'models/engines')
p.add_argument('--constant-door-mask',type=Path,default=ROOT/'annotations/constant_door_mask.json')
p.add_argument('--door-zone-config',type=Path,default=ROOT/'annotations/door_zone_multifeature.json')
p.add_argument('--active-door',choices=['can','front','rear','both','none'],default='can')
orin/scripts/initialize_camera_impl.sh
+2 -2
View File
@@ -1,8 +1,8 @@
#!/usr/bin/env bash
# Use the previously deployed 0920 drivers; do not force module compatibility.
# Initialize the camera using the packaged Orin drivers.
set -euo pipefail
root="$(cd "$(dirname "$0")/.." && pwd)"
driver_dir="${DRIVER_DIR:-$root/../vision-anti-pinch-0920/orin/bin/drivers}"
driver_dir="${DRIVER_DIR:-$root/orin/bin/drivers}"
device="${1:-/dev/video5}"
kernel="$(uname -r)"
for driver in max9295 max9296 sgx-yuv-gmsl2; do