| """ |
| 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) |
| TEXT_COLOR = (255, 255, 255) |
|
|
|
|
| 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 |
|
|
