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