391 lines
14 KiB
Python
391 lines
14 KiB
Python
"""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()
|