| import ctypes |
| import glob |
| import os |
| import sys |
| import time |
| from pathlib import Path |
| from typing import Dict, List, Optional, Tuple |
| import cv2 |
| import numpy as np |
|
|
| |
| def _preload_nvidia_cuda_libs(): |
| site_dirs = [p for p in sys.path if "site-packages" in p] |
| lib_dirs = [] |
| for sd in site_dirs: |
| for cand in sorted(glob.glob(os.path.join(sd, "nvidia", "*", "lib"))): |
| if os.path.isdir(cand): |
| lib_dirs.append(cand) |
| if lib_dirs: |
| cur_ld = os.environ.get("LD_LIBRARY_PATH", "") |
| os.environ["LD_LIBRARY_PATH"] = ":".join(lib_dirs + ([cur_ld] if cur_ld else [])) |
| for d in lib_dirs: |
| for so_file in sorted(glob.glob(os.path.join(d, "*.so*"))): |
| try: |
| ctypes.CDLL(so_file, mode=ctypes.RTLD_GLOBAL) |
| except Exception: |
| pass |
|
|
| _preload_nvidia_cuda_libs() |
|
|
| try: |
| import onnxruntime as ort |
| HAVE_ONNX = True |
| except ImportError: |
| HAVE_ONNX = False |
|
|
| MODEL_DIR = Path(__file__).resolve().parent / "models" |
| YOLO_MODEL_PATH = MODEL_DIR / "yolov8n.onnx" |
| POSE_MODEL_PATH = ( |
| MODEL_DIR / "pose" / "onnx" / "model.onnx" |
| if (MODEL_DIR / "pose" / "onnx" / "model.onnx").exists() |
| else MODEL_DIR / "pose" / "onnx" / "model_quantized.onnx" |
| ) |
|
|
| |
| COCO_CLASSES = [ |
| "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", |
| "traffic light", "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", |
| "dog", "horse", "sheep", "cow", "elephant", "bear", "zebra", "giraffe", "backpack", |
| "umbrella", "handbag", "tie", "suitcase", "frisbee", "skis", "snowboard", "sports ball", |
| "kite", "baseball bat", "baseball glove", "skateboard", "surfboard", "tennis racket", |
| "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple", |
| "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", |
| "couch", "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", |
| "keyboard", "cell phone", "microwave", "oven", "toaster", "sink", "refrigerator", "book", |
| "clock", "vase", "scissors", "teddy bear", "hair drier", "toothbrush" |
| ] |
|
|
| |
| KPT_NAMES = [ |
| "nose", "left_eye", "right_eye", "left_ear", "right_ear", |
| "left_shoulder", "right_shoulder", "left_elbow", "right_elbow", |
| "left_wrist", "right_wrist", "left_hip", "right_hip", |
| "left_knee", "right_knee", "left_ankle", "right_ankle" |
| ] |
|
|
| SKELETON_PAIRS = [ |
| (0, 1), (0, 2), (1, 3), (2, 4), |
| (5, 6), (5, 7), (7, 9), (6, 8), (8, 10), |
| (5, 11), (6, 12), (11, 12), |
| (11, 13), (13, 15), (12, 14), (14, 16) |
| ] |
|
|
| |
| COLOR_PALETTE = { |
| "person": (50, 225, 120), |
| "head / face": (0, 235, 255), |
| "left hand": (255, 200, 0), |
| "right hand": (255, 200, 0), |
| "left arm": (255, 175, 50), |
| "right arm": (255, 175, 50), |
| "torso / chest": (80, 220, 120), |
| "chest / torso": (80, 220, 120), |
| "hips / pelvic": (180, 105, 255), |
| "pelvic / hips": (180, 105, 255), |
| "left leg": (255, 120, 200), |
| "right leg": (255, 120, 200), |
| "left foot": (0, 165, 255), |
| "right foot": (0, 165, 255), |
| "cell phone": (255, 180, 0), |
| "laptop": (255, 140, 0), |
| "bottle": (180, 180, 180), |
| "cup": (200, 150, 255), |
| "chair": (255, 100, 180), |
| "cat": (0, 180, 255), |
| "dog": (0, 165, 255) |
| } |
| DEFAULT_COLOR = (0, 215, 255) |
|
|
| def compute_box_iou(boxA: List[int], boxB: List[int]) -> float: |
| xA = max(boxA[0], boxB[0]) |
| yA = max(boxA[1], boxB[1]) |
| xB = min(boxA[2], boxB[2]) |
| yB = min(boxA[3], boxB[3]) |
| inter = max(0, xB - xA) * max(0, yB - yA) |
| areaA = max(0, boxA[2] - boxA[0]) * max(0, boxA[3] - boxA[1]) |
| areaB = max(0, boxB[2] - boxB[0]) * max(0, boxB[3] - boxB[1]) |
| union = float(areaA + areaB - inter) |
| return inter / union if union > 0 else 0.0 |
|
|
| class TrackedBox: |
| def __init__(self, track_id: str, base_label: str, label: str, confidence: float, box: List[int]): |
| self.track_id = track_id |
| self.base_label = base_label |
| self.label = label |
| self.confidence = confidence |
| self.box = list(box) |
| self.missed_frames = 0 |
| self.hit_count = 1 |
|
|
| def update(self, label: str, confidence: float, new_box: List[int]): |
| alpha = 0.60 |
| self.label = label |
| self.box[0] = int(self.box[0] * (1 - alpha) + new_box[0] * alpha) |
| self.box[1] = int(self.box[1] * (1 - alpha) + new_box[1] * alpha) |
| self.box[2] = int(self.box[2] * (1 - alpha) + new_box[2] * alpha) |
| self.box[3] = int(self.box[3] * (1 - alpha) + new_box[3] * alpha) |
| self.confidence = round(self.confidence * (1 - alpha) + confidence * alpha, 1) |
| self.missed_frames = 0 |
| self.hit_count += 1 |
|
|
| class ObjectDetector: |
| def __init__(self): |
| self.yolo_session: Optional[any] = None |
| self.pose_session: Optional[any] = None |
| self.providers: List[str] = ["CPUExecutionProvider"] |
| self.active_provider: str = "CPUExecutionProvider" |
| self.gpu_active: bool = False |
|
|
| self.loaded_yolo = False |
| self.loaded_pose = False |
|
|
| self.confidence_threshold = 0.45 |
| self.nms_threshold = 0.45 |
|
|
| |
| self.tracked_objects: List[TrackedBox] = [] |
| self.max_missed_frames = 3 |
| self.next_track_num = 1 |
|
|
| |
| self.latest_detections: List[dict] = [] |
| self.latest_body_parts: List[dict] = [] |
| self.latest_keypoints: List[dict] = [] |
| self.latest_telemetry: dict = { |
| "gpu_active": False, |
| "provider": "CPUExecutionProvider", |
| "keypoints_count": 0, |
| "objects_count": 0, |
| "inference_ms": 0.0 |
| } |
| self.last_detection_time: float = 0.0 |
|
|
| self._load_models() |
|
|
| def _load_models(self): |
| if not HAVE_ONNX: |
| return |
|
|
| opts = ort.SessionOptions() |
| opts.intra_op_num_threads = min(4, os.cpu_count() or 2) |
| opts.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL |
| opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL |
|
|
| available_providers = ort.get_available_providers() |
| if "CUDAExecutionProvider" in available_providers: |
| self.providers = ["CUDAExecutionProvider", "CPUExecutionProvider"] |
| else: |
| self.providers = ["CPUExecutionProvider"] |
|
|
| |
| if YOLO_MODEL_PATH.exists(): |
| try: |
| self.yolo_session = ort.InferenceSession(str(YOLO_MODEL_PATH), opts, providers=self.providers) |
| self.loaded_yolo = True |
| sess_provs = self.yolo_session.get_providers() |
| if sess_provs: |
| self.active_provider = sess_provs[0] |
| self.gpu_active = ("CUDA" in self.active_provider) |
| except Exception: |
| try: |
| self.yolo_session = ort.InferenceSession(str(YOLO_MODEL_PATH), opts, providers=["CPUExecutionProvider"]) |
| self.loaded_yolo = True |
| except Exception: |
| pass |
|
|
| |
| if POSE_MODEL_PATH.exists(): |
| try: |
| self.pose_session = ort.InferenceSession(str(POSE_MODEL_PATH), opts, providers=self.providers) |
| self.loaded_pose = True |
| sess_provs = self.pose_session.get_providers() |
| if sess_provs and "CUDA" in sess_provs[0]: |
| self.active_provider = sess_provs[0] |
| self.gpu_active = True |
| except Exception: |
| try: |
| self.pose_session = ort.InferenceSession(str(POSE_MODEL_PATH), opts, providers=["CPUExecutionProvider"]) |
| self.loaded_pose = True |
| except Exception: |
| pass |
|
|
| self.latest_telemetry["gpu_active"] = self.gpu_active |
| self.latest_telemetry["provider"] = self.active_provider |
|
|
| def detect(self, frame: np.ndarray) -> Tuple[List[dict], dict]: |
| """Detects 80 COCO objects, specific body parts, and 17 pose keypoints.""" |
| t0 = time.perf_counter() |
| if frame is None or not HAVE_ONNX: |
| return [], dict(self.latest_telemetry) |
|
|
| h, w = frame.shape[:2] |
| raw_items: List[dict] = [] |
| skeletons: List[dict] = [] |
|
|
| img_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) |
| resized_640 = cv2.resize(img_rgb, (640, 640), interpolation=cv2.INTER_LINEAR) |
| blob_640 = resized_640.transpose(2, 0, 1)[None].astype(np.float32) / 255.0 |
|
|
| x_factor = w / 640.0 |
| y_factor = h / 640.0 |
|
|
| |
| if self.loaded_yolo and self.yolo_session: |
| try: |
| outputs = self.yolo_session.run(['output0'], {'images': blob_640})[0][0].T |
| boxes = outputs[:, :4] |
| scores = outputs[:, 4:] |
| class_ids = np.argmax(scores, axis=1) |
| confidences = np.max(scores, axis=1) |
|
|
| valid_mask = confidences >= self.confidence_threshold |
| boxes = boxes[valid_mask] |
| confidences = confidences[valid_mask] |
| class_ids = class_ids[valid_mask] |
|
|
| cv_boxes = [] |
| for box in boxes: |
| cx, cy, bw, bh = box |
| cv_boxes.append([ |
| int((cx - bw / 2.0) * x_factor), |
| int((cy - bh / 2.0) * y_factor), |
| int(bw * x_factor), |
| int(bh * y_factor) |
| ]) |
|
|
| indices = cv2.dnn.NMSBoxes(cv_boxes, confidences.tolist(), self.confidence_threshold, self.nms_threshold) |
| if len(indices) > 0: |
| for idx in indices: |
| i = idx if isinstance(idx, (int, np.integer)) else idx[0] |
| cid = int(class_ids[i]) |
| if 0 <= cid < len(COCO_CLASSES): |
| label = COCO_CLASSES[cid] |
| bx = cv_boxes[i] |
| raw_items.append({ |
| "base_label": label, |
| "label": label, |
| "confidence": round(float(confidences[i]) * 100, 1), |
| "box": [max(0, bx[0]), max(0, bx[1]), min(w - 1, bx[0] + bx[2]), min(h - 1, bx[1] + bx[3])], |
| "is_body_part": False |
| }) |
| except Exception: |
| pass |
|
|
| |
| if self.loaded_pose and self.pose_session: |
| try: |
| pose_out = self.pose_session.run(['output0'], {'images': blob_640})[0][0].T |
| p_boxes = pose_out[:, :4] |
| p_scores = pose_out[:, 4] |
| p_kpts = pose_out[:, 5:].reshape(-1, 17, 3) |
|
|
| p_mask = p_scores >= 0.35 |
| p_boxes = p_boxes[p_mask] |
| p_scores = p_scores[p_mask] |
| p_kpts = p_kpts[p_mask] |
|
|
| p_cv_boxes = [] |
| for b in p_boxes: |
| cx, cy, bw, bh = b |
| p_cv_boxes.append([ |
| int((cx - bw / 2.0) * x_factor), |
| int((cy - bh / 2.0) * y_factor), |
| int(bw * x_factor), |
| int(bh * y_factor) |
| ]) |
|
|
| p_indices = cv2.dnn.NMSBoxes(p_cv_boxes, p_scores.tolist(), 0.35, 0.45) |
| for idx in p_indices: |
| i = idx if isinstance(idx, (int, np.integer)) else idx[0] |
| kpts = p_kpts[i] |
| kpts_scaled = [] |
| for kx, ky, kc in kpts: |
| kpts_scaled.append([int(kx * x_factor), int(ky * y_factor), float(kc)]) |
|
|
| skeletons.append(kpts_scaled) |
|
|
| |
| head_pts = [kpts_scaled[j] for j in range(5) if kpts_scaled[j][2] > 0.30] |
| if len(head_pts) >= 2: |
| xs = [pt[0] for pt in head_pts] |
| ys = [pt[1] for pt in head_pts] |
| pad_x = max(25, int((max(xs) - min(xs)) * 0.7)) |
| pad_y = max(30, int((max(ys) - min(ys)) * 1.0)) |
| face_box = [ |
| max(0, min(xs) - pad_x), |
| max(0, min(ys) - pad_y), |
| min(w - 1, max(xs) + pad_x), |
| min(h - 1, max(ys) + pad_y // 2) |
| ] |
| conf = round(float(np.mean([pt[2] for pt in head_pts])) * 100, 1) |
| raw_items.append({"base_label": "head / face", "label": "head / face", "confidence": conf, "box": face_box, "is_body_part": True}) |
|
|
| |
| if kpts_scaled[9][2] > 0.35: |
| x, y = kpts_scaled[9][:2] |
| pad = 32 |
| raw_items.append({ |
| "base_label": "left hand", |
| "label": "left hand", |
| "confidence": round(kpts_scaled[9][2] * 100, 1), |
| "box": [max(0, x - pad), max(0, y - pad), min(w - 1, x + pad), min(h - 1, y + pad)], |
| "is_body_part": True |
| }) |
|
|
| |
| if kpts_scaled[10][2] > 0.35: |
| x, y = kpts_scaled[10][:2] |
| pad = 32 |
| raw_items.append({ |
| "base_label": "right hand", |
| "label": "right hand", |
| "confidence": round(kpts_scaled[10][2] * 100, 1), |
| "box": [max(0, x - pad), max(0, y - pad), min(w - 1, x + pad), min(h - 1, y + pad)], |
| "is_body_part": True |
| }) |
|
|
| |
| torso_pts = [kpts_scaled[j] for j in [5, 6, 11, 12] if kpts_scaled[j][2] > 0.20] |
| if len(torso_pts) >= 2: |
| xs = [pt[0] for pt in torso_pts] |
| ys = [pt[1] for pt in torso_pts] |
| span_x = max(120, max(xs) - min(xs)) |
| tx1 = max(0, min(xs) - int(span_x * 0.18)) |
| tx2 = min(w - 1, max(xs) + int(span_x * 0.18)) |
| ty1 = max(0, min(ys) - 20) |
| ty2 = min(h - 1, max(max(ys) + 25, min(ys) + int(span_x * 0.85))) |
| if ty2 - ty1 < 90: |
| ty1 = max(0, ty2 - 110) |
| raw_items.append({ |
| "base_label": "chest / torso", |
| "label": "chest / torso", |
| "confidence": round(float(np.mean([pt[2] for pt in torso_pts])) * 100, 1), |
| "box": [tx1, ty1, tx2, ty2], |
| "is_body_part": True |
| }) |
|
|
| |
| hip_pts = [kpts_scaled[j] for j in [11, 12] if kpts_scaled[j][2] > 0.18] |
| if len(hip_pts) >= 1: |
| xs = [pt[0] for pt in hip_pts] |
| ys = [pt[1] for pt in hip_pts] |
| hx1, hy1 = max(0, min(xs) - 55), max(0, min(ys) - 35) |
| hx2, hy2 = min(w - 1, max(xs) + 55), min(h - 1, max(ys) + 65) |
| if hy2 - hy1 < 75: |
| hy1 = max(0, hy2 - 85) |
| raw_items.append({ |
| "base_label": "pelvic / hips", |
| "label": "pelvic / hips", |
| "confidence": round(float(np.mean([pt[2] for pt in hip_pts])) * 100, 1), |
| "box": [hx1, hy1, hx2, hy2], |
| "is_body_part": True |
| }) |
|
|
| |
| for arm_name, idxs in [("left arm", [5, 7, 9]), ("right arm", [6, 8, 10])]: |
| apts = [kpts_scaled[j] for j in idxs if kpts_scaled[j][2] > 0.30] |
| if len(apts) >= 2: |
| xs, ys = [p[0] for p in apts], [p[1] for p in apts] |
| raw_items.append({ |
| "base_label": arm_name, |
| "label": arm_name, |
| "confidence": round(float(np.mean([p[2] for p in apts])) * 100, 1), |
| "box": [max(0, min(xs) - 18), max(0, min(ys) - 18), min(w - 1, max(xs) + 18), min(h - 1, max(ys) + 18)], |
| "is_body_part": True |
| }) |
|
|
| |
| for leg_name, idxs in [("left leg", [11, 13, 15]), ("right leg", [12, 14, 16])]: |
| lpts = [kpts_scaled[j] for j in idxs if kpts_scaled[j][2] > 0.25] |
| if len(lpts) >= 2: |
| xs, ys = [p[0] for p in lpts], [p[1] for p in lpts] |
| raw_items.append({ |
| "base_label": leg_name, |
| "label": leg_name, |
| "confidence": round(float(np.mean([p[2] for p in lpts])) * 100, 1), |
| "box": [max(0, min(xs) - 22), max(0, min(ys) - 18), min(w - 1, max(xs) + 22), min(h - 1, max(ys) + 18)], |
| "is_body_part": True |
| }) |
|
|
| |
| for foot_name, pt_idx in [("left foot", 15), ("right foot", 16)]: |
| if kpts_scaled[pt_idx][2] > 0.22: |
| fx, fy = kpts_scaled[pt_idx][:2] |
| fx1, fy1 = max(0, fx - 36), max(0, fy - 28) |
| fx2, fy2 = min(w - 1, fx + 36), min(h - 1, fy + 38) |
| raw_items.append({ |
| "base_label": foot_name, |
| "label": foot_name, |
| "confidence": round(kpts_scaled[pt_idx][2] * 100, 1), |
| "box": [fx1, fy1, fx2, fy2], |
| "is_body_part": True |
| }) |
| except Exception: |
| pass |
|
|
| self._update_tracks(raw_items) |
| self.latest_keypoints = skeletons |
|
|
| formatted = [ |
| { |
| "id": t.track_id, |
| "base_label": t.base_label, |
| "label": t.label, |
| "confidence": t.confidence, |
| "box": list(t.box), |
| "center": [ |
| round((t.box[0] + t.box[2]) / (2.0 * w), 3), |
| round((t.box[1] + t.box[3]) / (2.0 * h), 3) |
| ] |
| } |
| for t in self.tracked_objects |
| if t.missed_frames == 0 or t.hit_count >= 2 |
| ] |
|
|
| elapsed_ms = round((time.perf_counter() - t0) * 1000.0, 2) |
| telemetry = { |
| "gpu_active": self.gpu_active, |
| "provider": self.active_provider, |
| "keypoints_count": len(skeletons) * 17, |
| "objects_count": len(formatted), |
| "inference_ms": elapsed_ms |
| } |
|
|
| self.latest_detections = formatted |
| self.latest_telemetry = telemetry |
| self.last_detection_time = time.time() |
|
|
| return formatted, telemetry |
|
|
| def _update_tracks(self, raw_items: List[dict]): |
| matched_track_idx = set() |
| matched_raw_idx = set() |
|
|
| for r_idx, r in enumerate(raw_items): |
| r_base = r.get("base_label", r["label"]) |
| best_iou = 0.30 |
| best_t_idx = -1 |
| for t_idx, track in enumerate(self.tracked_objects): |
| if t_idx in matched_track_idx or track.base_label != r_base: |
| continue |
| iou = compute_box_iou(track.box, r["box"]) |
| if iou > best_iou: |
| best_iou = iou |
| best_t_idx = t_idx |
|
|
| if best_t_idx >= 0: |
| self.tracked_objects[best_t_idx].update( |
| r["label"], |
| r["confidence"], |
| r["box"] |
| ) |
| matched_track_idx.add(best_t_idx) |
| matched_raw_idx.add(r_idx) |
|
|
| for r_idx, r in enumerate(raw_items): |
| if r_idx not in matched_raw_idx: |
| track_id = f"obj_{self.next_track_num}" |
| self.next_track_num += 1 |
| self.tracked_objects.append( |
| TrackedBox( |
| track_id, |
| r.get("base_label", r["label"]), |
| r["label"], |
| r["confidence"], |
| r["box"] |
| ) |
| ) |
|
|
| self.tracked_objects = [ |
| t for i, t in enumerate(self.tracked_objects) |
| if (i in matched_track_idx) or (t.missed_frames + 1 <= self.max_missed_frames) |
| ] |
| for i, t in enumerate(self.tracked_objects): |
| if i not in matched_track_idx: |
| t.missed_frames += 1 |
|
|
| def render_overlay( |
| self, |
| frame: np.ndarray, |
| detections: Optional[List[dict]] = None, |
| telemetry: Optional[dict] = None, |
| crop_meta: Optional[Tuple[int, int, int, int]] = None, |
| active_target_id: Optional[str] = None |
| ) -> np.ndarray: |
| """Renders HUD overlays, skeletal wireframe, and targeted object highlight directly on frame.""" |
| if frame is None: |
| return frame |
|
|
| out = frame.copy() |
| h, w = out.shape[:2] |
|
|
| dets = detections if detections is not None else self.latest_detections |
| telem = telemetry if telemetry is not None else self.latest_telemetry |
|
|
| |
| if crop_meta is not None: |
| cx1, cy1, cw, ch = crop_meta |
| sx = w / float(cw) |
| sy = h / float(ch) |
| else: |
| cx1, cy1, cw, ch = 0, 0, w, h |
| sx, sy = 1.0, 1.0 |
|
|
| |
| for kpts in self.latest_keypoints: |
| |
| for p1, p2 in SKELETON_PAIRS: |
| x1, y1, c1 = kpts[p1] |
| x2, y2, c2 = kpts[p2] |
| if c1 > 0.35 and c2 > 0.35: |
| mx1 = int((x1 - cx1) * sx) |
| my1 = int((y1 - cy1) * sy) |
| mx2 = int((x2 - cx1) * sx) |
| my2 = int((y2 - cy1) * sy) |
| bone_color = (0, 240, 255) if p1 < 5 else ((255, 200, 0) if p1 < 11 else (50, 255, 100)) |
| cv2.line(out, (mx1, my1), (mx2, my2), bone_color, 2, cv2.LINE_AA) |
|
|
| |
| for kx, ky, kc in kpts: |
| if kc > 0.35: |
| mx = int((kx - cx1) * sx) |
| my = int((ky - cy1) * sy) |
| if 0 <= mx < w and 0 <= my < h: |
| cv2.circle(out, (mx, my), 4, (0, 255, 255), -1, cv2.LINE_AA) |
| cv2.circle(out, (mx, my), 2, (20, 20, 20), -1, cv2.LINE_AA) |
|
|
| |
| for d in dets: |
| label = d["label"] |
| conf = d["confidence"] |
| orig_box = d["box"] |
| is_active_target = (active_target_id and d.get("id") == active_target_id) |
|
|
| x1 = int((orig_box[0] - cx1) * sx) |
| y1 = int((orig_box[1] - cy1) * sy) |
| x2 = int((orig_box[2] - cx1) * sx) |
| y2 = int((orig_box[3] - cy1) * sy) |
|
|
| if x2 < 0 or x1 >= w or y2 < 0 or y1 >= h: |
| continue |
|
|
| x1 = max(0, min(w - 1, x1)) |
| y1 = max(0, min(h - 1, y1)) |
| x2 = max(0, min(w - 1, x2)) |
| y2 = max(0, min(h - 1, y2)) |
|
|
| base_lbl = d.get("base_label", label) |
| if is_active_target: |
| color = (0, 100, 255) |
| else: |
| color = COLOR_PALETTE.get(base_lbl, DEFAULT_COLOR) |
| thick = 3 if is_active_target else 2 |
|
|
| |
| cv2.rectangle(out, (x1, y1), (x2, y2), color, 1, lineType=cv2.LINE_AA) |
| c_len = min(22, max(8, int((x2 - x1) * 0.2))) |
| cv2.line(out, (x1, y1), (x1 + c_len, y1), color, thick) |
| cv2.line(out, (x1, y1), (x1, y1 + c_len), color, thick) |
| cv2.line(out, (x2, y1), (x2 - c_len, y1), color, thick) |
| cv2.line(out, (x2, y1), (x2, y1 + c_len), color, thick) |
| cv2.line(out, (x1, y2), (x1 + c_len, y2), color, thick) |
| cv2.line(out, (x1, y2), (x1, y2 - c_len), color, thick) |
| cv2.line(out, (x2, y2), (x2 - c_len, y2), color, thick) |
| cv2.line(out, (x2, y2), (x2, y2 - c_len), color, thick) |
|
|
| |
| if is_active_target: |
| mcx = (x1 + x2) // 2 |
| mcy = (y1 + y2) // 2 |
| cv2.drawMarker(out, (mcx, mcy), (0, 100, 255), cv2.MARKER_CROSS, 20, 2) |
| cv2.putText(out, "[LOCKED TARGET]", (x1, max(15, y1 - 22)), cv2.FONT_HERSHEY_SIMPLEX, 0.42, (0, 100, 255), 1, cv2.LINE_AA) |
|
|
| |
| text = f"{label.upper()} {int(conf)}%" |
| (tw, th), _ = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.42, 1) |
| ty = max(y1, th + 8) |
| cv2.rectangle(out, (x1, ty - th - 5), (x1 + tw + 8, ty + 4), color, -1) |
| cv2.putText(out, text, (x1 + 4, ty - 2), cv2.FONT_HERSHEY_SIMPLEX, 0.42, (10, 10, 10), 1, cv2.LINE_AA) |
|
|
| |
| is_gpu = telem.get("gpu_active", self.gpu_active) |
| accel_tag = "T4 CUDA GPU" if is_gpu else "ONNX ENGINE" |
| badge_text = f"AI TRACKER: {accel_tag} ({len(dets)} TARGETS)" |
| badge_bg = (35, 120, 45) |
| text_color = (220, 255, 220) |
|
|
| (bw, bh), _ = cv2.getTextSize(badge_text, cv2.FONT_HERSHEY_SIMPLEX, 0.48, 1) |
| badge_x = w - bw - 20 |
| badge_y = 28 |
|
|
| cv2.rectangle(out, (badge_x - 10, badge_y - bh - 6), (badge_x + bw + 10, badge_y + 6), (15, 15, 15), -1) |
| cv2.rectangle(out, (badge_x - 8, badge_y - bh - 4), (badge_x + bw + 8, badge_y + 4), badge_bg, -1) |
| cv2.putText(out, badge_text, (badge_x, badge_y), cv2.FONT_HERSHEY_SIMPLEX, 0.48, text_color, 1, cv2.LINE_AA) |
|
|
| return out |
|
|
| def detect_and_draw(self, frame: np.ndarray, draw_boxes: bool = True) -> Tuple[np.ndarray, List[dict]]: |
| dets, telem = self.detect(frame) |
| if draw_boxes: |
| out = self.render_overlay(frame, dets, telem) |
| return out, dets |
| return frame, dets |
|
|
