490 lines
15 KiB
Python
490 lines
15 KiB
Python
#!/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 |