348 lines
14 KiB
Python
348 lines
14 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
主界面窗口
|
||
包含视频显示区域、控制按钮和车牌号显示区域
|
||
"""
|
||
|
||
import sys
|
||
import os
|
||
from PyQt5.QtWidgets import (
|
||
QMainWindow, QWidget, QVBoxLayout, QHBoxLayout,
|
||
QLabel, QPushButton, QFrame, QTextEdit, QGroupBox,
|
||
QCheckBox, QSpinBox, QSlider, QGridLayout
|
||
)
|
||
from PyQt5.QtCore import Qt, QTimer, pyqtSignal
|
||
from PyQt5.QtGui import QFont, QPixmap, QPalette, QImage
|
||
|
||
from .video_widget import VideoWidget
|
||
from utils.video_capture import VideoCapture
|
||
from models.yolo_detector import YOLODetector
|
||
from models.plate_recognizer import PlateRecognizerManager
|
||
|
||
class MainWindow(QMainWindow):
|
||
"""主窗口类"""
|
||
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.video_capture = None
|
||
self.yolo_detector = None
|
||
self.plate_recognizer = PlateRecognizerManager('mock') # 车牌识别管理器
|
||
self.timer = QTimer()
|
||
self.use_camera = 1 # 1: 摄像头, 0: 视频文件
|
||
self.detected_plates = [] # 存储切割后的车牌图像数组
|
||
self.current_frame = None # 存储当前帧用于车牌矫正
|
||
|
||
self.init_ui()
|
||
self.init_detector()
|
||
self.init_video_capture()
|
||
self.connect_signals()
|
||
|
||
def init_ui(self):
|
||
"""初始化用户界面"""
|
||
self.setWindowTitle("车牌检测系统 - YOLO11s")
|
||
self.setGeometry(100, 100, 1200, 800)
|
||
|
||
# 创建中央widget
|
||
central_widget = QWidget()
|
||
self.setCentralWidget(central_widget)
|
||
|
||
# 主布局
|
||
main_layout = QHBoxLayout(central_widget)
|
||
|
||
# 左侧视频显示区域
|
||
self.create_video_area(main_layout)
|
||
|
||
# 右侧控制和信息显示区域
|
||
self.create_control_area(main_layout)
|
||
|
||
# 设置布局比例
|
||
main_layout.setStretch(0, 3) # 视频区域占3/4
|
||
main_layout.setStretch(1, 1) # 控制区域占1/4
|
||
|
||
def create_video_area(self, parent_layout):
|
||
"""创建视频显示区域"""
|
||
video_frame = QFrame()
|
||
video_frame.setFrameStyle(QFrame.StyledPanel)
|
||
video_layout = QVBoxLayout(video_frame)
|
||
|
||
# 视频显示widget
|
||
self.video_widget = VideoWidget()
|
||
video_layout.addWidget(self.video_widget)
|
||
|
||
parent_layout.addWidget(video_frame)
|
||
|
||
def create_control_area(self, parent_layout):
|
||
"""创建控制和信息显示区域"""
|
||
control_frame = QFrame()
|
||
control_frame.setFrameStyle(QFrame.StyledPanel)
|
||
control_frame.setMaximumWidth(300)
|
||
control_layout = QVBoxLayout(control_frame)
|
||
|
||
# 控制按钮组
|
||
self.create_control_buttons(control_layout)
|
||
|
||
# 检测信息显示
|
||
self.create_detection_info(control_layout)
|
||
|
||
# 车牌号显示区域
|
||
self.create_plate_display(control_layout)
|
||
|
||
# 系统状态显示
|
||
self.create_status_display(control_layout)
|
||
|
||
parent_layout.addWidget(control_frame)
|
||
|
||
def create_control_buttons(self, parent_layout):
|
||
"""创建控制按钮"""
|
||
button_group = QGroupBox("控制面板")
|
||
button_layout = QVBoxLayout(button_group)
|
||
|
||
# 开始/停止按钮
|
||
self.start_btn = QPushButton("开始检测")
|
||
self.start_btn.setMinimumHeight(40)
|
||
self.start_btn.clicked.connect(self.toggle_detection)
|
||
button_layout.addWidget(self.start_btn)
|
||
|
||
# 视频源切换
|
||
self.camera_checkbox = QCheckBox("使用摄像头")
|
||
self.camera_checkbox.setChecked(True)
|
||
self.camera_checkbox.stateChanged.connect(self.toggle_video_source)
|
||
button_layout.addWidget(self.camera_checkbox)
|
||
|
||
# 检测开关
|
||
self.detection_checkbox = QCheckBox("启用检测")
|
||
self.detection_checkbox.setChecked(True)
|
||
button_layout.addWidget(self.detection_checkbox)
|
||
|
||
parent_layout.addWidget(button_group)
|
||
|
||
def create_detection_info(self, parent_layout):
|
||
"""创建检测信息显示"""
|
||
info_group = QGroupBox("检测信息")
|
||
info_layout = QVBoxLayout(info_group)
|
||
|
||
# FPS显示
|
||
self.fps_label = QLabel("FPS: 0")
|
||
self.fps_label.setFont(QFont("Arial", 12, QFont.Bold))
|
||
info_layout.addWidget(self.fps_label)
|
||
|
||
# 检测数量
|
||
self.detection_count_label = QLabel("检测到车牌: 0")
|
||
info_layout.addWidget(self.detection_count_label)
|
||
|
||
# 模型信息
|
||
self.model_info_label = QLabel("模型: YOLO11s (ONNX)")
|
||
info_layout.addWidget(self.model_info_label)
|
||
|
||
parent_layout.addWidget(info_group)
|
||
|
||
def create_plate_display(self, parent_layout):
|
||
"""创建车牌号显示区域"""
|
||
plate_group = QGroupBox("车牌识别结果")
|
||
plate_layout = QVBoxLayout(plate_group)
|
||
|
||
# 当前识别的车牌号
|
||
self.current_plate_label = QLabel("当前车牌: 未识别")
|
||
self.current_plate_label.setFont(QFont("Arial", 14, QFont.Bold))
|
||
self.current_plate_label.setStyleSheet("color: blue; padding: 10px; border: 1px solid gray;")
|
||
plate_layout.addWidget(self.current_plate_label)
|
||
|
||
# 矫正后的车牌图像显示
|
||
self.plate_image_label = QLabel("矫正后车牌图像")
|
||
self.plate_image_label.setAlignment(Qt.AlignCenter)
|
||
self.plate_image_label.setMinimumHeight(100)
|
||
self.plate_image_label.setMaximumHeight(150)
|
||
self.plate_image_label.setStyleSheet("border: 1px solid gray; background-color: #f0f0f0;")
|
||
plate_layout.addWidget(self.plate_image_label)
|
||
|
||
# 历史车牌记录
|
||
history_label = QLabel("历史记录:")
|
||
plate_layout.addWidget(history_label)
|
||
|
||
self.plate_history = QTextEdit()
|
||
self.plate_history.setMaximumHeight(150)
|
||
self.plate_history.setReadOnly(True)
|
||
plate_layout.addWidget(self.plate_history)
|
||
|
||
# 预留接口说明
|
||
interface_label = QLabel("注: 车牌识别接口已预留，可接入OCR模型")
|
||
interface_label.setStyleSheet("color: gray; font-size: 10px;")
|
||
plate_layout.addWidget(interface_label)
|
||
|
||
parent_layout.addWidget(plate_group)
|
||
|
||
def create_status_display(self, parent_layout):
|
||
"""创建系统状态显示"""
|
||
status_group = QGroupBox("系统状态")
|
||
status_layout = QVBoxLayout(status_group)
|
||
|
||
self.status_label = QLabel("状态: 就绪")
|
||
status_layout.addWidget(self.status_label)
|
||
|
||
self.gpu_status_label = QLabel("GPU: 检测中...")
|
||
status_layout.addWidget(self.gpu_status_label)
|
||
|
||
parent_layout.addWidget(status_group)
|
||
|
||
# 添加弹性空间
|
||
parent_layout.addStretch()
|
||
|
||
def init_detector(self):
|
||
"""初始化YOLO检测器"""
|
||
try:
|
||
model_path = os.path.join(os.path.dirname(__file__), "..", "yolo11sth50.onnx")
|
||
self.yolo_detector = YOLODetector(model_path)
|
||
self.model_info_label.setText(f"模型: YOLO11s (ONNX) - GPU: {self.yolo_detector.use_gpu}")
|
||
self.gpu_status_label.setText(f"GPU: {'启用' if self.yolo_detector.use_gpu else '禁用'}")
|
||
except Exception as e:
|
||
self.status_label.setText(f"模型加载失败: {str(e)}")
|
||
|
||
def init_video_capture(self):
|
||
"""初始化视频捕获"""
|
||
try:
|
||
self.video_capture = VideoCapture()
|
||
self.status_label.setText("视频捕获初始化成功")
|
||
except Exception as e:
|
||
self.status_label.setText(f"视频捕获初始化失败: {str(e)}")
|
||
|
||
def connect_signals(self):
|
||
"""连接信号和槽"""
|
||
self.timer.timeout.connect(self.update_frame)
|
||
|
||
def toggle_detection(self):
|
||
"""切换检测状态"""
|
||
if self.timer.isActive():
|
||
self.stop_detection()
|
||
else:
|
||
self.start_detection()
|
||
|
||
def start_detection(self):
|
||
"""开始检测"""
|
||
if self.video_capture and self.video_capture.start_capture(self.use_camera):
|
||
# 根据视频源类型设置定时器间隔
|
||
video_fps = self.video_capture.get_video_fps()
|
||
timer_interval = int(1000 / video_fps) # 转换为毫秒
|
||
self.timer.start(timer_interval)
|
||
|
||
self.start_btn.setText("停止检测")
|
||
source_type = "摄像头" if self.use_camera else f"视频文件({video_fps:.1f}FPS)"
|
||
self.status_label.setText(f"检测中... - {source_type}")
|
||
else:
|
||
self.status_label.setText("启动失败")
|
||
|
||
def stop_detection(self):
|
||
"""停止检测"""
|
||
self.timer.stop()
|
||
if self.video_capture:
|
||
self.video_capture.stop_capture()
|
||
self.start_btn.setText("开始检测")
|
||
self.status_label.setText("已停止")
|
||
|
||
def toggle_video_source(self, state):
|
||
"""切换视频源"""
|
||
self.use_camera = 1 if state == Qt.Checked else 0
|
||
if self.timer.isActive():
|
||
self.stop_detection()
|
||
self.start_detection()
|
||
|
||
def update_frame(self):
|
||
"""更新帧"""
|
||
if not self.video_capture:
|
||
return
|
||
|
||
frame, fps = self.video_capture.get_frame()
|
||
if frame is None:
|
||
return
|
||
|
||
# 保存当前帧用于车牌矫正
|
||
self.current_frame = frame.copy()
|
||
|
||
# 更新FPS显示
|
||
self.fps_label.setText(f"FPS: {fps:.1f}")
|
||
|
||
# 进行检测
|
||
if self.detection_checkbox.isChecked() and self.yolo_detector:
|
||
detections = self.yolo_detector.detect(frame)
|
||
frame = self.yolo_detector.draw_detections(frame, detections)
|
||
|
||
# 切割车牌图像
|
||
if detections:
|
||
self.detected_plates = self.yolo_detector.crop_plates(frame, detections)
|
||
|
||
# 统计不同类型车牌数量
|
||
blue_count = sum(1 for d in detections if d['class_id'] == 0)
|
||
green_count = sum(1 for d in detections if d['class_id'] == 1)
|
||
total_count = len(detections)
|
||
|
||
self.detection_count_label.setText(f"检测到车牌: {total_count} (蓝牌:{blue_count}, 绿牌:{green_count})")
|
||
|
||
# 调用车牌识别接口（预留）
|
||
self.recognize_plates(self.detected_plates, detections)
|
||
else:
|
||
self.detection_count_label.setText("检测到车牌: 0")
|
||
|
||
# 显示帧
|
||
self.video_widget.update_frame(frame)
|
||
|
||
def recognize_plates(self, plate_images, detections):
|
||
"""车牌识别接口（预留）"""
|
||
# 这里是预留的车牌识别接口
|
||
# 可以接入OCR模型进行车牌号识别
|
||
if plate_images and detections and self.current_frame is not None:
|
||
# 获取最新检测到的车牌信息
|
||
latest_detection = detections[-1] # 取最后一个检测结果
|
||
plate_type = "Blue Plate" if latest_detection['class_id'] == 0 else "Green Plate"
|
||
confidence = latest_detection['confidence']
|
||
|
||
# 处理蓝色车牌的矫正
|
||
corrected_image = None
|
||
if latest_detection['class_id'] == 0: # 蓝色车牌
|
||
try:
|
||
bbox = latest_detection['bbox']
|
||
corrected_image = self.plate_recognizer.preprocess_blue_plate(
|
||
plate_images[-1], self.current_frame, bbox
|
||
)
|
||
self._display_plate_image(corrected_image)
|
||
except Exception as e:
|
||
print(f"蓝色车牌矫正失败: {e}")
|
||
self.plate_image_label.setText("蓝色车牌矫正失败")
|
||
elif latest_detection['class_id'] == 1: # 绿色车牌
|
||
# 绿色车牌处理预留
|
||
self.plate_image_label.setText("绿色车牌处理\n(待实现)")
|
||
|
||
# 模拟识别结果
|
||
plate_text = f"Mock {plate_type}-{len(plate_images)}"
|
||
self.current_plate_label.setText(f"Current Plate: {plate_text} (Confidence: {confidence:.2f})")
|
||
|
||
# 添加到历史记录
|
||
import datetime
|
||
timestamp = datetime.datetime.now().strftime("%H:%M:%S")
|
||
self.plate_history.append(f"[{timestamp}] {plate_text} (Confidence: {confidence:.2f})")
|
||
|
||
def _display_plate_image(self, image):
|
||
"""在界面上显示车牌图像"""
|
||
try:
|
||
# 将OpenCV图像转换为QPixmap
|
||
if len(image.shape) == 3:
|
||
height, width, channel = image.shape
|
||
bytes_per_line = 3 * width
|
||
q_image = QImage(image.data, width, height, bytes_per_line, QImage.Format_RGB888).rgbSwapped()
|
||
else:
|
||
height, width = image.shape
|
||
bytes_per_line = width
|
||
q_image = QImage(image.data, width, height, bytes_per_line, QImage.Format_Grayscale8)
|
||
|
||
# 缩放图像以适应标签大小
|
||
pixmap = QPixmap.fromImage(q_image)
|
||
scaled_pixmap = pixmap.scaled(self.plate_image_label.size(), Qt.KeepAspectRatio, Qt.SmoothTransformation)
|
||
|
||
self.plate_image_label.setPixmap(scaled_pixmap)
|
||
except Exception as e:
|
||
print(f"显示车牌图像失败: {e}")
|
||
self.plate_image_label.setText(f"图像显示失败: {str(e)}")
|
||
|
||
def closeEvent(self, event):
|
||
"""窗口关闭事件"""
|
||
self.stop_detection()
|
||
event.accept() |