PyQt5展示YOLOV8检测效果

最新推荐文章于 2025-05-04 20:59:10 发布
原创最新推荐文章于 2025-05-04 20:59:10 发布
· 263 阅读
8 ·
版权
文章标签：
#YOLO
结果示意：
代码：
import os
import sys
import time
import cv2
import numpy as np
from PyQt5 import QtCore, QtGui, QtWidgets
from PyQt5.QtCore import pyqtSignal, QThread, Qt
from PyQt5.QtGui import QImage, QPixmap
from PyQt5.QtWidgets import (QApplication, QMainWindow, QFileDialog,
                             QTableWidgetItem, QMessageBox)
from ultralytics import YOLO


class DetThread(QThread):
    send_imgraw = pyqtSignal(np.ndarray)
    send_imgresult = pyqtSignal(np.ndarray)
    send_result = pyqtSignal(object, dict)  # 修改信号类型

    def __init__(self):
        super().__init__()
        self.running = False
        self.source = None
        self.conf = 0.25
        self.iou = 0.45
        self.model = None
        self.cap = None
        self.mode = 'camera'

    def run(self):
        self.running = True
        try:
            if self.mode == 'image':
                self.process_image()
            else:
                self.process_video()
        except Exception as e:
            print(f"检测错误: {e}")
        finally:
            self.cleanup()

    def process_image(self):
        img = cv2.imread(self.source)
        if img is not None:
            start_time = time.time()
            results = self.model.predict(img, conf=self.conf, iou=self.iou)
            det_info = self._collect_det_info(results, img.shape, start_time)

            self.send_imgraw.emit(img)
            self.send_imgresult.emit(results[0].plot())
            self.send_result.emit(results[0].boxes, det_info)

    def process_video(self):
        self.cap = cv2.VideoCapture(self.source if self.mode == 'video' else 0)
        while self.running and self.cap.isOpened():
            ret, frame = self.cap.read()
            if not ret:
                break

            start_time = time.time()
            results = self.model.predict(frame, conf=self.conf, iou=self.iou)
            det_info = self._collect_det_info(results, frame.shape, start_time)

            self.send_imgraw.emit(frame)
            self.send_imgresult.emit(results[0].plot())
            self.send_result.emit(results[0].boxes, det_info)

            time.sleep(0.03)

    def _collect_det_info(self, results, shape, start_time):
        return {
            'shape': f"{shape[1]}x{shape[0]}",
            'objects': len(results[0]),
            'class_names': results[0].names,
            'speed': results[0].speed,
            'total_time': (time.time() - start_time) * 1000
        }

    def cleanup(self):
        if self.cap and self.cap.isOpened():
            self.cap.release()
        self.running = False

    def stop(self):
        self.running = False


class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setup_ui()
        self.init_variables()
        self.connect_signals()
    def setup_ui(self):
        self.setWindowTitle("检测系统")
        self.resize(1200, 800)
        self.setStyleSheet("background-color: #FFFAFA;")

        # 主布局组件
        self.central_widget = QtWidgets.QWidget(self)
        self.setCentralWidget(self.central_widget)
        self.main_layout = QtWidgets.QHBoxLayout(self.central_widget)

        # 左侧控制面板
        self.left_panel = QtWidgets.QWidget()
        self.left_panel.setFixedWidth(400)
        self.left_layout = QtWidgets.QVBoxLayout(self.left_panel)

        # 数据源设置
        self.setup_data_source()
        # 模型参数设置
        self.setup_model_params()
        # 结果显示
        self.setup_results()

        self.main_layout.addWidget(self.left_panel)

        # 右侧显示区域
        self.right_panel = QtWidgets.QWidget()
        self.right_layout = QtWidgets.QVBoxLayout(self.right_panel)

        self.img_label_raw = QtWidgets.QLabel("原始图像")
        self.img_label_raw.setStyleSheet("background: #C0C0C0;")
        self.img_label_raw.setAlignment(Qt.AlignCenter)
        self.img_label_raw.setMinimumSize(640, 360)

        self.img_label_result = QtWidgets.QLabel("检测结果")
        self.img_label_result.setStyleSheet("background: #C0C0C0;")
        self.img_label_result.setAlignment(Qt.AlignCenter)
        self.img_label_result.setMinimumSize(640, 360)

        self.right_layout.addWidget(self.img_label_raw)
        self.right_layout.addWidget(self.img_label_result)
        self.main_layout.addWidget(self.right_panel)

    def setup_data_source(self):
        group = QtWidgets.QGroupBox("数据来源设置")
        layout = QtWidgets.QVBoxLayout()

        # 相机设置
        self.cam_radio_on = QtWidgets.QRadioButton("开启相机")
        self.cam_radio_off = QtWidgets.QRadioButton("关闭相机", checked=True)
        radio_layout = QtWidgets.QHBoxLayout()
        radio_layout.addWidget(self.cam_radio_on)
        radio_layout.addWidget(self.cam_radio_off)
        layout.addLayout(radio_layout)

        # 文件选择
        self.btn_file = QtWidgets.QPushButton("选择图片/视频")
        self.file_path = QtWidgets.QLineEdit()
        self.file_path.setReadOnly(True)
        file_layout = QtWidgets.QHBoxLayout()
        file_layout.addWidget(self.btn_file)
        file_layout.addWidget(self.file_path)
        layout.addLayout(file_layout)

        # 参数设置
        param_layout = QtWidgets.QFormLayout()
        self.pixel_format = QtWidgets.QLineEdit("RGB")
        self.resolution = QtWidgets.QLineEdit("640x640")
        self.exposure = QtWidgets.QLineEdit("1000")
        param_layout.addRow("像素格式:", self.pixel_format)
        param_layout.addRow("标准尺寸:", self.resolution)
        param_layout.addRow("曝光时间:", self.exposure)
        layout.addLayout(param_layout)

        group.setLayout(layout)
        self.left_layout.addWidget(group)

    def setup_model_params(self):
        group = QtWidgets.QGroupBox("模型参数设置")
        layout = QtWidgets.QVBoxLayout()

        # 模型选择
        self.btn_model = QtWidgets.QPushButton("选择模型")
        self.model_path = QtWidgets.QLineEdit()
        self.model_path.setReadOnly(True)
        model_layout = QtWidgets.QHBoxLayout()
        model_layout.addWidget(self.btn_model)
        model_layout.addWidget(self.model_path)
        layout.addLayout(model_layout)

        # 设备选择
        self.cpu_radio = QtWidgets.QRadioButton("CPU")
        self.gpu_radio = QtWidgets.QRadioButton("GPU", checked=True)
        device_layout = QtWidgets.QHBoxLayout()
        device_layout.addWidget(self.cpu_radio)
        device_layout.addWidget(self.gpu_radio)
        layout.addLayout(device_layout)

        # 置信度设置
        self.conf_slider = QtWidgets.QSlider(Qt.Horizontal)
        self.conf_slider.setRange(0, 100)
        self.conf_spin = QtWidgets.QDoubleSpinBox()
        self.conf_spin.setRange(0.0, 1.0)
        conf_layout = QtWidgets.QHBoxLayout()
        conf_layout.addWidget(QtWidgets.QLabel("置信度:"))
        conf_layout.addWidget(self.conf_spin)
        conf_layout.addWidget(self.conf_slider)
        layout.addLayout(conf_layout)

        # IoU阈值设置
        self.iou_slider = QtWidgets.QSlider(Qt.Horizontal)
        self.iou_slider.setRange(0, 100)
        self.iou_spin = QtWidgets.QDoubleSpinBox()
        self.iou_spin.setRange(0.0, 1.0)
        iou_layout = QtWidgets.QHBoxLayout()
        iou_layout.addWidget(QtWidgets.QLabel("IoU阈值:"))
        iou_layout.addWidget(self.iou_spin)
        iou_layout.addWidget(self.iou_slider)
        layout.addLayout(iou_layout)

        group.setLayout(layout)
        self.left_layout.addWidget(group)


    def setup_results(self):
        group = QtWidgets.QGroupBox("检测结果")
        layout = QtWidgets.QVBoxLayout()

        self.result_table = QtWidgets.QTableWidget()
        self.result_table.setColumnCount(4)  # 修改为4列
        self.result_table.setHorizontalHeaderLabels(["类别", "中心坐标", "尺寸", "置信度"])  # 修改表头
        self.result_table.verticalHeader().setVisible(False)
        self.result_table.setEditTriggers(QtWidgets.QAbstractItemView.NoEditTriggers)
        layout.addWidget(self.result_table)

        # 控制按钮
        self.btn_start = QtWidgets.QPushButton("开始检测")
        self.btn_stop = QtWidgets.QPushButton("停止检测")
        btn_layout = QtWidgets.QHBoxLayout()
        btn_layout.addWidget(self.btn_start)
        btn_layout.addWidget(self.btn_stop)
        layout.addLayout(btn_layout)

        group.setLayout(layout)
        self.left_layout.addWidget(group)

    def init_variables(self):
        self.det_thread = None
        self.current_mode = None
        self.model = None

    def connect_signals(self):
         # 按钮信号
        self.btn_file.clicked.connect(self.select_file)
        self.btn_model.clicked.connect(self.select_model)
        self.btn_start.clicked.connect(self.start_detection)
        self.btn_stop.clicked.connect(self.stop_detection)

            # 参数联动
        self.conf_slider.valueChanged.connect(
            lambda v: self.conf_spin.setValue(v / 100))
        self.conf_spin.valueChanged.connect(
            lambda v: self.conf_slider.setValue(int(v * 100)))
        self.iou_slider.valueChanged.connect(
            lambda v: self.iou_spin.setValue(v / 100))
        self.iou_spin.valueChanged.connect(
            lambda v: self.iou_slider.setValue(int(v * 100)))

        # 相机模式切换
        self.cam_radio_on.toggled.connect(self.update_camera_mode)

    def select_file(self):
        file, _ = QFileDialog.getOpenFileName(
            self, "选择文件", "",
             "图片文件 (*.jpg *.png);;视频文件 (*.mp4 *.avi)")
        if file:
            self.file_path.setText(file)
            if file.lower().endswith(('.jpg', '.png')):
                self.current_mode = 'image'
            else:
                self.current_mode = 'video'

    def select_model(self):
        file, _ = QFileDialog.getOpenFileName(
            self, "选择模型", "", "模型文件 (*.pt)")
        if file:
            try:
                self.model = YOLO(file)
                self.model_path.setText(os.path.basename(file))
                QMessageBox.information(self, "成功", "模型加载成功！")
            except Exception as e:
                QMessageBox.critical(self, "错误", f"模型加载失败: {str(e)}")

    def update_camera_mode(self, checked):
        self.current_mode = 'camera' if checked else None

    def start_detection(self):
        if not self.model:
            QMessageBox.warning(self, "警告", "请先选择模型文件！")
            return

        if self.det_thread and self.det_thread.isRunning():
            self.det_thread.stop()

        self.det_thread = DetThread()
        self.det_thread.model = self.model
        self.det_thread.conf = self.conf_spin.value()
        self.det_thread.iou = self.iou_spin.value()

        if self.current_mode == 'camera':
            self.det_thread.mode = 'camera'
            self.det_thread.source = 0
        elif self.current_mode == 'image':
            self.det_thread.mode = 'image'
            self.det_thread.source = self.file_path.text()
        elif self.current_mode == 'video':
            self.det_thread.mode = 'video'
            self.det_thread.source = self.file_path.text()
        else:
            QMessageBox.warning(self, "警告", "请先选择数据源！")
            return

        # 连接信号
        self.det_thread.send_imgraw.connect(self.show_raw_image)
        self.det_thread.send_imgresult.connect(self.show_result_image)
        self.det_thread.send_result.connect(self.update_results)

        self.det_thread.start()

    def stop_detection(self):
        if self.det_thread and self.det_thread.isRunning():
            self.det_thread.stop()
            self.det_thread.quit()
            self.det_thread.wait()

    def show_raw_image(self, img):
        self.display_image(img, self.img_label_raw)

    def show_result_image(self, img):
        self.display_image(img, self.img_label_result)

    def display_image(self, img, label):
        try:
            h, w, ch = img.shape
            bytes_per_line = ch * w
            q_img = QImage(img.data, w, h, bytes_per_line,
                            QImage.Format_RGB888).rgbSwapped()
            pixmap = QPixmap.fromImage(q_img)
            label.setPixmap(pixmap.scaled(
                label.size(), Qt.KeepAspectRatio, Qt.SmoothTransformation))
        except Exception as e:
            print(f"显示图像错误: {e}")



    def update_results(self, boxes, det_info):
        try:
            # 更新状态栏信息
            speed_info = (
                f"Speed: {det_info['speed']['preprocess']:.1f}ms preprocess, "
                f"{det_info['speed']['inference']:.1f}ms inference, "
                f"{det_info['speed']['postprocess']:.1f}ms postprocess"
            )
            main_info = (
                f"{det_info['objects']} objects detected | "
                f"Image size: {det_info['shape']} | "
                f"Total: {det_info['total_time']:.1f}ms"
            )
            self.statusBar().showMessage(f"{main_info} | {speed_info}")

            # 更新表格数据
            self.result_table.setRowCount(det_info['objects'])
            for row, box in enumerate(boxes):
                class_id = int(box.cls[0])
                class_name = det_info['class_names'].get(class_id, "unknown")
                xywh = box.xywh[0].tolist()
                conf = box.conf[0].item()

                self.result_table.setItem(row, 0, QTableWidgetItem(class_name))
                self.result_table.setItem(row, 1, QTableWidgetItem(
                    f"({xywh[0]:.1f}, {xywh[1]:.1f})"))
                self.result_table.setItem(row, 2, QTableWidgetItem(
                    f"{xywh[2]:.1f}x{xywh[3]:.1f}"))
                self.result_table.setItem(row, 3, QTableWidgetItem(f"{conf:.2f}"))

        except Exception as e:
            print(f"更新结果错误: {e}")

    def closeEvent(self, event):
        self.stop_detection()
        super().closeEvent(event)


if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec_())