436 lines
15 KiB
Python
436 lines
15 KiB
Python
import sys
|
||
import cv2
|
||
import numpy as np
|
||
from PyQt5.QtWidgets import (
|
||
QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout,
|
||
QLabel, QPushButton, QScrollArea, QFrame, QSizePolicy
|
||
)
|
||
from PyQt5.QtCore import QTimer, Qt, pyqtSignal, QThread
|
||
from PyQt5.QtGui import QImage, QPixmap, QFont, QPainter, QPen, QColor
|
||
import os
|
||
from yolopart.detector import LicensePlateYOLO
|
||
|
||
#选择使用哪个模块
|
||
from LPRNET_part.lpr_interface import LPRNmodel_predict
|
||
from LPRNET_part.lpr_interface import LPRNinitialize_model
|
||
|
||
#使用OCR
|
||
#from OCR_part.ocr_interface import LPRNmodel_predict
|
||
#from OCR_part.ocr_interface import LPRNinitialize_model
|
||
# 使用CRNN
|
||
#from CRNN_part.crnn_interface import LPRNmodel_predict
|
||
#from CRNN_part.crnn_interface import LPRNinitialize_model
|
||
|
||
class CameraThread(QThread):
|
||
"""摄像头线程类"""
|
||
frame_ready = pyqtSignal(np.ndarray)
|
||
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.camera = None
|
||
self.running = False
|
||
|
||
def start_camera(self):
|
||
"""启动摄像头"""
|
||
self.camera = cv2.VideoCapture(0)
|
||
if self.camera.isOpened():
|
||
self.running = True
|
||
self.start()
|
||
return True
|
||
return False
|
||
|
||
def stop_camera(self):
|
||
"""停止摄像头"""
|
||
self.running = False
|
||
if self.camera:
|
||
self.camera.release()
|
||
self.quit()
|
||
self.wait()
|
||
|
||
def run(self):
|
||
"""线程运行函数"""
|
||
while self.running:
|
||
if self.camera and self.camera.isOpened():
|
||
ret, frame = self.camera.read()
|
||
if ret:
|
||
self.frame_ready.emit(frame)
|
||
self.msleep(30) # 约30fps
|
||
|
||
class LicensePlateWidget(QWidget):
|
||
"""单个车牌结果显示组件"""
|
||
|
||
def __init__(self, plate_id, class_name, corrected_image, plate_number):
|
||
super().__init__()
|
||
self.plate_id = plate_id
|
||
self.init_ui(class_name, corrected_image, plate_number)
|
||
|
||
def init_ui(self, class_name, corrected_image, plate_number):
|
||
layout = QHBoxLayout()
|
||
layout.setContentsMargins(10, 5, 10, 5)
|
||
|
||
# 车牌类型标签
|
||
type_label = QLabel(class_name)
|
||
type_label.setFixedWidth(60)
|
||
type_label.setAlignment(Qt.AlignCenter)
|
||
type_label.setStyleSheet(
|
||
"QLabel { "
|
||
"background-color: #4CAF50 if class_name == '绿牌' else #2196F3; "
|
||
"color: white; "
|
||
"border-radius: 5px; "
|
||
"padding: 5px; "
|
||
"font-weight: bold; "
|
||
"}"
|
||
)
|
||
if class_name == '绿牌':
|
||
type_label.setStyleSheet(
|
||
"QLabel { "
|
||
"background-color: #4CAF50; "
|
||
"color: white; "
|
||
"border-radius: 5px; "
|
||
"padding: 5px; "
|
||
"font-weight: bold; "
|
||
"}"
|
||
)
|
||
else:
|
||
type_label.setStyleSheet(
|
||
"QLabel { "
|
||
"background-color: #2196F3; "
|
||
"color: white; "
|
||
"border-radius: 5px; "
|
||
"padding: 5px; "
|
||
"font-weight: bold; "
|
||
"}"
|
||
)
|
||
|
||
# 矫正后的车牌图像
|
||
image_label = QLabel()
|
||
image_label.setFixedSize(120, 40)
|
||
image_label.setStyleSheet("border: 1px solid #ddd; background-color: white;")
|
||
|
||
if corrected_image is not None:
|
||
# 转换numpy数组为QPixmap
|
||
h, w = corrected_image.shape[:2]
|
||
if len(corrected_image.shape) == 3:
|
||
bytes_per_line = 3 * w
|
||
q_image = QImage(corrected_image.data, w, h, bytes_per_line, QImage.Format_RGB888).rgbSwapped()
|
||
else:
|
||
bytes_per_line = w
|
||
q_image = QImage(corrected_image.data, w, h, bytes_per_line, QImage.Format_Grayscale8)
|
||
|
||
pixmap = QPixmap.fromImage(q_image)
|
||
scaled_pixmap = pixmap.scaled(120, 40, Qt.KeepAspectRatio, Qt.SmoothTransformation)
|
||
image_label.setPixmap(scaled_pixmap)
|
||
else:
|
||
image_label.setText("车牌未完全\n进入摄像头")
|
||
image_label.setAlignment(Qt.AlignCenter)
|
||
image_label.setStyleSheet("border: 1px solid #ddd; background-color: #f5f5f5; color: #666;")
|
||
|
||
# 车牌号标签
|
||
number_label = QLabel(plate_number)
|
||
number_label.setFixedWidth(150)
|
||
number_label.setAlignment(Qt.AlignCenter)
|
||
number_label.setStyleSheet(
|
||
"QLabel { "
|
||
"border: 1px solid #ddd; "
|
||
"background-color: white; "
|
||
"padding: 8px; "
|
||
"font-family: 'Courier New'; "
|
||
"font-size: 14px; "
|
||
"font-weight: bold; "
|
||
"}"
|
||
)
|
||
|
||
layout.addWidget(type_label)
|
||
layout.addWidget(image_label)
|
||
layout.addWidget(number_label)
|
||
layout.addStretch()
|
||
|
||
self.setLayout(layout)
|
||
self.setStyleSheet(
|
||
"QWidget { "
|
||
"background-color: white; "
|
||
"border: 1px solid #e0e0e0; "
|
||
"border-radius: 8px; "
|
||
"margin: 2px; "
|
||
"}"
|
||
)
|
||
|
||
class MainWindow(QMainWindow):
|
||
"""主窗口类"""
|
||
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.detector = None
|
||
self.camera_thread = None
|
||
self.current_frame = None
|
||
self.detections = []
|
||
|
||
self.init_ui()
|
||
self.init_detector()
|
||
self.init_camera()
|
||
|
||
# 初始化OCR/CRNN模型（函数名改成一样的了，所以不要修改这里了，想用哪个模块直接导入）
|
||
LPRNinitialize_model()
|
||
|
||
|
||
def init_ui(self):
|
||
"""初始化用户界面"""
|
||
self.setWindowTitle("车牌识别系统")
|
||
self.setGeometry(100, 100, 1200, 800)
|
||
|
||
# 创建中央widget
|
||
central_widget = QWidget()
|
||
self.setCentralWidget(central_widget)
|
||
|
||
# 创建主布局
|
||
main_layout = QHBoxLayout(central_widget)
|
||
|
||
# 左侧摄像头显示区域
|
||
left_frame = QFrame()
|
||
left_frame.setFrameStyle(QFrame.StyledPanel)
|
||
left_frame.setStyleSheet("QFrame { background-color: #f0f0f0; border: 2px solid #ddd; }")
|
||
left_layout = QVBoxLayout(left_frame)
|
||
|
||
# 摄像头显示标签
|
||
self.camera_label = QLabel()
|
||
self.camera_label.setMinimumSize(640, 480)
|
||
self.camera_label.setStyleSheet("QLabel { background-color: black; border: 1px solid #ccc; }")
|
||
self.camera_label.setAlignment(Qt.AlignCenter)
|
||
self.camera_label.setText("摄像头未启动")
|
||
self.camera_label.setScaledContents(True)
|
||
|
||
# 控制按钮
|
||
button_layout = QHBoxLayout()
|
||
self.start_button = QPushButton("启动摄像头")
|
||
self.stop_button = QPushButton("停止摄像头")
|
||
self.start_button.clicked.connect(self.start_camera)
|
||
self.stop_button.clicked.connect(self.stop_camera)
|
||
self.stop_button.setEnabled(False)
|
||
|
||
button_layout.addWidget(self.start_button)
|
||
button_layout.addWidget(self.stop_button)
|
||
button_layout.addStretch()
|
||
|
||
left_layout.addWidget(self.camera_label)
|
||
left_layout.addLayout(button_layout)
|
||
|
||
# 右侧结果显示区域
|
||
right_frame = QFrame()
|
||
right_frame.setFrameStyle(QFrame.StyledPanel)
|
||
right_frame.setFixedWidth(400)
|
||
right_frame.setStyleSheet("QFrame { background-color: #fafafa; border: 2px solid #ddd; }")
|
||
right_layout = QVBoxLayout(right_frame)
|
||
|
||
# 标题
|
||
title_label = QLabel("检测结果")
|
||
title_label.setAlignment(Qt.AlignCenter)
|
||
title_label.setFont(QFont("Arial", 16, QFont.Bold))
|
||
title_label.setStyleSheet("QLabel { color: #333; padding: 10px; }")
|
||
|
||
# 车牌数量显示
|
||
self.count_label = QLabel("识别到的车牌数量: 0")
|
||
self.count_label.setAlignment(Qt.AlignCenter)
|
||
self.count_label.setFont(QFont("Arial", 12))
|
||
self.count_label.setStyleSheet(
|
||
"QLabel { "
|
||
"background-color: #e3f2fd; "
|
||
"border: 1px solid #2196f3; "
|
||
"border-radius: 5px; "
|
||
"padding: 8px; "
|
||
"color: #1976d2; "
|
||
"font-weight: bold; "
|
||
"}"
|
||
)
|
||
|
||
# 滚动区域用于显示车牌结果
|
||
scroll_area = QScrollArea()
|
||
scroll_area.setWidgetResizable(True)
|
||
scroll_area.setStyleSheet("QScrollArea { border: none; background-color: transparent; }")
|
||
|
||
self.results_widget = QWidget()
|
||
self.results_layout = QVBoxLayout(self.results_widget)
|
||
self.results_layout.setAlignment(Qt.AlignTop)
|
||
|
||
scroll_area.setWidget(self.results_widget)
|
||
|
||
right_layout.addWidget(title_label)
|
||
right_layout.addWidget(self.count_label)
|
||
right_layout.addWidget(scroll_area)
|
||
|
||
# 添加到主布局
|
||
main_layout.addWidget(left_frame, 2)
|
||
main_layout.addWidget(right_frame, 1)
|
||
|
||
# 设置样式
|
||
self.setStyleSheet("""
|
||
QMainWindow {
|
||
background-color: #f5f5f5;
|
||
}
|
||
QPushButton {
|
||
background-color: #2196F3;
|
||
color: white;
|
||
border: none;
|
||
padding: 8px 16px;
|
||
border-radius: 4px;
|
||
font-weight: bold;
|
||
}
|
||
QPushButton:hover {
|
||
background-color: #1976D2;
|
||
}
|
||
QPushButton:pressed {
|
||
background-color: #0D47A1;
|
||
}
|
||
QPushButton:disabled {
|
||
background-color: #cccccc;
|
||
color: #666666;
|
||
}
|
||
""")
|
||
|
||
def init_detector(self):
|
||
"""初始化检测器"""
|
||
model_path = os.path.join(os.path.dirname(__file__), "yolopart", "yolo11s-pose42.pt")
|
||
self.detector = LicensePlateYOLO(model_path)
|
||
|
||
def init_camera(self):
|
||
"""初始化摄像头线程"""
|
||
self.camera_thread = CameraThread()
|
||
self.camera_thread.frame_ready.connect(self.process_frame)
|
||
|
||
def start_camera(self):
|
||
"""启动摄像头"""
|
||
if self.camera_thread.start_camera():
|
||
self.start_button.setEnabled(False)
|
||
self.stop_button.setEnabled(True)
|
||
self.camera_label.setText("摄像头启动中...")
|
||
else:
|
||
self.camera_label.setText("摄像头启动失败")
|
||
|
||
def stop_camera(self):
|
||
"""停止摄像头"""
|
||
self.camera_thread.stop_camera()
|
||
self.start_button.setEnabled(True)
|
||
self.stop_button.setEnabled(False)
|
||
self.camera_label.setText("摄像头已停止")
|
||
self.camera_label.clear()
|
||
|
||
def process_frame(self, frame):
|
||
"""处理摄像头帧"""
|
||
self.current_frame = frame.copy()
|
||
|
||
# 进行车牌检测
|
||
self.detections = self.detector.detect_license_plates(frame)
|
||
|
||
# 在图像上绘制检测结果
|
||
display_frame = self.draw_detections(frame.copy())
|
||
|
||
# 转换为Qt格式并显示
|
||
self.display_frame(display_frame)
|
||
|
||
# 更新右侧结果显示
|
||
self.update_results_display()
|
||
|
||
def draw_detections(self, frame):
|
||
"""在图像上绘制检测结果"""
|
||
return self.detector.draw_detections(frame, self.detections)
|
||
|
||
def display_frame(self, frame):
|
||
"""显示帧到界面"""
|
||
rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
||
h, w, ch = rgb_frame.shape
|
||
bytes_per_line = ch * w
|
||
qt_image = QImage(rgb_frame.data, w, h, bytes_per_line, QImage.Format_RGB888)
|
||
|
||
pixmap = QPixmap.fromImage(qt_image)
|
||
scaled_pixmap = pixmap.scaled(self.camera_label.size(), Qt.KeepAspectRatio, Qt.SmoothTransformation)
|
||
self.camera_label.setPixmap(scaled_pixmap)
|
||
|
||
def update_results_display(self):
|
||
"""更新右侧结果显示"""
|
||
# 更新车牌数量
|
||
count = len(self.detections)
|
||
self.count_label.setText(f"识别到的车牌数量: {count}")
|
||
|
||
# 清除之前的结果
|
||
for i in reversed(range(self.results_layout.count())):
|
||
child = self.results_layout.itemAt(i).widget()
|
||
if child:
|
||
child.setParent(None)
|
||
|
||
# 添加新的结果
|
||
for i, detection in enumerate(self.detections):
|
||
# 矫正车牌图像
|
||
corrected_image = self.correct_license_plate(detection)
|
||
|
||
# 获取车牌号，传入车牌类型信息
|
||
plate_number = self.recognize_plate_number(corrected_image, detection['class_name'])
|
||
|
||
# 创建车牌显示组件
|
||
plate_widget = LicensePlateWidget(
|
||
i + 1,
|
||
detection['class_name'],
|
||
corrected_image,
|
||
plate_number
|
||
)
|
||
|
||
self.results_layout.addWidget(plate_widget)
|
||
|
||
def correct_license_plate(self, detection):
|
||
"""矫正车牌图像"""
|
||
if self.current_frame is None:
|
||
return None
|
||
|
||
# 检查是否为不完整检测
|
||
if detection.get('incomplete', False):
|
||
return None
|
||
|
||
# 使用检测器的矫正方法
|
||
return self.detector.correct_license_plate(
|
||
self.current_frame,
|
||
detection['keypoints']
|
||
)
|
||
|
||
def recognize_plate_number(self, corrected_image, class_name):
|
||
"""识别车牌号"""
|
||
if corrected_image is None:
|
||
return "识别失败"
|
||
|
||
try:
|
||
# 预测函数（来自模块）
|
||
# 函数名改成一样的了，所以不要修改这里了，想用哪个模块直接导入
|
||
result = LPRNmodel_predict(corrected_image)
|
||
|
||
# 将字符列表转换为字符串，支持8位车牌号
|
||
if isinstance(result, list) and len(result) >= 7:
|
||
# 根据车牌类型决定显示位数
|
||
if class_name == '绿牌' and len(result) >= 8:
|
||
# 绿牌显示8位，过滤掉空字符占位符
|
||
plate_chars = [char for char in result[:8] if char != '']
|
||
# 如果过滤后确实有8位，显示8位；否则显示7位
|
||
if len(plate_chars) == 8:
|
||
return ''.join(plate_chars)
|
||
else:
|
||
return ''.join(plate_chars[:7])
|
||
else:
|
||
# 蓝牌或其他类型显示前7位，过滤掉空字符
|
||
plate_chars = [char for char in result[:7] if char != '']
|
||
return ''.join(plate_chars)
|
||
else:
|
||
return "识别失败"
|
||
except Exception as e:
|
||
print(f"车牌号识别失败: {e}")
|
||
return "识别失败"
|
||
|
||
def closeEvent(self, event):
|
||
"""窗口关闭事件"""
|
||
if self.camera_thread:
|
||
self.camera_thread.stop_camera()
|
||
event.accept()
|
||
|
||
def main():
|
||
app = QApplication(sys.argv)
|
||
window = MainWindow()
|
||
window.show()
|
||
sys.exit(app.exec_())
|
||
|
||
if __name__ == "__main__":
|
||
main() |