cleanup and files added
This commit is contained in:
308
detection_openvino_fixed.py
Normal file
308
detection_openvino_fixed.py
Normal file
@@ -0,0 +1,308 @@
|
||||
# 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
|
||||
|
||||
# --- 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 = [
|
||||
'person', 'bicycle', 'car', 'motorcycle', 'bus', 'truck',
|
||||
'traffic light', 'stop sign', 'parking meter'
|
||||
]
|
||||
|
||||
# --- 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
|
||||
|
||||
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)
|
||||
print(f"🎯 OpenVINOVehicleDetector: Using model: {self.model_path}")
|
||||
|
||||
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):
|
||||
# If a specific model path is provided, use it directly
|
||||
if model_path and Path(model_path).exists():
|
||||
print(f"🎯 Using provided model path: {model_path}")
|
||||
return str(model_path)
|
||||
|
||||
# If no model path provided, extract model name from path or default to yolo11x
|
||||
model_name = "yolo11x" # Default fallback
|
||||
if model_path:
|
||||
# Try to extract model name from path
|
||||
path_obj = Path(model_path)
|
||||
if "yolo11n" in str(path_obj).lower():
|
||||
model_name = "yolo11n"
|
||||
elif "yolo11s" in str(path_obj).lower():
|
||||
model_name = "yolo11s"
|
||||
elif "yolo11m" in str(path_obj).lower():
|
||||
model_name = "yolo11m"
|
||||
elif "yolo11l" in str(path_obj).lower():
|
||||
model_name = "yolo11l"
|
||||
elif "yolo11x" in str(path_obj).lower():
|
||||
model_name = "yolo11x"
|
||||
|
||||
print(f"🔍 Searching for {model_name} model files...")
|
||||
|
||||
# Priority: quantized IR > IR > .pt
|
||||
search_paths = [
|
||||
Path(f"{model_name}_openvino_int8_model/{model_name}.xml") if use_quantized else None,
|
||||
Path(f"{model_name}_openvino_model/{model_name}.xml"),
|
||||
Path(f"rcb/{model_name}_openvino_model/{model_name}.xml"),
|
||||
Path(f"{model_name}.xml"),
|
||||
Path(f"rcb/{model_name}.xml"),
|
||||
Path(f"{model_name}.pt"),
|
||||
Path(f"rcb/{model_name}.pt")
|
||||
]
|
||||
|
||||
for p in search_paths:
|
||||
if p and p.exists():
|
||||
print(f"✅ Found model: {p}")
|
||||
return str(p)
|
||||
|
||||
# Fallback to any yolo11x if specific model not found
|
||||
fallback_paths = [
|
||||
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 fallback_paths:
|
||||
if p and p.exists():
|
||||
print(f"⚠️ Using fallback model: {p}")
|
||||
return str(p)
|
||||
|
||||
raise FileNotFoundError(f"No suitable {model_name} 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
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Test the detector with YOLOv11n model
|
||||
detector = OpenVINOVehicleDetector(model_path="yolo11n_openvino_model/yolo11n.xml")
|
||||
print(f"Detector initialized with model: {detector.model_path}")
|
||||
Reference in New Issue
Block a user