348 lines
13 KiB
Python
348 lines
13 KiB
Python
"""
|
|
车辆检测服务适配器
|
|
支持车辆检测、跟踪和违停检测功能（基于YOLOv8）
|
|
"""
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import logging
|
|
import threading
|
|
import time
|
|
from typing import Dict, List, Optional, Tuple
|
|
from dataclasses import dataclass
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class VehicleTrackingInfo:
|
|
"""车辆跟踪信息"""
|
|
track_id: int
|
|
bbox: List[float]
|
|
center: Tuple[float, float]
|
|
first_seen: float
|
|
last_seen: float
|
|
plate_number: Optional[str] = None
|
|
is_illegal_parking: bool = False
|
|
trajectory: List[Tuple[float, float]] = None
|
|
|
|
def __post_init__(self):
|
|
if self.trajectory is None:
|
|
self.trajectory = []
|
|
|
|
|
|
class VehicleDetectionService:
|
|
"""车辆检测服务（YOLOv8模式）"""
|
|
|
|
# COCO数据集中车辆相关类别映射
|
|
VEHICLE_CLASSES = {2: 'car', 3: 'motorcycle', 5: 'bus', 7: 'truck'}
|
|
VEHICLE_LABELS = {'car': '小汽车', 'motorcycle': '摩托车', 'bus': '公交车', 'truck': '卡车'}
|
|
|
|
def __init__(self, yolo_model=None):
|
|
self.model_name = "vehicle_detection"
|
|
self.threshold = 0.1
|
|
self._lock = threading.Lock()
|
|
self._yolo_model = yolo_model
|
|
|
|
# 车辆跟踪信息
|
|
self.vehicle_tracks: Dict[int, VehicleTrackingInfo] = {}
|
|
self.track_id_counter = 0
|
|
|
|
# 违停检测配置
|
|
self.illegal_parking_time = 5.0
|
|
self.illegal_parking_region = None
|
|
|
|
self.available = yolo_model is not None
|
|
if self.available:
|
|
logger.info("车辆检测服务初始化完成（YOLOv8模式）")
|
|
else:
|
|
logger.error("车辆检测服务初始化失败：未提供YOLO模型")
|
|
|
|
def detect_image(self, image: np.ndarray, threshold: float = None) -> Dict:
|
|
"""
|
|
检测图片中的车辆
|
|
|
|
Args:
|
|
image: OpenCV 图片 (BGR格式)
|
|
threshold: 置信度阈值
|
|
|
|
Returns:
|
|
检测结果字典
|
|
"""
|
|
if threshold is None:
|
|
threshold = self.threshold
|
|
|
|
if not self.available or self._yolo_model is None:
|
|
return {
|
|
'success': False,
|
|
'message': '车辆检测服务不可用',
|
|
'detections': [],
|
|
'stats': None
|
|
}
|
|
|
|
try:
|
|
with self._lock:
|
|
start_time = time.time()
|
|
|
|
# 使用YOLO模型推理
|
|
inference_start = time.time()
|
|
results = self._yolo_model(image, conf=threshold, verbose=False)
|
|
inference_time = time.time() - inference_start
|
|
|
|
# 解析检测结果，只保留车辆类别
|
|
detections = []
|
|
result = results[0]
|
|
if result.boxes is not None:
|
|
for box in result.boxes:
|
|
cls_id = int(box.cls[0])
|
|
if cls_id in self.VEHICLE_CLASSES:
|
|
confidence = float(box.conf[0])
|
|
x1, y1, x2, y2 = box.xyxy[0].tolist()
|
|
|
|
center_x = (x1 + x2) / 2
|
|
center_y = (y1 + y2) / 2
|
|
|
|
class_name = self.VEHICLE_CLASSES[cls_id]
|
|
label = self.VEHICLE_LABELS[class_name]
|
|
|
|
detections.append({
|
|
'class': class_name,
|
|
'label': label,
|
|
'confidence': round(confidence, 3),
|
|
'bbox': [int(x1), int(y1), int(x2), int(y2)],
|
|
'center': [round(center_x, 2), round(center_y, 2)]
|
|
})
|
|
|
|
total_time = time.time() - start_time
|
|
|
|
return {
|
|
'success': True,
|
|
'message': '检测完成',
|
|
'detections': detections,
|
|
'stats': {
|
|
'total_detections': len(detections),
|
|
'model_used': 'yolov8s',
|
|
'threshold': threshold,
|
|
'processing_time': round(total_time, 3),
|
|
'inference_time': round(inference_time, 3)
|
|
}
|
|
}
|
|
|
|
except Exception as e:
|
|
import traceback
|
|
logger.error(f"检测失败: {e}")
|
|
logger.error(f"错误堆栈: {traceback.format_exc()}")
|
|
return {
|
|
'success': False,
|
|
'message': f'检测失败: {e}',
|
|
'detections': [],
|
|
'stats': None
|
|
}
|
|
|
|
def detect_illegal_parking(self, image: np.ndarray, threshold: float = None,
|
|
illegal_parking_time: float = 5.0,
|
|
region_polygon: List[Tuple[int, int]] = None) -> Dict:
|
|
"""
|
|
检测违停车辆
|
|
|
|
Args:
|
|
image: OpenCV 图片
|
|
threshold: 置信度阈值
|
|
illegal_parking_time: 违停时间阈值（秒）
|
|
region_polygon: 违停区域多边形点集 [(x1,y1), (x2,y2), ...]
|
|
|
|
Returns:
|
|
违停检测结果
|
|
"""
|
|
if threshold is None:
|
|
threshold = self.threshold
|
|
|
|
# 规范化 region_polygon 格式（支持 [{x,y},...] 和 [[x,y],...] 两种格式）
|
|
if region_polygon:
|
|
normalized = []
|
|
for p in region_polygon:
|
|
if isinstance(p, dict):
|
|
normalized.append((int(p['x']), int(p['y'])))
|
|
elif isinstance(p, (list, tuple)) and len(p) >= 2:
|
|
normalized.append((int(p[0]), int(p[1])))
|
|
region_polygon = normalized
|
|
|
|
# 更新违停配置
|
|
self.illegal_parking_time = illegal_parking_time
|
|
self.illegal_parking_region = region_polygon
|
|
|
|
# 基础车辆检测
|
|
detection_result = self.detect_image(image, threshold)
|
|
|
|
if not detection_result['success']:
|
|
return {
|
|
'success': False,
|
|
'message': detection_result['message'],
|
|
'illegal_parking': [],
|
|
'vehicles': []
|
|
}
|
|
|
|
current_time = time.time()
|
|
current_detections = detection_result['detections']
|
|
|
|
# 更新车辆跟踪信息
|
|
illegal_parking_vehicles = []
|
|
|
|
for detection in current_detections:
|
|
bbox = detection['bbox']
|
|
center = detection['center']
|
|
|
|
# 简单的跟踪（基于位置匹配）
|
|
matched_track_id = self._match_vehicle_to_track(center, bbox)
|
|
|
|
if matched_track_id is None:
|
|
# 新车辆
|
|
self.track_id_counter += 1
|
|
matched_track_id = self.track_id_counter
|
|
self.vehicle_tracks[matched_track_id] = VehicleTrackingInfo(
|
|
track_id=matched_track_id,
|
|
bbox=bbox,
|
|
center=center,
|
|
first_seen=current_time,
|
|
last_seen=current_time,
|
|
trajectory=[center]
|
|
)
|
|
else:
|
|
# 更新现有车辆
|
|
track_info = self.vehicle_tracks[matched_track_id]
|
|
track_info.bbox = bbox
|
|
track_info.center = center
|
|
track_info.last_seen = current_time
|
|
track_info.trajectory.append(center)
|
|
|
|
# 检查违停条件（对新车辆和已有车辆都检查）
|
|
track_info = self.vehicle_tracks[matched_track_id]
|
|
if self._check_illegal_parking(track_info, region_polygon):
|
|
track_info.is_illegal_parking = True
|
|
illegal_parking_vehicles.append({
|
|
'track_id': matched_track_id,
|
|
'bbox': bbox,
|
|
'center': center,
|
|
'parking_duration': round(current_time - track_info.first_seen, 2),
|
|
'plate_number': track_info.plate_number
|
|
})
|
|
|
|
# 清理长时间未出现的车辆
|
|
self._cleanup_old_tracks(current_time)
|
|
|
|
return {
|
|
'success': True,
|
|
'message': '违停检测完成',
|
|
'illegal_parking': illegal_parking_vehicles,
|
|
'total_vehicles': len(current_detections),
|
|
'stats': detection_result['stats']
|
|
}
|
|
|
|
def _match_vehicle_to_track(self, center: Tuple[float, float],
|
|
bbox: List[float]) -> Optional[int]:
|
|
"""将检测到的车辆匹配到已有轨迹"""
|
|
x, y = center
|
|
|
|
for track_id, track_info in self.vehicle_tracks.items():
|
|
track_x, track_y = track_info.center
|
|
|
|
# 计算距离
|
|
distance = np.sqrt((x - track_x) ** 2 + (y - track_y) ** 2)
|
|
|
|
# 距离阈值（基于检测框大小）
|
|
bbox_width = bbox[2] - bbox[0]
|
|
bbox_height = bbox[3] - bbox[1]
|
|
max_dim = max(bbox_width, bbox_height)
|
|
|
|
if distance < max_dim * 0.5: # 距离小于检测框最大尺寸的一半
|
|
return track_id
|
|
|
|
return None
|
|
|
|
def _check_illegal_parking(self, track_info: VehicleTrackingInfo,
|
|
region_polygon: List[Tuple[int, int]] = None) -> bool:
|
|
"""检查是否违停"""
|
|
current_time = time.time()
|
|
parking_duration = current_time - track_info.first_seen
|
|
|
|
# 检查时间是否超过阈值
|
|
if parking_duration < self.illegal_parking_time:
|
|
return False
|
|
|
|
# 如果指定了违停区域，检查车辆中心是否在多边形内
|
|
if region_polygon is not None:
|
|
return self._point_in_polygon(track_info.center, region_polygon)
|
|
|
|
# 未指定区域时，无法判定是否违停（避免全量误报）
|
|
return False
|
|
|
|
def _point_in_polygon(self, point: Tuple[float, float],
|
|
polygon: List[Tuple[int, int]]) -> bool:
|
|
"""判断点是否在多边形内（射线法）"""
|
|
x, y = point
|
|
n = len(polygon)
|
|
inside = False
|
|
|
|
p1x, p1y = polygon[0]
|
|
for i in range(n + 1):
|
|
p2x, p2y = polygon[i % n]
|
|
if y > min(p1y, p2y):
|
|
if y <= max(p1y, p2y):
|
|
if x <= max(p1x, p2x):
|
|
if p1y != p2y:
|
|
xinters = (y - p1y) * (p2x - p1x) / (p2y - p1y) + p1x
|
|
if p1x == p2x or x <= xinters:
|
|
inside = not inside
|
|
p1x, p1y = p2x, p2y
|
|
|
|
return inside
|
|
|
|
def _cleanup_old_tracks(self, current_time: float):
|
|
"""清理长时间未出现的车辆轨迹"""
|
|
timeout = 10.0 # 10秒未出现则删除
|
|
|
|
tracks_to_remove = []
|
|
for track_id, track_info in self.vehicle_tracks.items():
|
|
if current_time - track_info.last_seen > timeout:
|
|
tracks_to_remove.append(track_id)
|
|
|
|
for track_id in tracks_to_remove:
|
|
del self.vehicle_tracks[track_id]
|
|
logger.debug(f"清理车辆轨迹: {track_id}")
|
|
|
|
def get_performance_info(self) -> Dict:
|
|
"""获取性能信息"""
|
|
return {
|
|
'mode': 'local',
|
|
'environment': 'YOLOv8',
|
|
'detector_loaded': self._yolo_model is not None,
|
|
'available': self.available,
|
|
'active_tracks': len(self.vehicle_tracks)
|
|
}
|
|
|
|
|
|
# 兼容性包装，保持与 YOLO 模型相同的接口
|
|
class VehicleDetectionModel:
|
|
"""车辆检测模型包装器，兼容 YOLO 接口"""
|
|
|
|
def __init__(self, yolo_model=None):
|
|
self._yolo_model = yolo_model
|
|
self.service = VehicleDetectionService(yolo_model=yolo_model)
|
|
# 保留 names 以兼容 YOLO 接口（实际类别由底层YOLO模型提供）
|
|
self.names = {0: 'vehicle'}
|
|
|
|
def __call__(self, image, conf=0.1, iou=0.45, verbose=False):
|
|
"""
|
|
直接委托给底层YOLO模型，返回标准YOLO结果
|
|
"""
|
|
if self._yolo_model is not None:
|
|
return self._yolo_model(image, conf=conf, iou=iou, verbose=verbose)
|
|
raise RuntimeError("YOLO模型未初始化")
|
|
|
|
def detect_illegal_parking(self, image, conf=0.1, illegal_parking_time=5.0,
|
|
region_polygon=None):
|
|
"""违停检测接口"""
|
|
return self.service.detect_illegal_parking(
|
|
image, conf, illegal_parking_time, region_polygon
|
|
)
|