Initial commit
This commit is contained in:
490
yolopart/models/plate_recognizer.py
Normal file
490
yolopart/models/plate_recognizer.py
Normal file
@@ -0,0 +1,490 @@
|
||||
#!/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
|
||||
Reference in New Issue
Block a user