File size: 2,165 Bytes
4bab068 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 | """
OpenCV drawing layer. Pure image-in / image-out — no model logic here,
so it's easy to unit test and reuse (e.g. for a CCTV batch pipeline later).
"""
import cv2
import numpy as np
BOX_COLOR = (255, 140, 0) # BGR - orange for detected objects
TEXT_COLOR = (255, 255, 255) # white
def _hex_to_bgr(hex_color: str):
hex_color = hex_color.lstrip("#")
r, g, b = tuple(int(hex_color[i:i + 2], 16) for i in (0, 2, 4))
return (b, g, r)
def draw_detections(image: np.ndarray, detections: list) -> np.ndarray:
"""Draw YOLO bounding boxes + class labels."""
out = image.copy()
for det in detections:
x1, y1, x2, y2 = [int(v) for v in det["box"]]
label = f'{det["class"]} {det["conf"]:.0%}'
cv2.rectangle(out, (x1, y1), (x2, y2), BOX_COLOR, 2)
(tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
cv2.rectangle(out, (x1, y1 - th - 8), (x1 + tw + 6, y1), BOX_COLOR, -1)
cv2.putText(out, label, (x1 + 3, y1 - 5),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, TEXT_COLOR, 1, cv2.LINE_AA)
return out
def draw_risk_badge(image: np.ndarray, risk_score: int, risk_level: str, color_hex: str) -> np.ndarray:
"""Draw a top-left risk badge: e.g. 'RISK 92/100 - CRITICAL'."""
out = image.copy()
h, w = out.shape[:2]
color_bgr = _hex_to_bgr(color_hex)
text = f"RISK {risk_score}/100 - {risk_level}"
font_scale = max(0.6, w / 1000)
thickness = 2
(tw, th), baseline = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, font_scale, thickness)
pad = 12
overlay = out.copy()
cv2.rectangle(overlay, (0, 0), (tw + pad * 2, th + baseline + pad * 2), color_bgr, -1)
out = cv2.addWeighted(overlay, 0.85, out, 0.15, 0)
cv2.putText(out, text, (pad, th + pad),
cv2.FONT_HERSHEY_SIMPLEX, font_scale, TEXT_COLOR, thickness, cv2.LINE_AA)
return out
def annotate(image_bgr: np.ndarray, detections: list, risk_score: int,
risk_level: str, color_hex: str) -> np.ndarray:
out = draw_detections(image_bgr, detections)
out = draw_risk_badge(out, risk_score, risk_level, color_hex)
return out
|