yolorestart
This commit is contained in:
411
main.py
Normal file
411
main.py
Normal file
@@ -0,0 +1,411 @@
|
||||
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 OCR_part.ocr_interface import ocr_predict
|
||||
#from CRNN_part.crnn_interface import crnn_predict（不使用CRNN）
|
||||
|
||||
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()
|
||||
|
||||
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)
|
||||
|
||||
# 创建车牌显示组件
|
||||
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):
|
||||
"""识别车牌号"""
|
||||
if corrected_image is None:
|
||||
return "识别失败"
|
||||
|
||||
try:
|
||||
# 使用OCR接口进行识别
|
||||
# 可以根据需要切换为CRNN: crnn_predict(corrected_image)
|
||||
result = ocr_predict(corrected_image)
|
||||
|
||||
# 将字符列表转换为字符串
|
||||
if isinstance(result, list) and len(result) >= 7:
|
||||
return ''.join(result[:7])
|
||||
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()
|
||||
Reference in New Issue
Block a user