# Detection logic using OpenVINO models (YOLO, etc.) import os import sys import time import cv2 import numpy as np from pathlib import Path from typing import List, Dict, Tuple, Optional from red_light_violation_pipeline import RedLightViolationPipeline # --- Install required packages if missing --- try: import openvino as ov except ImportError: print("Installing openvino...") os.system('pip install --quiet "openvino>=2024.0.0"') import openvino as ov try: from ultralytics import YOLO except ImportError: print("Installing ultralytics...") os.system('pip install --quiet "ultralytics==8.3.0"') from ultralytics import YOLO try: import nncf except ImportError: print("Installing nncf...") os.system('pip install --quiet "nncf>=2.9.0"') import nncf # --- COCO dataset class names --- COCO_CLASSES = { 0: 'person', 1: 'bicycle', 2: 'car', 3: 'motorcycle', 4: 'airplane', 5: 'bus', 6: 'train', 7: 'truck', 8: 'boat', 9: 'traffic light', 10: 'fire hydrant', 11: 'stop sign', 12: 'parking meter', 13: 'bench', 14: 'bird', 15: 'cat', 16: 'dog', 17: 'horse', 18: 'sheep', 19: 'cow', 20: 'elephant', 21: 'bear', 22: 'zebra', 23: 'giraffe', 24: 'backpack', 25: 'umbrella', 26: 'handbag', 27: 'tie', 28: 'suitcase', 29: 'frisbee', 30: 'skis', 31: 'snowboard', 32: 'sports ball', 33: 'kite', 34: 'baseball bat', 35: 'baseball glove', 36: 'skateboard', 37: 'surfboard', 38: 'tennis racket', 39: 'bottle', 40: 'wine glass', 41: 'cup', 42: 'fork', 43: 'knife', 44: 'spoon', 45: 'bowl', 46: 'banana', 47: 'apple', 48: 'sandwich', 49: 'orange', 50: 'broccoli', 51: 'carrot', 52: 'hot dog', 53: 'pizza', 54: 'donut', 55: 'cake', 56: 'chair', 57: 'couch', 58: 'potted plant', 59: 'bed', 60: 'dining table', 61: 'toilet', 62: 'tv', 63: 'laptop', 64: 'mouse', 65: 'remote', 66: 'keyboard', 67: 'cell phone', 68: 'microwave', 69: 'oven', 70: 'toaster', 71: 'sink', 72: 'refrigerator', 73: 'book', 74: 'clock', 75: 'vase', 76: 'scissors', 77: 'teddy bear', 78: 'hair drier', 79: 'toothbrush' } # Traffic-related classes we're interested in (using standard COCO indices) TRAFFIC_CLASS_NAMES = COCO_CLASSES # --- Model Conversion and Quantization --- def convert_yolo_to_openvino(model_name: str = "yolo11x", half: bool = True) -> Path: """Convert YOLOv11x PyTorch model to OpenVINO IR format.""" pt_path = Path(f"{model_name}.pt") ov_dir = Path(f"{model_name}_openvino_model") ov_xml = ov_dir / f"{model_name}.xml" if not ov_xml.exists(): print(f"Exporting {pt_path} to OpenVINO IR...") model = YOLO(str(pt_path)) model.export(format="openvino", dynamic=True, half=half) else: print(f"OpenVINO IR already exists: {ov_xml}") return ov_xml def quantize_openvino_model(ov_xml: Path, model_name: str = "yolo11x") -> Path: """Quantize OpenVINO IR model to INT8 using NNCF.""" int8_dir = Path(f"{model_name}_openvino_int8_model") int8_xml = int8_dir / f"{model_name}.xml" if int8_xml.exists(): print(f"INT8 model already exists: {int8_xml}") return int8_xml print("Quantization requires a calibration dataset. Skipping actual quantization in this demo.") return ov_xml # Return FP32 if no quantization # --- OpenVINO Inference Pipeline --- class OpenVINOYOLODetector: def __init__(self, model_xml: Path, device: str = "AUTO"): self.core = ov.Core() self.device = device self.model = self.core.read_model(model_xml) self.input_shape = self.model.inputs[0].shape self.input_height = self.input_shape[2] self.input_width = self.input_shape[3] self.ov_config = {} if device != "CPU": self.model.reshape({0: [1, 3, 640, 640]}) if "GPU" in device or ("AUTO" in device and "GPU" in self.core.available_devices): self.ov_config = {"GPU_DISABLE_WINOGRAD_CONVOLUTION": "YES"} self.compiled_model = self.core.compile_model(model=self.model, device_name=self.device, config=self.ov_config) self.output_layer = self.compiled_model.output(0) def preprocess(self, frame: np.ndarray) -> np.ndarray: img = cv2.resize(frame, (self.input_width, self.input_height)) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img = img.astype(np.float32) / 255.0 img = img.transpose(2, 0, 1)[None] return img def infer(self, frame: np.ndarray, conf_threshold: float = 0.25) -> List[Dict]: input_tensor = self.preprocess(frame) output = self.compiled_model([input_tensor])[self.output_layer] return self.postprocess(output, frame.shape, conf_threshold) def postprocess(self, output: np.ndarray, frame_shape, conf_threshold: float) -> List[Dict]: # Output: (1, 84, 8400) or (84, 8400) or (8400, 84) if output.ndim == 3: output = np.squeeze(output) if output.shape[0] == 84: output = output.T # (8400, 84) boxes = output[:, :4] scores = output[:, 4:] class_ids = np.argmax(scores, axis=1) confidences = np.max(scores, axis=1) detections = [] h, w = frame_shape[:2] for i, (box, score, class_id) in enumerate(zip(boxes, confidences, class_ids)): if score < conf_threshold: continue x_c, y_c, bw, bh = box # If normalized, scale to input size if all(0.0 <= v <= 1.0 for v in box): x_c *= self.input_width y_c *= self.input_height bw *= self.input_width bh *= self.input_height # Scale to original frame size scale_x = w / self.input_width scale_y = h / self.input_height x_c *= scale_x y_c *= scale_y bw *= scale_x bh *= scale_y x1 = int(round(x_c - bw / 2)) y1 = int(round(y_c - bh / 2)) x2 = int(round(x_c + bw / 2)) y2 = int(round(y_c + bh / 2)) x1 = max(0, min(x1, w - 1)) y1 = max(0, min(y1, h - 1)) x2 = max(0, min(x2, w - 1)) y2 = max(0, min(y2, h - 1)) if x2 <= x1 or y2 <= y1: continue # Only keep class 9 as traffic light, rename if found if class_id == 9: class_name = "traffic light" elif class_id < len(TRAFFIC_CLASS_NAMES): class_name = TRAFFIC_CLASS_NAMES[class_id] else: continue # Remove unknown/other classes detections.append({ 'bbox': [x1, y1, x2, y2], 'confidence': float(score), 'class_id': int(class_id), 'class_name': class_name }) return detections def draw(self, frame: np.ndarray, detections: List[Dict], box_thickness: int = 2) -> np.ndarray: # 80+ visually distinct colors for COCO classes (BGR) COCO_COLORS = [ (255, 56, 56), (255, 157, 151), (255, 112, 31), (255, 178, 29), (207, 210, 49), (72, 249, 10), (146, 204, 23), (61, 219, 134), (26, 147, 52), (0, 212, 187), (44, 153, 168), (0, 194, 255), (52, 69, 147), (100, 115, 255), (0, 24, 236), (132, 56, 255), (82, 0, 133), (203, 56, 255), (255, 149, 200), (255, 55, 199), (255, 255, 56), (255, 255, 151), (255, 255, 31), (255, 255, 29), (207, 255, 49), (72, 255, 10), (146, 255, 23), (61, 255, 134), (26, 255, 52), (0, 255, 187), (44, 255, 168), (0, 255, 255), (52, 255, 147), (100, 255, 255), (0, 255, 236), (132, 255, 255), (82, 255, 133), (203, 255, 255), (255, 255, 200), (255, 255, 199), (56, 255, 255), (157, 255, 151), (112, 255, 31), (178, 255, 29), (210, 255, 49), (249, 255, 10), (204, 255, 23), (219, 255, 134), (147, 255, 52), (212, 255, 187), (153, 255, 168), (194, 255, 255), (69, 255, 147), (115, 255, 255), (24, 255, 236), (56, 132, 255), (157, 82, 151), (112, 203, 31), (178, 255, 29), (210, 255, 49), (249, 72, 10), (204, 146, 23), (219, 61, 134), (147, 26, 52), (212, 0, 187), (153, 44, 168), (194, 0, 255), (69, 52, 147), (115, 100, 255), (24, 0, 236), (56, 132, 255), (157, 82, 151), (112, 203, 31), (178, 255, 29), (210, 255, 49), (249, 72, 10), (204, 146, 23), (219, 61, 134), (147, 26, 52), (212, 0, 187), (153, 44, 168), (194, 0, 255), (69, 52, 147), (115, 100, 255), (24, 0, 236), (56, 132, 255), (157, 82, 151), (112, 203, 31), (178, 255, 29), (210, 255, 49) ] for det in detections: x1, y1, x2, y2 = det['bbox'] label = f"{det['class_name']} {det['confidence']:.2f}" color = COCO_COLORS[det['class_id'] % len(COCO_COLORS)] cv2.rectangle(frame, (x1, y1), (x2, y2), color, box_thickness) cv2.putText(frame, label, (x1, max(y1 - 10, 0)), cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2) return frame # --- Video/Image/Live Inference --- def run_inference(detector: OpenVINOYOLODetector, source=0, conf_threshold=0.25, flip=False, use_popup=False, video_width=None): if isinstance(source, str) and not os.path.exists(source): print(f"Downloading sample video: {source}") import requests url = "https://storage.openvinotoolkit.org/repositories/openvino_notebooks/data/data/video/people.mp4" r = requests.get(url) with open(source, 'wb') as f: f.write(r.content) cap = cv2.VideoCapture(source) if not cap.isOpened(): print(f"Failed to open video source: {source}") return window_name = "YOLOv11x + OpenVINO Detection" if use_popup: cv2.namedWindow(window_name, cv2.WINDOW_GUI_NORMAL | cv2.WINDOW_AUTOSIZE) frame_count = 0 times = [] while True: ret, frame = cap.read() if not ret: break if flip: frame = cv2.flip(frame, 1) if video_width: scale = video_width / max(frame.shape[:2]) frame = cv2.resize(frame, None, fx=scale, fy=scale, interpolation=cv2.INTER_AREA) start = time.time() detections = detector.infer(frame, conf_threshold=conf_threshold) frame = detector.draw(frame, detections) elapsed = time.time() - start times.append(elapsed) if len(times) > 200: times.pop(0) fps = 1.0 / np.mean(times) if times else 0 cv2.putText(frame, f"FPS: {fps:.1f}", (20, 40), cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0, 0, 255), 2) if use_popup: cv2.imshow(window_name, frame) if cv2.waitKey(1) & 0xFF == 27: break else: cv2.imshow(window_name, frame) if cv2.waitKey(1) & 0xFF == 27: break frame_count += 1 cap.release() cv2.destroyAllWindows() # --- Main Entrypoint --- if __name__ == "__main__": # Choose model: yolo11x or yolo11n, etc. MODEL_NAME = "yolo11x" DEVICE = "AUTO" # or "CPU", "GPU" # Step 1: Convert model if needed ov_xml = convert_yolo_to_openvino(MODEL_NAME) # Step 2: Quantize (optional, demo skips actual quantization) ov_xml = quantize_openvino_model(ov_xml, MODEL_NAME) # Step 3: Create detector detector = OpenVINOYOLODetector(ov_xml, device=DEVICE) # Step 4: Run on webcam, video, or image # Webcam: source=0, Video: source="video.mp4", Image: source="image.jpg" run_inference(detector, source=0, conf_threshold=0.25, flip=True, use_popup=True, video_width=1280) # To run on a video file: run_inference(detector, source="people.mp4", conf_threshold=0.25) # To run on an image: run_inference(detector, source="image.jpg", conf_threshold=0.25) # To run async or batch, extend the OpenVINOYOLODetector class with async API as needed. import numpy as np import cv2 def postprocess_openvino_yolo(output, conf_threshold=0.4, iou_threshold=0.5, input_shape=(640, 640), original_shape=None): """ output: OpenVINO raw output tensor (e.g., shape [1, 25200, 85]) conf_threshold: minimum confidence iou_threshold: for NMS input_shape: model input size (w, h) original_shape: original image size (w, h) """ # 1. Squeeze batch dimension output = np.squeeze(output) # [25200, 85] # 2. Split predictions boxes = output[:, :4] obj_conf = output[:, 4] class_scores = output[:, 5:] # 3. Get class with highest score class_ids = np.argmax(class_scores, axis=1) class_conf = class_scores[np.arange(len(class_scores)), class_ids] # 4. Multiply objectness confidence with class confidence scores = obj_conf * class_conf # 5. Filter by confidence threshold mask = scores > conf_threshold boxes = boxes[mask] scores = scores[mask] class_ids = class_ids[mask] if original_shape is not None: # Rescale boxes from input_shape to original image shape input_w, input_h = input_shape orig_w, orig_h = original_shape scale_x = orig_w / input_w scale_y = orig_h / input_h boxes[:, 0] *= scale_x # x1 boxes[:, 1] *= scale_y # y1 boxes[:, 2] *= scale_x # x2 boxes[:, 3] *= scale_y # y2 # 6. Convert boxes to [x, y, w, h] format for OpenCV NMS boxes_xywh = [] for box in boxes: x1, y1, x2, y2 = box boxes_xywh.append([x1, y1, x2 - x1, y2 - y1]) # 7. Apply NMS indices = cv2.dnn.NMSBoxes(boxes_xywh, scores.tolist(), conf_threshold, iou_threshold) # 8. Return filtered boxes result_boxes = [] result_scores = [] result_classes = [] if len(boxes) > 0 and len(scores) > 0: indices = cv2.dnn.NMSBoxes(boxes_xywh, scores.tolist(), conf_threshold, iou_threshold) if len(indices) > 0: indices = np.array(indices).flatten() for i in indices: i = int(i) result_boxes.append(boxes[i]) result_scores.append(scores[i]) result_classes.append(class_ids[i]) return result_boxes, result_scores, result_classes import os import time import numpy as np import cv2 from pathlib import Path from typing import List, Dict, Optional # Only traffic-related classes for detection TRAFFIC_CLASS_NAMES = [ 'person', 'bicycle', 'car', 'motorcycle', 'bus', 'truck', 'traffic light', 'stop sign', 'parking meter' ] class OpenVINOVehicleDetector: def __init__(self, model_path: str = None, device: str = "AUTO", use_quantized: bool = False, enable_ocr: bool = False, confidence_threshold: float = 0.4): import openvino as ov self.device = device self.confidence_threshold = confidence_threshold self.ocr_reader = None self.class_names = TRAFFIC_CLASS_NAMES self.performance_stats = { 'fps': 0, 'avg_inference_time': 0, 'frames_processed': 0, 'backend': f"OpenVINO-{device}", 'total_detections': 0, 'detection_rate': 0 } self._inference_times = [] self._start_time = time.time() self._frame_count = 0 # Model selection logic self.model_path = self._find_best_model(model_path, use_quantized) self.core = ov.Core() self.model = self.core.read_model(self.model_path) # Always reshape to static shape before accessing .shape self.model.reshape({0: [1, 3, 640, 640]}) self.input_shape = self.model.inputs[0].shape self.input_height = self.input_shape[2] self.input_width = self.input_shape[3] self.ov_config = {} if device != "CPU": # Already reshaped above, so nothing more needed here pass if "GPU" in device or ("AUTO" in device and "GPU" in self.core.available_devices): self.ov_config = {"GPU_DISABLE_WINOGRAD_CONVOLUTION": "YES"} self.compiled_model = self.core.compile_model(model=self.model, device_name=self.device, config=self.ov_config) self.output_layer = self.compiled_model.output(0) def _find_best_model(self, model_path, use_quantized): # Priority: quantized IR > IR > .pt search_paths = [ Path(model_path) if model_path else None, Path("yolo11x_openvino_int8_model/yolo11x.xml") if use_quantized else None, Path("yolo11x_openvino_model/yolo11x.xml"), Path("rcb/yolo11x_openvino_model/yolo11x.xml"), Path("yolo11x.xml"), Path("rcb/yolo11x.xml"), Path("yolo11x.pt"), Path("rcb/yolo11x.pt") ] for p in search_paths: if p and p.exists(): return str(p) raise FileNotFoundError("No suitable YOLOv11x model found for OpenVINO.") def detect_vehicles(self, frame: np.ndarray, conf_threshold: float = None) -> List[Dict]: if conf_threshold is None: conf_threshold = 0.1 # Lowered for debugging start = time.time() input_tensor = self._preprocess(frame) output = self.compiled_model([input_tensor])[self.output_layer] # Debug: print raw output shape print(f"[DEBUG] Model output shape: {output.shape}") detections = self._postprocess(output, frame.shape, conf_threshold) print(f"[DEBUG] Detections after postprocess: {len(detections)}") elapsed = time.time() - start self._inference_times.append(elapsed) self._frame_count += 1 self.performance_stats['frames_processed'] = self._frame_count self.performance_stats['total_detections'] += len(detections) if len(self._inference_times) > 100: self._inference_times.pop(0) self.performance_stats['avg_inference_time'] = float(np.mean(self._inference_times)) if self._inference_times else 0 total_time = time.time() - self._start_time self.performance_stats['fps'] = self._frame_count / total_time if total_time > 0 else 0 return detections def _preprocess(self, frame: np.ndarray) -> np.ndarray: img = cv2.resize(frame, (self.input_width, self.input_height)) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img = img.astype(np.float32) / 255.0 img = img.transpose(2, 0, 1)[None] return img def _postprocess(self, output: np.ndarray, frame_shape, conf_threshold: float) -> List[Dict]: # Output: (1, 84, 8400) or (84, 8400) or (8400, 84) if output.ndim == 3: output = np.squeeze(output) if output.shape[0] == 84: output = output.T # (8400, 84) boxes = output[:, :4] scores = output[:, 4:] class_ids = np.argmax(scores, axis=1) confidences = np.max(scores, axis=1) detections = [] h, w = frame_shape[:2] for i, (box, score, class_id) in enumerate(zip(boxes, confidences, class_ids)): if score < conf_threshold: continue x_c, y_c, bw, bh = box # If normalized, scale to input size if all(0.0 <= v <= 1.0 for v in box): x_c *= self.input_width y_c *= self.input_height bw *= self.input_width bh *= self.input_height # Scale to original frame size scale_x = w / self.input_width scale_y = h / self.input_height x_c *= scale_x y_c *= scale_y bw *= scale_x bh *= scale_y x1 = int(round(x_c - bw / 2)) y1 = int(round(y_c - bh / 2)) x2 = int(round(x_c + bw / 2)) y2 = int(round(y_c + bh / 2)) x1 = max(0, min(x1, w - 1)) y1 = max(0, min(y1, h - 1)) x2 = max(0, min(x2, w - 1)) y2 = max(0, min(y2, h - 1)) if x2 <= x1 or y2 <= y1: continue # Only keep class 9 as traffic light, rename if found if class_id == 9: class_name = "traffic light" elif class_id < len(TRAFFIC_CLASS_NAMES): class_name = TRAFFIC_CLASS_NAMES[class_id] else: continue # Remove unknown/other classes detections.append({ 'bbox': [x1, y1, x2, y2], 'confidence': float(score), 'class_id': int(class_id), 'class_name': class_name }) print(f"[DEBUG] Raw detections before NMS: {len(detections)}") # Apply NMS if len(detections) > 0: boxes = np.array([det['bbox'] for det in detections]) scores = np.array([det['confidence'] for det in detections]) indices = cv2.dnn.NMSBoxes(boxes.tolist(), scores.tolist(), conf_threshold, 0.5) if isinstance(indices, (list, tuple)) and len(indices) > 0: indices = np.array(indices).flatten() elif isinstance(indices, np.ndarray) and indices.size > 0: indices = indices.flatten() else: indices = [] detections = [detections[int(i)] for i in indices] if len(indices) > 0 else [] print(f"[DEBUG] Detections after NMS: {len(detections)}") return detections def draw(self, frame: np.ndarray, detections: List[Dict], box_thickness: int = 2) -> np.ndarray: # 80+ visually distinct colors for COCO classes (BGR) COCO_COLORS = [ (255, 56, 56), (255, 157, 151), (255, 112, 31), (255, 178, 29), (207, 210, 49), (72, 249, 10), (146, 204, 23), (61, 219, 134), (26, 147, 52), (0, 212, 187), (44, 153, 168), (0, 194, 255), (52, 69, 147), (100, 115, 255), (0, 24, 236), (132, 56, 255), (82, 0, 133), (203, 56, 255), (255, 149, 200), (255, 55, 199), (255, 255, 56), (255, 255, 151), (255, 255, 31), (255, 255, 29), (207, 255, 49), (72, 255, 10), (146, 255, 23), (61, 255, 134), (26, 255, 52), (0, 255, 187), (44, 255, 168), (0, 255, 255), (52, 255, 147), (100, 255, 255), (0, 255, 236), (132, 255, 255), (82, 255, 133), (203, 255, 255), (255, 255, 200), (255, 255, 199), (56, 255, 255), (157, 255, 151), (112, 255, 31), (178, 255, 29), (210, 255, 49), (249, 255, 10), (204, 255, 23), (219, 255, 134), (147, 255, 52), (212, 255, 187), (153, 255, 168), (194, 255, 255), (69, 255, 147), (115, 255, 255), (24, 255, 236), (56, 132, 255), (157, 82, 151), (112, 203, 31), (178, 255, 29), (210, 255, 49), (249, 72, 10), (204, 146, 23), (219, 61, 134), (147, 26, 52), (212, 0, 187), (153, 44, 168), (194, 0, 255), (69, 52, 147), (115, 100, 255), (24, 0, 236), (56, 132, 255), (157, 82, 151), (112, 203, 31), (178, 255, 29), (210, 255, 49), (249, 72, 10), (204, 146, 23), (219, 61, 134), (147, 26, 52), (212, 0, 187), (153, 44, 168), (194, 0, 255), (69, 52, 147), (115, 100, 255), (24, 0, 236), (56, 132, 255), (157, 82, 151), (112, 203, 31), (178, 255, 29), (210, 255, 49) ] for det in detections: x1, y1, x2, y2 = det['bbox'] label = f"{det['class_name']} {det['confidence']:.2f}" color = COCO_COLORS[det['class_id'] % len(COCO_COLORS)] cv2.rectangle(frame, (x1, y1), (x2, y2), color, box_thickness) cv2.putText(frame, label, (x1, max(y1 - 10, 0)), cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2) return frame # --- Video/Image/Live Inference --- def run_inference(detector: OpenVINOYOLODetector, source=0, conf_threshold=0.25, flip=False, use_popup=False, video_width=None): if isinstance(source, str) and not os.path.exists(source): print(f"Downloading sample video: {source}") import requests url = "https://storage.openvinotoolkit.org/repositories/openvino_notebooks/data/data/video/people.mp4" r = requests.get(url) with open(source, 'wb') as f: f.write(r.content) cap = cv2.VideoCapture(source) if not cap.isOpened(): print(f"Failed to open video source: {source}") return window_name = "YOLOv11x + OpenVINO Detection" if use_popup: cv2.namedWindow(window_name, cv2.WINDOW_GUI_NORMAL | cv2.WINDOW_AUTOSIZE) frame_count = 0 times = [] while True: ret, frame = cap.read() if not ret: break if flip: frame = cv2.flip(frame, 1) if video_width: scale = video_width / max(frame.shape[:2]) frame = cv2.resize(frame, None, fx=scale, fy=scale, interpolation=cv2.INTER_AREA) start = time.time() detections = detector.infer(frame, conf_threshold=conf_threshold) frame = detector.draw(frame, detections) elapsed = time.time() - start times.append(elapsed) if len(times) > 200: times.pop(0) fps = 1.0 / np.mean(times) if times else 0 cv2.putText(frame, f"FPS: {fps:.1f}", (20, 40), cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0, 0, 255), 2) if use_popup: cv2.imshow(window_name, frame) if cv2.waitKey(1) & 0xFF == 27: break else: cv2.imshow(window_name, frame) if cv2.waitKey(1) & 0xFF == 27: break frame_count += 1 cap.release() cv2.destroyAllWindows() # --- Main Entrypoint --- if __name__ == "__main__": # Choose model: yolo11x or yolo11n, etc. MODEL_NAME = "yolo11x" DEVICE = "AUTO" # or "CPU", "GPU" # Step 1: Convert model if needed ov_xml = convert_yolo_to_openvino(MODEL_NAME) # Step 2: Quantize (optional, demo skips actual quantization) ov_xml = quantize_openvino_model(ov_xml, MODEL_NAME) # Step 3: Create detector detector = OpenVINOYOLODetector(ov_xml, device=DEVICE) # Step 4: Run on webcam, video, or image # Webcam: source=0, Video: source="video.mp4", Image: source="image.jpg" run_inference(detector, source=0, conf_threshold=0.25, flip=True, use_popup=True, video_width=1280) # To run on a video file: run_inference(detector, source="people.mp4", conf_threshold=0.25) # To run on an image: run_inference(detector, source="image.jpg", conf_threshold=0.25) # To run async or batch, extend the OpenVINOYOLODetector class with async API as needed. import numpy as np import cv2 def postprocess_openvino_yolo(output, conf_threshold=0.4, iou_threshold=0.5, input_shape=(640, 640), original_shape=None): """ output: OpenVINO raw output tensor (e.g., shape [1, 25200, 85]) conf_threshold: minimum confidence iou_threshold: for NMS input_shape: model input size (w, h) original_shape: original image size (w, h) """ # 1. Squeeze batch dimension output = np.squeeze(output) # [25200, 85] # 2. Split predictions boxes = output[:, :4] obj_conf = output[:, 4] class_scores = output[:, 5:] # 3. Get class with highest score class_ids = np.argmax(class_scores, axis=1) class_conf = class_scores[np.arange(len(class_scores)), class_ids] # 4. Multiply objectness confidence with class confidence scores = obj_conf * class_conf # 5. Filter by confidence threshold mask = scores > conf_threshold boxes = boxes[mask] scores = scores[mask] class_ids = class_ids[mask] if original_shape is not None: # Rescale boxes from input_shape to original image shape input_w, input_h = input_shape orig_w, orig_h = original_shape scale_x = orig_w / input_w scale_y = orig_h / input_h boxes[:, 0] *= scale_x # x1 boxes[:, 1] *= scale_y # y1 boxes[:, 2] *= scale_x # x2 boxes[:, 3] *= scale_y # y2 # 6. Convert boxes to [x, y, w, h] format for OpenCV NMS boxes_xywh = [] for box in boxes: x1, y1, x2, y2 = box boxes_xywh.append([x1, y1, x2 - x1, y2 - y1]) # 7. Apply NMS indices = cv2.dnn.NMSBoxes(boxes_xywh, scores.tolist(), conf_threshold, iou_threshold) # 8. Return filtered boxes result_boxes = [] result_scores = [] result_classes = [] if len(boxes) > 0 and len(scores) > 0: indices = cv2.dnn.NMSBoxes(boxes_xywh, scores.tolist(), conf_threshold, iou_threshold) if len(indices) > 0: indices = np.array(indices).flatten() for i in indices: i = int(i) result_boxes.append(boxes[i]) result_scores.append(scores[i]) result_classes.append(class_ids[i]) return result_boxes, result_scores, result_classes import os import time import numpy as np import cv2 from pathlib import Path from typing import List, Dict, Optional # Only traffic-related classes for detection TRAFFIC_CLASS_NAMES = [ 'person', 'bicycle', 'car', 'motorcycle', 'bus', 'truck', 'traffic light', 'stop sign', 'parking meter' ] class OpenVINOVehicleDetector: def __init__(self, model_path: str = None, device: str = "AUTO", use_quantized: bool = False, enable_ocr: bool = False, confidence_threshold: float = 0.4): import openvino as ov self.device = device self.confidence_threshold = confidence_threshold self.ocr_reader = None self.class_names = TRAFFIC_CLASS_NAMES self.performance_stats = { 'fps': 0, 'avg_inference_time': 0, 'frames_processed': 0, 'backend': f"OpenVINO-{device}", 'total_detections': 0, 'detection_rate': 0 } self._inference_times = [] self._start_time = time.time() self._frame_count = 0 # Model selection logic self.model_path = self._find_best_model(model_path, use_quantized) self.core = ov.Core() self.model = self.core.read_model(self.model_path) # Always reshape to static shape before accessing .shape self.model.reshape({0: [1, 3, 640, 640]}) self.input_shape = self.model.inputs[0].shape self.input_height = self.input_shape[2] self.input_width = self.input_shape[3] self.ov_config = {} if device != "CPU": # Already reshaped above, so nothing more needed here pass if "GPU" in device or ("AUTO" in device and "GPU" in self.core.available_devices): self.ov_config = {"GPU_DISABLE_WINOGRAD_CONVOLUTION": "YES"} self.compiled_model = self.core.compile_model(model=self.model, device_name=self.device, config=self.ov_config) self.output_layer = self.compiled_model.output(0) def _find_best_model(self, model_path, use_quantized): # Priority: quantized IR > IR > .pt search_paths = [ Path(model_path) if model_path else None, Path("yolo11x_openvino_int8_model/yolo11x.xml") if use_quantized else None, Path("yolo11x_openvino_model/yolo11x.xml"), Path("rcb/yolo11x_openvino_model/yolo11x.xml"), Path("yolo11x.xml"), Path("rcb/yolo11x.xml"), Path("yolo11x.pt"), Path("rcb/yolo11x.pt") ] for p in search_paths: if p and p.exists(): return str(p) raise FileNotFoundError("No suitable YOLOv11x model found for OpenVINO.") def detect_vehicles(self, frame: np.ndarray, conf_threshold: float = None) -> List[Dict]: if conf_threshold is None: conf_threshold = 0.1 # Lowered for debugging start = time.time() input_tensor = self._preprocess(frame) output = self.compiled_model([input_tensor])[self.output_layer] # Debug: print raw output shape print(f"[DEBUG] Model output shape: {output.shape}") detections = self._postprocess(output, frame.shape, conf_threshold) print(f"[DEBUG] Detections after postprocess: {len(detections)}") elapsed = time.time() - start self._inference_times.append(elapsed) self._frame_count += 1 self.performance_stats['frames_processed'] = self._frame_count self.performance_stats['total_detections'] += len(detections) if len(self._inference_times) > 100: self._inference_times.pop(0) self.performance_stats['avg_inference_time'] = float(np.mean(self._inference_times)) if self._inference_times else 0 total_time = time.time() - self._start_time self.performance_stats['fps'] = self._frame_count / total_time if total_time > 0 else 0 return detections def _preprocess(self, frame: np.ndarray) -> np.ndarray: img = cv2.resize(frame, (self.input_width, self.input_height)) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img = img.astype(np.float32) / 255.0 img = img.transpose(2, 0, 1)[None] return img def _postprocess(self, output: np.ndarray, frame_shape, conf_threshold: float) -> List[Dict]: # Output: (1, 84, 8400) or (84, 8400) or (8400, 84) if output.ndim == 3: output = np.squeeze(output) if output.shape[0] == 84: output = output.T # (8400, 84) boxes = output[:, :4] scores = output[:, 4:] class_ids = np.argmax(scores, axis=1) confidences = np.max(scores, axis=1) detections = [] h, w = frame_shape[:2] for i, (box, score, class_id) in enumerate(zip(boxes, confidences, class_ids)): if score < conf_threshold: continue x_c, y_c, bw, bh = box # If normalized, scale to input size if all(0.0 <= v <= 1.0 for v in box): x_c *= self.input_width y_c *= self.input_height bw *= self.input_width bh *= self.input_height # Scale to original frame size scale_x = w / self.input_width scale_y = h / self.input_height x_c *= scale_x y_c *= scale_y bw *= scale_x bh *= scale_y x1 = int(round(x_c - bw / 2)) y1 = int(round(y_c - bh / 2)) x2 = int(round(x_c + bw / 2)) y2 = int(round(y_c + bh / 2)) x1 = max(0, min(x1, w - 1)) y1 = max(0, min(y1, h - 1)) x2 = max(0, min(x2, w - 1)) y2 = max(0, min(y2, h - 1)) if x2 <= x1 or y2 <= y1: continue # Only keep class 9 as traffic light, rename if found if class_id == 9: class_name = "traffic light" elif class_id < len(TRAFFIC_CLASS_NAMES): class_name = TRAFFIC_CLASS_NAMES[class_id] else: continue # Remove unknown/other classes detections.append({ 'bbox': [x1, y1, x2, y2], 'confidence': float(score), 'class_id': int(class_id), 'class_name': class_name }) print(f"[DEBUG] Raw detections before NMS: {len(detections)}") # Apply NMS if len(detections) > 0: boxes = np.array([det['bbox'] for det in detections]) scores = np.array([det['confidence'] for det in detections]) indices = cv2.dnn.NMSBoxes(boxes.tolist(), scores.tolist(), conf_threshold, 0.5) if isinstance(indices, (list, tuple)) and len(indices) > 0: indices = np.array(indices).flatten() elif isinstance(indices, np.ndarray) and indices.size > 0: indices = indices.flatten() else: indices = [] detections = [detections[int(i)] for i in indices] if len(indices) > 0 else [] print(f"[DEBUG] Detections after NMS: {len(detections)}") return detections def draw(self, frame: np.ndarray, detections: List[Dict], box_thickness: int = 2) -> np.ndarray: # 80+ visually distinct colors for COCO classes (BGR) COCO_COLORS = [ (255, 56, 56), (255, 157, 151), (255, 112, 31), (255, 178, 29), (207, 210, 49), (72, 249, 10), (146, 204, 23), (61, 219, 134), (26, 147, 52), (0, 212, 187), (44, 153, 168), (0, 194, 255), (52, 69, 147), (100, 115, 255), (0, 24, 236), (132, 56, 255), (82, 0, 133), (203, 56, 255), (255, 149, 200), (255, 55, 199), (255, 255, 56), (255, 255, 151), (255, 255, 31), (255, 255, 29), (207, 255, 49), (72, 255, 10), (146, 255, 23), (61, 255, 134), (26, 255, 52), (0, 255, 187), (44, 255, 168), (0, 255, 255), (52, 255, 147), (100, 255, 255), (0, 255, 236), (132, 255, 255), (82, 255, 133), (203, 255, 255), (255, 255, 200), (255, 255, 199), (56, 255, 255), (157, 255, 151), (112, 255, 31), (178, 255, 29), (210, 255, 49), (249, 255, 10), (204, 255, 23), (219, 255, 134), (147, 255, 52), (212, 255, 187), (153, 255, 168), (194, 255, 255), (69, 255, 147), (115, 255, 255), (24, 255, 236), (56, 132, 255), (157, 82, 151), (112, 203, 31), (178, 255, 29), (210, 255, 49), (249, 72, 10), (204, 146, 23), (219, 61, 134), (147, 26, 52), (212, 0, 187), (153, 44, 168), (194, 0, 255), (69, 52, 147), (115, 100, 255), (24, 0, 236), (56, 132, 255), (157, 82, 151), (112, 203, 31), (178, 255, 29), (210, 255, 49), (249, 72, 10), (204, 146, 23), (219, 61, 134), (147, 26, 52), (212, 0, 187), (153, 44, 168), (194, 0, 255), (69, 52, 147), (115, 100, 255), (24, 0, 236), (56, 132, 255), (157, 82, 151), (112, 203, 31), (178, 255, 29), (210, 255, 49) ] for det in detections: x1, y1, x2, y2 = det['bbox'] label = f"{det['class_name']} {det['confidence']:.2f}" color = COCO_COLORS[det['class_id'] % len(COCO_COLORS)] cv2.rectangle(frame, (x1, y1), (x2, y2), color, box_thickness) cv2.putText(frame, label, (x1, max(y1 - 10, 0)), cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2) return frame # --- Video/Image/Live Inference --- def run_inference(detector: OpenVINOYOLODetector, source=0, conf_threshold=0.25, flip=False, use_popup=False, video_width=None): if isinstance(source, str) and not os.path.exists(source): print(f"Downloading sample video: {source}") import requests url = "https://storage.openvinotoolkit.org/repositories/openvino_notebooks/data/data/video/people.mp4" r = requests.get(url) with open(source, 'wb') as f: f.write(r.content) cap = cv2.VideoCapture(source) if not cap.isOpened(): print(f"Failed to open video source: {source}") return window_name = "YOLOv11x + OpenVINO Detection" if use_popup: cv2.namedWindow(window_name, cv2.WINDOW_GUI_NORMAL | cv2.WINDOW_AUTOSIZE) frame_count = 0 times = [] while True: ret, frame = cap.read() if not ret: break if flip: frame = cv2.flip(frame, 1) if video_width: scale = video_width / max(frame.shape[:2]) frame = cv2.resize(frame, None, fx=scale, fy=scale, interpolation=cv2.INTER_AREA) start = time.time() detections = detector.infer(frame, conf_threshold=conf_threshold) frame = detector.draw(frame, detections) elapsed = time.time() - start times.append(elapsed) if len(times) > 200: times.pop(0) fps = 1.0 / np.mean(times) if times else 0 cv2.putText(frame, f"FPS: {fps:.1f}", (20, 40), cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0, 0, 255), 2) if use_popup: cv2.imshow(window_name, frame) if cv2.waitKey(1) & 0xFF == 27: break else: cv2.imshow(window_name, frame) if cv2.waitKey(1) & 0xFF == 27: break frame_count += 1 cap.release() cv2.destroyAllWindows() # --- Main Entrypoint --- if __name__ == "__main__": # Choose model: yolo11x or yolo11n, etc. MODEL_NAME = "yolo11x" DEVICE = "AUTO" # or "CPU", "GPU" # Step 1: Convert model if needed ov_xml = convert_yolo_to_openvino(MODEL_NAME) # Step 2: Quantize (optional, demo skips actual quantization) ov_xml = quantize_openvino_model(ov_xml, MODEL_NAME) # Step 3: Create detector detector = OpenVINOYOLODetector(ov_xml, device=DEVICE) # Step 4: Run on webcam, video, or image # Webcam: source=0, Video: source="video.mp4", Image: source="image.jpg" run_inference(detector, source=0, conf_threshold=0.25, flip=True, use_popup=True, video_width=1280) # To run on a video file: run_inference(detector, source="people.mp4", conf_threshold=0.25) # To run on an image: run_inference(detector, source="image.jpg", conf_threshold=0.25) # To run async or batch, extend the OpenVINOYOLODetector class with async API as needed. import numpy as np import cv2 def postprocess_openvino_yolo(output, conf_threshold=0.4, iou_threshold=0.5, input_shape=(640, 640), original_shape=None): """ output: OpenVINO raw output tensor (e.g., shape [1, 25200, 85]) conf_threshold: minimum confidence iou_threshold: for NMS input_shape: model input size (w, h) original_shape: original image size (w, h) """ # 1. Squeeze batch dimension output = np.squeeze(output) # [25200, 85] # 2. Split predictions boxes = output[:, :4] obj_conf = output[:, 4] class_scores = output[:, 5:] # 3. Get class with highest score class_ids = np.argmax(class_scores, axis=1) class_conf = class_scores[np.arange(len(class_scores)), class_ids] # 4. Multiply objectness confidence with class confidence scores = obj_conf * class_conf # 5. Filter by confidence threshold mask = scores > conf_threshold boxes = boxes[mask] scores = scores[mask] class_ids = class_ids[mask] if original_shape is not None: # Rescale boxes from input_shape to original image shape input_w, input_h = input_shape orig_w, orig_h = original_shape scale_x = orig_w / input_w scale_y = orig_h / input_h boxes[:, 0] *= scale_x # x1 boxes[:, 1] *= scale_y # y1 boxes[:, 2] *= scale_x # x2 boxes[:, 3] *= scale_y # y2 # 6. Convert boxes to [x, y, w, h] format for OpenCV NMS boxes_xywh = [] for box in boxes: x1, y1, x2, y2 = box boxes_xywh.append([x1, y1, x2 - x1, y2 - y1]) # 7. Apply NMS indices = cv2.dnn.NMSBoxes(boxes_xywh, scores.tolist(), conf_threshold, iou_threshold) # 8. Return filtered boxes result_boxes = [] result_scores = [] result_classes = [] if len(boxes) > 0 and len(scores) > 0: indices = cv2.dnn.NMSBoxes(boxes_xywh, scores.tolist(), conf_threshold, iou_threshold) if len(indices) > 0: indices = np.array(indices).flatten() for i in indices: i = int(i) result_boxes.append(boxes[i]) result_scores.append(scores[i]) result_classes.append(class_ids[i]) return result_boxes, result_scores, result_classes import os import time import numpy as np import cv2 from pathlib import Path from typing import List, Dict, Optional # Only traffic-related classes for detection TRAFFIC_CLASS_NAMES = [ 'person', 'bicycle', 'car', 'motorcycle', 'bus', 'truck', 'traffic light', 'stop sign', 'parking meter' ] class OpenVINOVehicleDetector: def __init__(self, model_path: str = None, device: str = "AUTO", use_quantized: bool = False, enable_ocr: bool = False, confidence_threshold: float = 0.4): import openvino as ov self.device = device self.confidence_threshold = confidence_threshold self.ocr_reader = None self.class_names = TRAFFIC_CLASS_NAMES self.performance_stats = { 'fps': 0, 'avg_inference_time': 0, 'frames_processed': 0, 'backend': f"OpenVINO-{device}", 'total_detections': 0, 'detection_rate': 0 } self._inference_times = [] self._start_time = time.time() self._frame_count = 0 # Model selection logic self.model_path = self._find_best_model(model_path, use_quantized) self.core = ov.Core() self.model = self.core.read_model(self.model_path) # Always reshape to static shape before accessing .shape self.model.reshape({0: [1, 3, 640, 640]}) self.input_shape = self.model.inputs[0].shape self.input_height = self.input_shape[2] self.input_width = self.input_shape[3] self.ov_config = {} if device != "CPU": # Already reshaped above, so nothing more needed here pass if "GPU" in device or ("AUTO" in device and "GPU" in self.core.available_devices): self.ov_config = {"GPU_DISABLE_WINOGRAD_CONVOLUTION": "YES"} self.compiled_model = self.core.compile_model(model=self.model, device_name=self.device, config=self.ov_config) self.output_layer = self.compiled_model.output(0) def _find_best_model(self, model_path, use_quantized): # Priority: quantized IR > IR > .pt search_paths = [ Path(model_path) if model_path else None, Path("yolo11x_openvino_int8_model/yolo11x.xml") if use_quantized else None, Path("yolo11x_openvino_model/yolo11x.xml"), Path("rcb/yolo11x_openvino_model/yolo11x.xml"), Path("yolo11x.xml"), Path("rcb/yolo11x.xml"), Path("yolo11x.pt"), Path("rcb/yolo11x.pt") ] for p in search_paths: if p and p.exists(): return str(p) raise FileNotFoundError("No suitable YOLOv11x model found for OpenVINO.") def detect_vehicles(self, frame: np.ndarray, conf_threshold: float = None) -> List[Dict]: if conf_threshold is None: conf_threshold = 0.1 # Lowered for debugging start = time.time() input_tensor = self._preprocess(frame) output = self.compiled_model([input_tensor])[self.output_layer] # Debug: print raw output shape print(f"[DEBUG] Model output shape: {output.shape}") detections = self._postprocess(output, frame.shape, conf_threshold) print(f"[DEBUG] Detections after postprocess: {len(detections)}") elapsed = time.time() - start self._inference_times.append(elapsed) self._frame_count += 1 self.performance_stats['frames_processed'] = self._frame_count self.performance_stats['total_detections'] += len(detections) if len(self._inference_times) > 100: self._inference_times.pop(0) self.performance_stats['avg_inference_time'] = float(np.mean(self._inference_times)) if self._inference_times else 0 total_time = time.time() - self._start_time self.performance_stats['fps'] = self._frame_count / total_time if total_time > 0 else 0 return detections def _preprocess(self, frame: np.ndarray) -> np.ndarray: img = cv2.resize(frame, (self.input_width, self.input_height)) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img = img.astype(np.float32) / 255.0 img = img.transpose(2, 0, 1)[None] return img def _postprocess(self, output: np.ndarray, frame_shape, conf_threshold: float) -> List[Dict]: # Output: (1, 84, 8400) or (84, 8400) or (8400, 84) if output.ndim == 3: output = np.squeeze(output) if output.shape[0] == 84: output = output.T # (8400, 84) boxes = output[:, :4] scores = output[:, 4:] class_ids = np.argmax(scores, axis=1) confidences = np.max(scores, axis=1) detections = [] h, w = frame_shape[:2] for i, (box, score, class_id) in enumerate(zip(boxes, confidences, class_ids)): if score < conf_threshold: continue x_c, y_c, bw, bh = box # If normalized, scale to input size if all(0.0 <= v <= 1.0 for v in box): x_c *= self.input_width y_c *= self.input_height bw *= self.input_width bh *= self.input_height # Scale to original frame size scale_x = w / self.input_width scale_y = h / self.input_height x_c *= scale_x y_c *= scale_y bw *= scale_x bh *= scale_y x1 = int(round(x_c - bw / 2)) y1 = int(round(y_c - bh / 2)) x2 = int(round(x_c + bw / 2)) y2 = int(round(y_c + bh / 2)) x1 = max(0, min(x1, w - 1)) y1 = max(0, min(y1, h - 1)) x2 = max(0, min(x2, w - 1)) y2 = max(0, min(y2, h - 1)) if x2 <= x1 or y2 <= y1: continue # Only keep class 9 as traffic light, rename if found if class_id == 9: class_name = "traffic light" elif class_id < len(TRAFFIC_CLASS_NAMES): class_name = TRAFFIC_CLASS_NAMES[class_id] else: continue # Remove unknown/other classes detections.append({ 'bbox': [x1, y1, x2, y2], 'confidence': float(score), 'class_id': int(class_id), 'class_name': class_name }) print(f"[DEBUG] Raw detections before NMS: {len(detections)}") # Apply NMS if len(detections) > 0: boxes = np.array([det['bbox'] for det in detections]) scores = np.array([det['confidence'] for det in detections]) indices = cv2.dnn.NMSBoxes(boxes.tolist(), scores.tolist(), conf_threshold, 0.5) if isinstance(indices, (list, tuple)) and len(indices) > 0: indices = np.array(indices).flatten() elif isinstance(indices, np.ndarray) and indices.size > 0: indices = indices.flatten() else: indices = [] detections = [detections[int(i)] for i in indices] if len(indices) > 0 else [] print(f"[DEBUG] Detections after NMS: {len(detections)}") return detections def draw(self, frame: np.ndarray, detections: List[Dict], box_thickness: int = 2) -> np.ndarray: # 80+ visually distinct colors for COCO classes (BGR) COCO_COLORS = [ (255, 56, 56), (255, 157, 151), (255, 112, 31), (255, 178, 29), (207, 210, 49), (72, 249, 10), (146, 204, 23), (61, 219, 134), (26, 147, 52), (0, 212, 187), (44, 153, 168), (0, 194, 255), (52, 69, 147), (100, 115, 255), (0, 24, 236), (132, 56, 255), (82, 0, 133), (203, 56, 255), (255, 149, 200), (255, 55, 199), (255, 255, 56), (255, 255, 151), (255, 255, 31), (255, 255, 29), (207, 255, 49), (72, 255, 10), (146, 255, 23), (61, 255, 134), (26, 255, 52), (0, 255, 187), (44, 255, 168), (0, 255, 255), (52, 255, 147), (100, 255, 255), (0, 255, 236), (132, 255, 255), (82, 255, 133), (203, 255, 255), (255, 255, 200), (255, 255, 199), (56, 255, 255), (157, 255, 151), (112, 255, 31), (178, 255, 29), (210, 255, 49), (249, 255, 10), (204, 255, 23), (219, 255, 134), (147, 255, 52), (212, 255, 187), (153, 255, 168), (194, 255, 255), (69, 255, 147), (115, 255, 255), (24, 255, 236), (56, 132, 255), (157, 82, 151), (112, 203, 31), (178, 255, 29), (210, 255, 49), (249, 72, 10), (204, 146, 23), (219, 61, 134), (147, 26, 52), (212, 0, 187), (153, 44, 168), (194, 0, 255), (69, 52, 147), (115, 100, 255), (24, 0, 236), (56, 132, 255), (157, 82, 151), (112, 203, 31), (178, 255, 29), (210, 255, 49), (249, 72, 10), (204, 146, 23), (219, 61, 134), (147, 26, 52), (212, 0, 187), (153, 44, 168), (194, 0, 255), (69, 52, 147), (115, 100, 255), (24, 0, 236), (56, 132, 255), (157, 82, 151), (112, 203, 31), (178, 255, 29), (210, 255, 49) ] for det in detections: x1, y1, x2, y2 = det['bbox'] label = f"{det['class_name']} {det['confidence']:.2f}" color = COCO_COLORS[det['class_id'] % len(COCO_COLORS)] cv2.rectangle(frame, (x1, y1), (x2, y2), color, box_thickness) cv2.putText(frame, label, (x1, max(y1 - 10, 0)), cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2) return frame