Coverage for orchestr_ant_ion / yolo / ui / draw.py: 0%
210 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-03-19 08:36 +0000
1"""Overlay drawing utilities for detections and telemetry."""
3from __future__ import annotations
5from typing import TYPE_CHECKING
7import cv2
8import numpy as np
9from loguru import logger
11from orchestr_ant_ion.yolo.core.constants import CLASS_NAMES, COLORS
14if TYPE_CHECKING:
15 from collections import deque
17 from orchestr_ant_ion.pipeline.types import (
18 PerformanceMetrics,
19 SystemStats,
20 Track,
21 )
24def _track_color(track_id: int) -> tuple[int, int, int]:
25 """Return a deterministic color for a track id."""
26 r = (track_id * 97) % 255
27 g = (track_id * 57) % 255
28 b = (track_id * 17) % 255
29 return int(b), int(g), int(r)
32def draw_2d_running_map(
33 frame: np.ndarray,
34 tracks: dict[int, Track],
35 *,
36 map_size: int = 260,
37 margin: int = 10,
38) -> None:
39 """Draw a simple top-down minimap of tracked centroids with motion trails."""
40 if frame is None or frame.size == 0:
41 return
43 h, w = frame.shape[:2]
44 size = int(map_size)
45 x0 = max(margin, w - size - margin)
46 y0 = margin
47 x1 = min(w - margin, x0 + size)
48 y1 = min(h - margin, y0 + size)
50 cv2.rectangle(frame, (x0, y0), (x1, y1), (0, 0, 0), -1)
51 cv2.rectangle(frame, (x0, y0), (x1, y1), (100, 100, 100), 1)
52 cv2.putText(
53 frame,
54 "2D running (persons)",
55 (x0 + 8, y0 + 18),
56 cv2.FONT_HERSHEY_SIMPLEX,
57 0.5,
58 (200, 200, 200),
59 1,
60 cv2.LINE_AA,
61 )
63 for tick in (0.25, 0.5, 0.75):
64 xt = int(x0 + tick * (x1 - x0))
65 yt = int(y0 + tick * (y1 - y0))
66 cv2.line(frame, (xt, y0 + 24), (xt, y1), (35, 35, 35), 1)
67 cv2.line(frame, (x0, yt), (x1, yt), (35, 35, 35), 1)
69 usable_top = y0 + 26
70 usable_h = max(1, y1 - usable_top)
71 usable_w = max(1, x1 - x0)
73 for tid, track in tracks.items():
74 pts = list(track.points_norm)
75 if not pts:
76 continue
78 color = _track_color(tid)
79 poly: list[tuple[int, int]] = []
80 for xn, yn in pts:
81 px = int(x0 + np.clip(xn, 0.0, 1.0) * (usable_w - 1))
82 py = int(usable_top + np.clip(yn, 0.0, 1.0) * (usable_h - 1))
83 poly.append((px, py))
85 if len(poly) >= 2:
86 cv2.polylines(
87 frame,
88 [np.array(poly, dtype=np.int32)],
89 isClosed=False,
90 color=color,
91 thickness=2,
92 )
94 cx, cy = poly[-1]
95 cv2.circle(frame, (cx, cy), 4, color, -1)
96 cv2.putText(
97 frame,
98 str(tid),
99 (cx + 6, cy - 6),
100 cv2.FONT_HERSHEY_SIMPLEX,
101 0.45,
102 color,
103 1,
104 cv2.LINE_AA,
105 )
108def draw_cpu_process_history_plot(
109 frame: np.ndarray,
110 cpu_history: deque[float],
111 *,
112 x: int,
113 y: int,
114 w: int,
115 h: int,
116) -> None:
117 """Draw a simple 2D line chart of process CPU% history inside a rectangle."""
118 if frame is None or frame.size == 0:
119 return
121 if w < 30 or h < 20:
122 return
124 if cpu_history is None or len(cpu_history) < 1:
125 return
127 x0, y0 = int(x), int(y)
128 x1, y1 = int(x + w), int(y + h)
130 cv2.rectangle(frame, (x0, y0), (x1, y1), (0, 0, 0), -1)
131 cv2.rectangle(frame, (x0, y0), (x1, y1), (100, 100, 100), 1)
133 for tick in (0.25, 0.5, 0.75):
134 xt = int(x0 + tick * (x1 - x0))
135 yt = int(y0 + tick * (y1 - y0))
136 cv2.line(frame, (xt, y0), (xt, y1), (35, 35, 35), 1)
137 cv2.line(frame, (x0, yt), (x1, yt), (35, 35, 35), 1)
139 values = list(cpu_history)
140 n = len(values)
142 values = [float(np.clip(v, 0.0, 100.0)) for v in values]
144 poly: list[tuple[int, int]] = []
145 for i, value in enumerate(values):
146 xi = x0 + round(i * (w - 1) / max(1, n - 1))
147 yi = y1 - round((value / 100.0) * (h - 1))
148 poly.append((xi, yi))
150 last_v = values[-1]
151 color = get_color_by_percent(percent=last_v)
153 if len(poly) >= 2:
154 cv2.polylines(
155 frame,
156 [np.array(poly, dtype=np.int32)],
157 isClosed=False,
158 color=color,
159 thickness=2,
160 )
161 cv2.circle(frame, poly[-1], 3, color, -1)
163 label = f"Proc CPU history ({n}): {last_v:.1f}%"
164 cv2.putText(
165 frame,
166 label,
167 (x0 + 6, y0 - 6),
168 cv2.FONT_HERSHEY_SIMPLEX,
169 0.45,
170 (200, 200, 200),
171 1,
172 cv2.LINE_AA,
173 )
176def get_color_by_percent(
177 *, percent: float, invert: bool = False
178) -> tuple[int, int, int]:
179 """Return a color for a percentage value."""
180 if invert:
181 if percent > 50:
182 return (0, 255, 0)
183 if percent > 20:
184 return (0, 165, 255)
185 return (0, 0, 255)
186 if percent < 70:
187 return (0, 255, 0)
188 if percent < 90:
189 return (0, 165, 255)
190 return (0, 0, 255)
193def _draw_detection_boxes( # noqa: C901
194 frame: np.ndarray,
195 detections: list[dict[str, object]],
196 *,
197 debug_boxes: bool,
198) -> None:
199 h, w = frame.shape[:2]
201 def _clamp_int(value: float, low: int, high: int) -> int:
202 return max(low, min(high, int(value)))
204 for det in detections:
205 bbox_obj = det.get("bbox")
206 if not isinstance(bbox_obj, list | tuple) or len(bbox_obj) < 4:
207 continue
208 x1_obj, y1_obj, x2_obj, y2_obj = bbox_obj[:4]
209 if not isinstance(x1_obj, int | float):
210 continue
211 if not isinstance(y1_obj, int | float):
212 continue
213 if not isinstance(x2_obj, int | float):
214 continue
215 if not isinstance(y2_obj, int | float):
216 continue
217 x1 = float(x1_obj)
218 y1 = float(y1_obj)
219 x2 = float(x2_obj)
220 y2 = float(y2_obj)
221 class_id_obj = det.get("class_id", 0)
222 score_obj = det.get("score", 0.0)
224 if isinstance(class_id_obj, int | float | str):
225 class_id = int(class_id_obj)
226 else:
227 class_id = 0
229 score = float(score_obj) if isinstance(score_obj, int | float | str) else 0.0
231 x1 = _clamp_int(x1, 0, w - 1)
232 y1 = _clamp_int(y1, 0, h - 1)
233 x2 = _clamp_int(x2, 0, w - 1)
234 y2 = _clamp_int(y2, 0, h - 1)
235 if x2 <= x1 or y2 <= y1:
236 if debug_boxes:
237 logger.info(
238 "Skipping invalid bbox: {}",
239 [x1, y1, x2, y2],
240 )
241 continue
243 color = tuple(map(int, COLORS[class_id % len(COLORS)]))
244 if class_id < len(CLASS_NAMES):
245 label = f"{CLASS_NAMES[class_id]}: {score:.2f}"
246 else:
247 label = f"class {class_id}: {score:.2f}"
249 cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
251 (tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
252 cv2.rectangle(frame, (x1, y1 - th - 10), (x1 + tw, y1), color, -1)
253 cv2.putText(
254 frame,
255 label,
256 (x1, y1 - 5),
257 cv2.FONT_HERSHEY_SIMPLEX,
258 0.5,
259 (255, 255, 255),
260 1,
261 )
264def _draw_stats_panel(
265 frame: np.ndarray,
266 perf_metrics: PerformanceMetrics,
267 sys_stats: SystemStats,
268 proc_stats: dict[str, float | int],
269 camera_info: dict[str, object],
270) -> tuple[tuple[int, int, int], int, int, int]:
271 panel_height = 320
272 panel_width = 450
273 cv2.rectangle(frame, (5, 5), (panel_width, panel_height), (0, 0, 0), -1)
274 cv2.rectangle(frame, (5, 5), (panel_width, panel_height), (100, 100, 100), 1)
276 y_offset = 25
277 line_height = 22
279 backend_display = camera_info["backend"]
280 cv2.putText(
281 frame,
282 f"--- Capture: {backend_display} ---",
283 (10, y_offset),
284 cv2.FONT_HERSHEY_SIMPLEX,
285 0.5,
286 (150, 150, 255),
287 1,
288 )
289 y_offset += line_height
291 cv2.putText(
292 frame,
293 "--- Performance ---",
294 (10, y_offset),
295 cv2.FONT_HERSHEY_SIMPLEX,
296 0.5,
297 (150, 150, 150),
298 1,
299 )
300 y_offset += line_height
302 cv2.putText(
303 frame,
304 f"Camera Input: {perf_metrics.camera_fps:.1f} FPS",
305 (10, y_offset),
306 cv2.FONT_HERSHEY_SIMPLEX,
307 0.55,
308 (0, 255, 0),
309 1,
310 )
311 y_offset += line_height
313 cv2.putText(
314 frame,
315 f"Inference Latency: {perf_metrics.inference_ms:.1f} ms/frame",
316 (10, y_offset),
317 cv2.FONT_HERSHEY_SIMPLEX,
318 0.55,
319 (0, 255, 255),
320 1,
321 )
322 y_offset += line_height
324 budget_color = get_color_by_percent(percent=perf_metrics.frame_budget_percent)
325 cv2.putText(
326 frame,
327 f"Frame Budget Used: {perf_metrics.frame_budget_percent:.1f}%",
328 (10, y_offset),
329 cv2.FONT_HERSHEY_SIMPLEX,
330 0.55,
331 budget_color,
332 1,
333 )
334 y_offset += line_height
336 headroom = 100 - perf_metrics.frame_budget_percent
337 headroom_color = get_color_by_percent(percent=headroom, invert=True)
338 cv2.putText(
339 frame,
340 (
341 "GPU Headroom: "
342 f"{headroom:.0f}% (capacity: {perf_metrics.inference_capacity_fps:.0f} FPS)"
343 ),
344 (10, y_offset),
345 cv2.FONT_HERSHEY_SIMPLEX,
346 0.55,
347 headroom_color,
348 1,
349 )
350 y_offset += line_height + 5
352 cv2.putText(
353 frame,
354 "--- System ---",
355 (10, y_offset),
356 cv2.FONT_HERSHEY_SIMPLEX,
357 0.5,
358 (150, 150, 150),
359 1,
360 )
361 y_offset += line_height
363 cv2.putText(
364 frame,
365 f"System CPU: {sys_stats.cpu_percent:.1f}%",
366 (10, y_offset),
367 cv2.FONT_HERSHEY_SIMPLEX,
368 0.55,
369 get_color_by_percent(percent=sys_stats.cpu_percent),
370 1,
371 )
372 y_offset += line_height
374 cv2.putText(
375 frame,
376 (
377 f"System RAM: {sys_stats.ram_used_gb:.1f}/"
378 f"{sys_stats.ram_total_gb:.1f} GB ({sys_stats.ram_percent:.1f}%)"
379 ),
380 (10, y_offset),
381 cv2.FONT_HERSHEY_SIMPLEX,
382 0.55,
383 get_color_by_percent(percent=sys_stats.ram_percent),
384 1,
385 )
386 y_offset += line_height
388 if sys_stats.gpu_name != "N/A":
389 cv2.putText(
390 frame,
391 (
392 f"GPU Load: {sys_stats.gpu_percent:.0f}% | "
393 f"Temp: {sys_stats.gpu_temp_celsius:.0f}C | "
394 f"{sys_stats.gpu_power_watts:.0f}W"
395 ),
396 (10, y_offset),
397 cv2.FONT_HERSHEY_SIMPLEX,
398 0.55,
399 get_color_by_percent(percent=sys_stats.gpu_percent),
400 1,
401 )
402 y_offset += line_height
404 cv2.putText(
405 frame,
406 (
407 f"VRAM: {sys_stats.gpu_memory_used_gb:.1f}/"
408 f"{sys_stats.gpu_memory_total_gb:.1f} GB "
409 f"({sys_stats.gpu_memory_percent:.0f}%)"
410 ),
411 (10, y_offset),
412 cv2.FONT_HERSHEY_SIMPLEX,
413 0.55,
414 get_color_by_percent(percent=sys_stats.gpu_memory_percent),
415 1,
416 )
417 y_offset += line_height + 5
419 cv2.putText(
420 frame,
421 "--- Process ---",
422 (10, y_offset),
423 cv2.FONT_HERSHEY_SIMPLEX,
424 0.5,
425 (150, 150, 150),
426 1,
427 )
428 y_offset += line_height
430 cv2.putText(
431 frame,
432 (
433 f"CPU: {proc_stats['cpu_percent']:.1f}% | "
434 f"RAM: {proc_stats['memory_mb']:.0f}MB | "
435 f"Threads: {proc_stats['threads']}"
436 ),
437 (10, y_offset),
438 cv2.FONT_HERSHEY_SIMPLEX,
439 0.55,
440 get_color_by_percent(percent=min(proc_stats["cpu_percent"], 100)),
441 1,
442 )
444 bar_y = panel_height - 15
445 bar_width = panel_width - 20
446 bar_height = 8
447 return budget_color, bar_y, bar_width, bar_height
450def _draw_detection_panel(
451 frame: np.ndarray,
452 detections: list[dict[str, object]],
453 classification: dict[str, object] | None,
454) -> None:
455 det_text = f"Detections: {len(detections)}"
456 (tw, _th), _ = cv2.getTextSize(det_text, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)
457 cv2.putText(
458 frame,
459 det_text,
460 (frame.shape[1] - tw - 10, 30),
461 cv2.FONT_HERSHEY_SIMPLEX,
462 0.6,
463 (0, 255, 255),
464 2,
465 )
467 if classification is not None:
468 class_label = classification.get("label", "unknown")
469 class_score = classification.get("score", 0.0)
470 cls_text = f"Class: {class_label} ({class_score:.2f})"
471 (ctw, _cth), _ = cv2.getTextSize(cls_text, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)
472 cv2.putText(
473 frame,
474 cls_text,
475 (frame.shape[1] - ctw - 10, 60),
476 cv2.FONT_HERSHEY_SIMPLEX,
477 0.6,
478 (0, 255, 0),
479 2,
480 )
483def _draw_budget_bar(
484 frame: np.ndarray,
485 *,
486 budget_color: tuple[int, int, int],
487 bar_y: int,
488 bar_width: int,
489 bar_height: int,
490 frame_budget_percent: float,
491) -> None:
492 cv2.rectangle(
493 frame, (10, bar_y), (10 + bar_width, bar_y + bar_height), (50, 50, 50), -1
494 )
495 fill_width = int(bar_width * min(frame_budget_percent, 100) / 100)
496 cv2.rectangle(
497 frame,
498 (10, bar_y),
499 (10 + fill_width, bar_y + bar_height),
500 budget_color,
501 -1,
502 )
503 cv2.rectangle(
504 frame,
505 (10, bar_y),
506 (10 + bar_width, bar_y + bar_height),
507 (100, 100, 100),
508 1,
509 )
512def draw_detections(
513 frame: np.ndarray,
514 detections: list[dict[str, object]],
515 perf_metrics: PerformanceMetrics,
516 sys_stats: SystemStats,
517 proc_stats: dict[str, float | int],
518 camera_info: dict[str, object],
519 *,
520 cpu_history: deque[float] | None = None,
521 classification: dict[str, object] | None = None,
522 tracks: dict[int, Track] | None = None,
523 map_size: int = 260,
524 debug_boxes: bool = False,
525 show_stats_panel: bool = True,
526 show_detection_panel: bool = True,
527) -> np.ndarray:
528 """Draw bounding boxes, labels, FPS, and system stats on frame."""
529 _draw_detection_boxes(frame, detections, debug_boxes=debug_boxes)
531 budget_color = (0, 0, 0)
532 bar_y = 0
533 bar_width = 0
534 bar_height = 0
535 if show_stats_panel:
536 budget_color, bar_y, bar_width, bar_height = _draw_stats_panel(
537 frame,
538 perf_metrics,
539 sys_stats,
540 proc_stats,
541 camera_info,
542 )
544 if cpu_history is not None and show_stats_panel:
545 overlay_w = 320
546 overlay_h = 110
547 margin = 10
548 overlay_x = max(margin, frame.shape[1] - overlay_w - margin)
549 overlay_y = 40
550 draw_cpu_process_history_plot(
551 frame,
552 cpu_history,
553 x=overlay_x,
554 y=overlay_y,
555 w=overlay_w,
556 h=overlay_h,
557 )
559 if show_detection_panel:
560 _draw_detection_panel(frame, detections, classification)
562 if show_stats_panel:
563 _draw_budget_bar(
564 frame,
565 budget_color=budget_color,
566 bar_y=bar_y,
567 bar_width=bar_width,
568 bar_height=bar_height,
569 frame_budget_percent=perf_metrics.frame_budget_percent,
570 )
572 if tracks:
573 draw_2d_running_map(frame, tracks, map_size=map_size)
575 return frame