spdis/License_plate_recognition
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

yolorestart

Browse Source
This commit is contained in:
spdis
2025-08-30 11:53:20 +08:00
parent 090c44433f
commit 19176c71b7
26 changed files with 939 additions and 1825 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
yolopart/models/__init__.py
View File

@@ -1 +0,0 @@
# 模型模块初始化文件

490
yolopart/models/plate_recognizer.py
View File

@@ -1,490 +0,0 @@
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
车牌识别接口模块
预留接口，可接入各种OCR模型进行车牌号识别
"""
import cv2
import numpy as np
from typing import List, Optional, Dict, Any
from abc import ABC, abstractmethod
class PlateRecognizerInterface(ABC):
"""车牌识别接口基类"""
@abstractmethod
def recognize(self, plate_image: np.ndarray) -> Dict[str, Any]:
"""
识别车牌号
Args:
plate_image: 车牌图像 (BGR格式)
Returns:
识别结果字典，包含:
{
'text': str, # 识别的车牌号
'confidence': float, # 置信度 (0-1)
'success': bool # 是否识别成功
}
"""
pass
@abstractmethod
def batch_recognize(self, plate_images: List[np.ndarray]) -> List[Dict[str, Any]]:
"""
批量识别车牌号
Args:
plate_images: 车牌图像列表
Returns:
识别结果列表
"""
pass
class MockPlateRecognizer(PlateRecognizerInterface):
"""模拟车牌识别器（用于测试）"""
def __init__(self):
self.mock_plates = [
"京A12345", "沪B67890", "粤C11111", "川D22222",
"鲁E33333", "苏F44444", "浙G55555", "闽H66666"
]
self.call_count = 0
def recognize(self, plate_image: np.ndarray) -> Dict[str, Any]:
"""
模拟识别单个车牌
Args:
plate_image: 车牌图像
Returns:
模拟识别结果
"""
# 模拟处理时间
import time
time.sleep(0.01) # 10ms模拟处理时间
# 简单的图像质量检查
if plate_image is None or plate_image.size == 0:
return {
'text': '',
'confidence': 0.0,
'success': False
}
# 检查图像尺寸
height, width = plate_image.shape[:2]
if width < 50 or height < 20:
return {
'text': '',
'confidence': 0.3,
'success': False
}
# 模拟识别结果
plate_text = self.mock_plates[self.call_count % len(self.mock_plates)]
confidence = 0.85 + (self.call_count % 10) * 0.01 # 0.85-0.94
self.call_count += 1
return {
'text': plate_text,
'confidence': confidence,
'success': True
}
def batch_recognize(self, plate_images: List[np.ndarray]) -> List[Dict[str, Any]]:
"""
批量识别车牌
Args:
plate_images: 车牌图像列表
Returns:
识别结果列表
"""
results = []
for plate_image in plate_images:
result = self.recognize(plate_image)
results.append(result)
return results
class PaddleOCRRecognizer(PlateRecognizerInterface):
"""PaddleOCR车牌识别器（示例实现）"""
def __init__(self, use_gpu: bool = True):
"""
初始化PaddleOCR识别器
Args:
use_gpu: 是否使用GPU
"""
self.use_gpu = use_gpu
self.ocr = None
self._init_ocr()
def _init_ocr(self):
"""初始化OCR模型"""
try:
# 这里可以接入PaddleOCR
# from paddleocr import PaddleOCR
# self.ocr = PaddleOCR(use_angle_cls=True, lang='ch', use_gpu=self.use_gpu)
print("PaddleOCR初始化完成（示例代码，需要安装PaddleOCR）")
except ImportError:
print("PaddleOCR未安装，使用模拟识别器")
self.ocr = None
def recognize(self, plate_image: np.ndarray) -> Dict[str, Any]:
"""
使用PaddleOCR识别车牌
Args:
plate_image: 车牌图像
Returns:
识别结果
"""
if self.ocr is None:
# 回退到模拟识别
mock_recognizer = MockPlateRecognizer()
return mock_recognizer.recognize(plate_image)
try:
# 使用PaddleOCR进行识别
results = self.ocr.ocr(plate_image, cls=True)
if results and len(results) > 0 and results[0]:
# 提取文本和置信度
text_results = []
for line in results[0]:
text = line[1][0]
confidence = line[1][1]
text_results.append((text, confidence))
# 选择置信度最高的结果
if text_results:
best_result = max(text_results, key=lambda x: x[1])
return {
'text': best_result[0],
'confidence': best_result[1],
'success': True
}
except Exception as e:
print(f"PaddleOCR识别失败: {e}")
return {
'text': '',
'confidence': 0.0,
'success': False
}
def batch_recognize(self, plate_images: List[np.ndarray]) -> List[Dict[str, Any]]:
"""
批量识别
Args:
plate_images: 车牌图像列表
Returns:
识别结果列表
"""
results = []
for plate_image in plate_images:
result = self.recognize(plate_image)
results.append(result)
return results
class TesseractRecognizer(PlateRecognizerInterface):
"""Tesseract车牌识别器（示例实现）"""
def __init__(self, lang: str = 'chi_sim+eng'):
"""
初始化Tesseract识别器
Args:
lang: 识别语言
"""
self.lang = lang
self.tesseract_available = self._check_tesseract()
def _check_tesseract(self) -> bool:
"""检查Tesseract是否可用"""
try:
import pytesseract
return True
except ImportError:
print("pytesseract未安装，使用模拟识别器")
return False
def recognize(self, plate_image: np.ndarray) -> Dict[str, Any]:
"""
使用Tesseract识别车牌
Args:
plate_image: 车牌图像
Returns:
识别结果
"""
if not self.tesseract_available:
# 回退到模拟识别
mock_recognizer = MockPlateRecognizer()
return mock_recognizer.recognize(plate_image)
try:
import pytesseract
# 图像预处理
processed_image = self._preprocess_image(plate_image)
# 使用Tesseract识别
text = pytesseract.image_to_string(
processed_image,
lang=self.lang,
config='--psm 8 --oem 3 -c tessedit_char_whitelist=0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ京沪粤川鲁苏浙闽'
)
# 清理识别结果
text = text.strip().replace(' ', '').replace('\n', '')
if text and len(text) >= 5: # 车牌号至少5位
return {
'text': text,
'confidence': 0.8, # Tesseract不直接提供置信度
'success': True
}
except Exception as e:
print(f"Tesseract识别失败: {e}")
return {
'text': '',
'confidence': 0.0,
'success': False
}
def _preprocess_image(self, image: np.ndarray) -> np.ndarray:
"""图像预处理"""
# 转换为灰度图
if len(image.shape) == 3:
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
else:
gray = image
# 调整尺寸
height, width = gray.shape
if width < 200:
scale = 200 / width
new_width = int(width * scale)
new_height = int(height * scale)
gray = cv2.resize(gray, (new_width, new_height))
# 二值化
_, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
return binary
def batch_recognize(self, plate_images: List[np.ndarray]) -> List[Dict[str, Any]]:
"""
批量识别
Args:
plate_images: 车牌图像列表
Returns:
识别结果列表
"""
results = []
for plate_image in plate_images:
result = self.recognize(plate_image)
results.append(result)
return results
class PlateRecognizerManager:
"""车牌识别管理器"""
def __init__(self, recognizer_type: str = 'mock'):
"""
初始化识别管理器
Args:
recognizer_type: 识别器类型 ('mock', 'paddleocr', 'tesseract')
"""
self.recognizer_type = recognizer_type
self.recognizer = self._create_recognizer(recognizer_type)
def _create_recognizer(self, recognizer_type: str) -> PlateRecognizerInterface:
"""创建识别器"""
if recognizer_type == 'mock':
return MockPlateRecognizer()
elif recognizer_type == 'paddleocr':
return PaddleOCRRecognizer()
elif recognizer_type == 'tesseract':
return TesseractRecognizer()
else:
print(f"未知的识别器类型: {recognizer_type}，使用模拟识别器")
return MockPlateRecognizer()
def recognize_plates(self, plate_images: List[np.ndarray]) -> List[Dict[str, Any]]:
"""
识别车牌列表
Args:
plate_images: 车牌图像列表
Returns:
识别结果列表
"""
if not plate_images:
return []
return self.recognizer.batch_recognize(plate_images)
def switch_recognizer(self, recognizer_type: str):
"""
切换识别器
Args:
recognizer_type: 新的识别器类型
"""
self.recognizer_type = recognizer_type
self.recognizer = self._create_recognizer(recognizer_type)
print(f"已切换到识别器: {recognizer_type}")
def get_recognizer_info(self) -> Dict[str, Any]:
"""
获取识别器信息
Returns:
识别器信息
"""
return {
'type': self.recognizer_type,
'class': self.recognizer.__class__.__name__
}
def preprocess_blue_plate(self, plate_image: np.ndarray, original_image: np.ndarray, bbox: List[int]) -> np.ndarray:
"""
蓝色车牌预处理：倾斜矫正
Args:
plate_image: 切割后的车牌图像
original_image: 原始图像
bbox: 边界框坐标 [x1, y1, x2, y2]
Returns:
矫正后的车牌图像
"""
try:
# 从原图中提取车牌区域
x1, y1, x2, y2 = bbox
roi = original_image[y1:y2, x1:x2]
# 获取蓝色车牌的二值图像
bin_img = self._get_blue_img_bin(roi)
# 倾斜矫正
corrected_img = self._deskew_plate(bin_img, roi)
return corrected_img
except Exception as e:
print(f"蓝色车牌预处理失败: {e}")
return plate_image
def _get_blue_img_bin(self, img: np.ndarray) -> np.ndarray:
"""
获取蓝色车牌的二值图像
"""
# 掩膜：BGR通道，若像素B分量在 100~255 且 G分量在 0~190 且 R分量在 0~140 置255（白色），否则置0（黑色）
mask_bgr = cv2.inRange(img, (100, 0, 0), (255, 190, 140))
# 转换成 HSV 颜色空间
img_hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
h, s, v = cv2.split(img_hsv) # 分离通道 色调(H)，饱和度(S)，明度(V)
mask_s = cv2.inRange(s, 80, 255) # 取饱和度通道进行掩膜得到二值图像
# 与操作，两个二值图像都为白色才保留，否则置黑
rgbs = mask_bgr & mask_s
# 核的横向分量大，使车牌数字尽量连在一起
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (15, 3))
img_rgbs_dilate = cv2.dilate(rgbs, kernel, 3) # 膨胀，减小车牌空洞
return img_rgbs_dilate
def _order_points(self, pts: np.ndarray) -> np.ndarray:
"""
将四点按 左上、右上、右下、左下 排序
"""
rect = np.zeros((4, 2), dtype="float32")
s = pts.sum(axis=1)
rect[0] = pts[np.argmin(s)] # 左上
rect[2] = pts[np.argmax(s)] # 右下
diff = np.diff(pts, axis=1)
rect[1] = pts[np.argmin(diff)] # 右上
rect[3] = pts[np.argmax(diff)] # 左下
return rect
def _deskew_plate(self, bin_img: np.ndarray, original_roi: np.ndarray) -> np.ndarray:
"""
车牌倾斜矫正
Args:
bin_img: 二值图像
original_roi: 原始ROI区域
Returns:
矫正后的原始图像（未被掩模，但经过旋转和切割）
"""
try:
# 找最大轮廓
cnts, _ = cv2.findContours(bin_img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
if not cnts:
return original_roi
c = max(cnts, key=cv2.contourArea)
# 最小外接矩形
rect = cv2.minAreaRect(c)
box = cv2.boxPoints(rect)
box = np.array(box, dtype="float32")
# 排序四个点
pts_src = self._order_points(box)
# 计算目标矩形宽高
(tl, tr, br, bl) = pts_src
widthA = np.linalg.norm(br - bl)
widthB = np.linalg.norm(tr - tl)
maxWidth = int(max(widthA, widthB))
heightA = np.linalg.norm(tr - br)
heightB = np.linalg.norm(tl - bl)
maxHeight = int(max(heightA, heightB))
# 确保尺寸合理
if maxWidth < 10 or maxHeight < 10:
return original_roi
# 目标点集合
pts_dst = np.array([
[0, 0],
[maxWidth - 1, 0],
[maxWidth - 1, maxHeight - 1],
[0, maxHeight - 1]], dtype="float32")
# 透视变换
M = cv2.getPerspectiveTransform(pts_src, pts_dst)
warped = cv2.warpPerspective(original_roi, M, (maxWidth, maxHeight))
return warped
except Exception as e:
print(f"车牌矫正失败: {e}")
return original_roi

368
yolopart/models/yolo_detector.py
View File

@@ -1,368 +0,0 @@
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
YOLO车牌检测器
基于ONNX Runtime的YOLO11s模型推理
"""
import cv2
import numpy as np
import onnxruntime as ort
import time
from typing import List, Tuple, Optional
class YOLODetector:
"""YOLO车牌检测器"""
def __init__(self, model_path: str, conf_threshold: float = 0.25, nms_threshold: float = 0.4):
"""
初始化YOLO检测器
Args:
model_path: ONNX模型文件路径
conf_threshold: 置信度阈值
nms_threshold: NMS阈值
"""
self.model_path = model_path
self.conf_threshold = conf_threshold
self.nms_threshold = nms_threshold
self.input_size = (640, 640) # YOLO11s输入尺寸
self.use_gpu = False
# 初始化ONNX Runtime会话
self._init_session()
# 获取模型输入输出信息
self.input_name = self.session.get_inputs()[0].name
self.output_names = [output.name for output in self.session.get_outputs()]
print(f"YOLO检测器初始化完成")
print(f"模型路径: {model_path}")
print(f"输入尺寸: {self.input_size}")
print(f"GPU加速: {self.use_gpu}")
def _init_session(self):
"""初始化ONNX Runtime会话"""
# 获取可用的providers
available_providers = ort.get_available_providers()
print(f"可用的执行提供者: {available_providers}")
# 优先使用GPU，如果可用的话
providers = []
if 'CUDAExecutionProvider' in available_providers:
providers.append('CUDAExecutionProvider')
self.use_gpu = True
print("检测到CUDA支持，将使用GPU加速")
elif 'TensorrtExecutionProvider' in available_providers:
providers.append('TensorrtExecutionProvider')
self.use_gpu = True
print("检测到TensorRT支持，将使用GPU加速")
else:
self.use_gpu = False
print("未检测到GPU支持，将使用CPU")
# 添加CPU作为备选
providers.append('CPUExecutionProvider')
print(f"使用的执行提供者: {providers}")
# 创建会话
session_options = ort.SessionOptions()
session_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
try:
self.session = ort.InferenceSession(
self.model_path,
sess_options=session_options,
providers=providers
)
# 检查实际使用的provider
actual_providers = self.session.get_providers()
print(f"实际使用的执行提供者: {actual_providers}")
if 'CUDAExecutionProvider' in actual_providers or 'TensorrtExecutionProvider' in actual_providers:
self.use_gpu = True
print("✅ GPU加速已启用")
else:
self.use_gpu = False
print("⚠️ 使用CPU执行")
except Exception as e:
print(f"模型加载失败: {e}")
raise
def preprocess(self, image: np.ndarray) -> Tuple[np.ndarray, float, float]:
"""
图像预处理
Args:
image: 输入图像 (BGR格式)
Returns:
preprocessed_image: 预处理后的图像
scale_x: X轴缩放比例
scale_y: Y轴缩放比例
"""
original_height, original_width = image.shape[:2]
target_width, target_height = self.input_size
# 计算缩放比例
scale_x = target_width / original_width
scale_y = target_height / original_height
# 等比例缩放
scale = min(scale_x, scale_y)
new_width = int(original_width * scale)
new_height = int(original_height * scale)
# 缩放图像
resized_image = cv2.resize(image, (new_width, new_height))
# 创建目标尺寸的图像并居中放置
padded_image = np.full((target_height, target_width, 3), 114, dtype=np.uint8)
# 计算填充位置
start_x = (target_width - new_width) // 2
start_y = (target_height - new_height) // 2
padded_image[start_y:start_y + new_height, start_x:start_x + new_width] = resized_image
# 转换为RGB并归一化
rgb_image = cv2.cvtColor(padded_image, cv2.COLOR_BGR2RGB)
normalized_image = rgb_image.astype(np.float32) / 255.0
# 转换为NCHW格式
input_tensor = np.transpose(normalized_image, (2, 0, 1))
input_tensor = np.expand_dims(input_tensor, axis=0)
return input_tensor, scale, scale
def postprocess(self, outputs: List[np.ndarray], scale_x: float, scale_y: float,
original_shape: Tuple[int, int]) -> List[dict]:
"""
后处理检测结果
Args:
outputs: 模型输出
scale_x: X轴缩放比例
scale_y: Y轴缩放比例
original_shape: 原始图像尺寸 (height, width)
Returns:
检测结果列表
"""
detections = []
if len(outputs) == 0:
return detections
# 获取输出张量
output = outputs[0]
# YOLO11输出格式: [batch, 6, 8400] -> [batch, 8400, 6]
if len(output.shape) == 3:
output = output.transpose(0, 2, 1)
# 处理每个检测结果
for detection in output[0]: # 取第一个batch
# 前4个值是边界框坐标，后2个是类别概率
x_center, y_center, width, height = detection[:4]
class_scores = detection[4:] # 类别概率 [蓝牌概率, 绿牌概率]
# 获取最高概率的类别
class_id = np.argmax(class_scores)
confidence = class_scores[class_id] # 使用类别概率作为置信度
# 过滤低置信度检测
if confidence < self.conf_threshold:
continue
# 转换坐标到原始图像尺寸
original_height, original_width = original_shape
# 计算实际缩放比例和偏移
scale = min(self.input_size[0] / original_width, self.input_size[1] / original_height)
pad_x = (self.input_size[0] - original_width * scale) / 2
pad_y = (self.input_size[1] - original_height * scale) / 2
# 转换坐标
x_center = (x_center - pad_x) / scale
y_center = (y_center - pad_y) / scale
width = width / scale
height = height / scale
# 计算边界框
x1 = int(x_center - width / 2)
y1 = int(y_center - height / 2)
x2 = int(x_center + width / 2)
y2 = int(y_center + height / 2)
# 确保坐标在图像范围内
x1 = max(0, min(x1, original_width - 1))
y1 = max(0, min(y1, original_height - 1))
x2 = max(0, min(x2, original_width - 1))
y2 = max(0, min(y2, original_height - 1))
# 定义类别名称
class_names = ['blue_plate', 'green_plate'] # 0: 蓝牌, 1: 绿牌
class_name = class_names[class_id] if class_id < len(class_names) else 'unknown'
detections.append({
'bbox': [x1, y1, x2, y2],
'confidence': float(confidence),
'class_id': int(class_id),
'class_name': class_name
})
# 应用NMS
if detections:
detections = self._apply_nms(detections)
return detections
def _apply_nms(self, detections: List[dict]) -> List[dict]:
"""
应用非极大值抑制
Args:
detections: 检测结果列表
Returns:
NMS后的检测结果
"""
if len(detections) == 0:
return detections
# 提取边界框和置信度
boxes = np.array([det['bbox'] for det in detections])
scores = np.array([det['confidence'] for det in detections])
# 应用NMS
indices = cv2.dnn.NMSBoxes(
boxes.tolist(),
scores.tolist(),
self.conf_threshold,
self.nms_threshold
)
# 返回保留的检测结果
if len(indices) > 0:
indices = indices.flatten()
return [detections[i] for i in indices]
else:
return []
def detect(self, image: np.ndarray) -> List[dict]:
"""
检测车牌
Args:
image: 输入图像 (BGR格式)
Returns:
检测结果列表
"""
try:
# 预处理
input_tensor, scale_x, scale_y = self.preprocess(image)
# 推理
outputs = self.session.run(self.output_names, {self.input_name: input_tensor})
# 调试输出
print(f"模型输出数量: {len(outputs)}")
for i, output in enumerate(outputs):
print(f"输出 {i} 形状: {output.shape}")
print(f"输出 {i} 数据范围: [{output.min():.4f}, {output.max():.4f}]")
# 后处理
detections = self.postprocess(outputs, scale_x, scale_y, image.shape[:2])
print(f"检测到的目标数量: {len(detections)}")
for i, det in enumerate(detections):
print(f"检测 {i}: 类别={det['class_name']}, 置信度={det['confidence']:.3f}")
return detections
except Exception as e:
print(f"检测过程出错: {e}")
return []
def draw_detections(self, image: np.ndarray, detections: List[dict]) -> np.ndarray:
"""
在图像上绘制检测结果
Args:
image: 输入图像
detections: 检测结果列表
Returns:
绘制了检测框的图像
"""
result_image = image.copy()
for detection in detections:
bbox = detection['bbox']
confidence = detection['confidence']
class_id = detection['class_id']
class_name = detection['class_name']
x1, y1, x2, y2 = bbox
# 根据车牌类型选择颜色
if class_id == 0: # 蓝牌
color = (255, 0, 0) # 蓝色 (BGR格式)
plate_type = "Blue Plate"
elif class_id == 1: # 绿牌
color = (0, 255, 0) # 绿色 (BGR格式)
plate_type = "Green Plate"
else:
color = (0, 255, 255) # 黄色 (BGR格式)
plate_type = "Unknown"
# 绘制边界框
cv2.rectangle(result_image, (x1, y1), (x2, y2), color, 2)
# 绘制置信度标签
label = f"{plate_type}: {confidence:.2f}"
label_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)[0]
# 绘制标签背景
cv2.rectangle(result_image,
(x1, y1 - label_size[1] - 10),
(x1 + label_size[0], y1),
color, -1)
# 绘制标签文字
cv2.putText(result_image, label,
(x1, y1 - 5),
cv2.FONT_HERSHEY_SIMPLEX, 0.6,
(255, 255, 255), 2)
return result_image
def crop_plates(self, image: np.ndarray, detections: List[dict]) -> List[np.ndarray]:
"""
切割车牌图像
Args:
image: 原始图像
detections: 检测结果列表
Returns:
切割后的车牌图像列表
"""
plate_images = []
for detection in detections:
bbox = detection['bbox']
x1, y1, x2, y2 = bbox
# 确保坐标有效
if x2 > x1 and y2 > y1:
# 切割车牌区域
plate_image = image[y1:y2, x1:x2]
if plate_image.size > 0:
plate_images.append(plate_image)
return plate_images