65 lines
2.1 KiB
Python
65 lines
2.1 KiB
Python
"""OpenCV overlays for tracked people and on-screen counters."""
|
|
from __future__ import annotations
|
|
|
|
from typing import List, Tuple
|
|
|
|
import cv2
|
|
|
|
from saqr.core.compliance import split_wearing_missing
|
|
from saqr.core.detection import STATUSES
|
|
|
|
GREEN = (0, 200, 0)
|
|
YELLOW = (0, 200, 255)
|
|
RED = (0, 0, 220)
|
|
WHITE = (255, 255, 255)
|
|
BLACK = (0, 0, 0)
|
|
GRAY = (120, 120, 120)
|
|
CYAN = (200, 200, 0)
|
|
|
|
|
|
def status_color(status: str) -> Tuple:
|
|
return {"SAFE": GREEN, "PARTIAL": YELLOW, "UNSAFE": RED}.get(status, GRAY)
|
|
|
|
|
|
def draw_track(frame, track):
|
|
x1, y1, x2, y2 = track.bbox
|
|
color = status_color(track.status)
|
|
cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
|
|
|
|
wearing, missing, _unknown = split_wearing_missing(track.items)
|
|
line1 = f"ID {track.track_id:04d} {track.status}"
|
|
w_str = ", ".join(wearing) if wearing else "none"
|
|
m_str = ", ".join(missing) if missing else "-"
|
|
line2 = f"W:{w_str} M:{m_str}"
|
|
|
|
(tw1, th1), _ = cv2.getTextSize(line1, cv2.FONT_HERSHEY_SIMPLEX, 0.55, 1)
|
|
(tw2, th2), _ = cv2.getTextSize(line2, cv2.FONT_HERSHEY_SIMPLEX, 0.40, 1)
|
|
tw = max(tw1, tw2) + 8
|
|
total_h = th1 + th2 + 12
|
|
y_top = max(0, y1 - total_h - 2)
|
|
|
|
cv2.rectangle(frame, (x1, y_top), (x1 + tw, y1), color, -1)
|
|
cv2.putText(frame, line1, (x1 + 4, y_top + th1 + 2),
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.55, WHITE, 1, cv2.LINE_AA)
|
|
cv2.putText(frame, line2, (x1 + 4, y_top + th1 + th2 + 8),
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.40, WHITE, 1, cv2.LINE_AA)
|
|
|
|
|
|
def draw_counters(frame, tracks: List, fps: float):
|
|
counts = {s: 0 for s in STATUSES}
|
|
for t in tracks:
|
|
counts[t.status] += 1
|
|
|
|
lines = [
|
|
(f"FPS: {fps:.1f}", WHITE),
|
|
(f"SAFE {counts['SAFE']}", GREEN),
|
|
(f"PARTIAL {counts['PARTIAL']}", YELLOW),
|
|
(f"UNSAFE {counts['UNSAFE']}", RED),
|
|
(f"TRACKS {len(tracks)}", CYAN),
|
|
]
|
|
y = 24
|
|
for text, color in lines:
|
|
cv2.putText(frame, text, (10, y), cv2.FONT_HERSHEY_SIMPLEX, 0.7, BLACK, 4, cv2.LINE_AA)
|
|
cv2.putText(frame, text, (10, y), cv2.FONT_HERSHEY_SIMPLEX, 0.7, color, 2, cv2.LINE_AA)
|
|
y += 28
|