yolorestart
This commit is contained in:
@@ -1 +0,0 @@
|
||||
# UI模块初始化文件
|
||||
@@ -1,348 +0,0 @@
|
||||
#!/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()
|
||||
@@ -1,59 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
视频显示组件
|
||||
用于显示视频帧和检测结果
|
||||
"""
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
from PyQt5.QtWidgets import QLabel
|
||||
from PyQt5.QtCore import Qt
|
||||
from PyQt5.QtGui import QImage, QPixmap, QPainter, QPen, QFont
|
||||
|
||||
class VideoWidget(QLabel):
|
||||
"""视频显示组件"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.setMinimumSize(640, 480)
|
||||
self.setStyleSheet("border: 1px solid gray; background-color: black;")
|
||||
self.setAlignment(Qt.AlignCenter)
|
||||
self.setText("视频显示区域\n点击'开始检测'开始")
|
||||
self.setScaledContents(True)
|
||||
|
||||
def update_frame(self, frame):
|
||||
"""更新显示帧"""
|
||||
if frame is None:
|
||||
return
|
||||
|
||||
# 转换BGR到RGB
|
||||
rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
||||
h, w, ch = rgb_frame.shape
|
||||
bytes_per_line = ch * w
|
||||
|
||||
# 创建QImage
|
||||
qt_image = QImage(rgb_frame.data, w, h, bytes_per_line, QImage.Format_RGB888)
|
||||
|
||||
# 转换为QPixmap并显示
|
||||
pixmap = QPixmap.fromImage(qt_image)
|
||||
|
||||
# 缩放以适应widget大小，保持宽高比
|
||||
scaled_pixmap = pixmap.scaled(
|
||||
self.size(),
|
||||
Qt.KeepAspectRatio,
|
||||
Qt.SmoothTransformation
|
||||
)
|
||||
|
||||
self.setPixmap(scaled_pixmap)
|
||||
|
||||
def paintEvent(self, event):
|
||||
"""绘制事件"""
|
||||
super().paintEvent(event)
|
||||
|
||||
# 如果没有图像，显示提示文本
|
||||
if not self.pixmap():
|
||||
painter = QPainter(self)
|
||||
painter.setPen(QPen(Qt.white))
|
||||
painter.setFont(QFont("Arial", 16))
|
||||
painter.drawText(self.rect(), Qt.AlignCenter, "视频显示区域\n点击'开始检测'开始")
|
||||
Reference in New Issue
Block a user