Initial commit

yolopart/ui/__init__.py

@@ -0,0 +1 @@
+# UI模块初始化文件

yolopart/ui/main_window.py

@@ -0,0 +1,348 @@
+#!/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()

yolopart/ui/video_widget.py

@@ -0,0 +1,59 @@
+#!/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点击'开始检测'开始")