117 lines
3.9 KiB
Python
117 lines
3.9 KiB
Python
"""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
|