多线程界面化、ONNXRUNTIME+YOLOv8推理演示

OpenCV学堂 2023-05-21 22:24

TI MCU方案：电动汽车实时控制 构建AI未来，Arm计算平台无处不在

点击上方蓝字关注我们

微信公众号：OpenCV学堂
关注获取更多计算机视觉与深度学习知识

2023年一月份跟二月份创建了一个PyQT5人工智能软件开发系列的文章系列，过去的两个月都没怎么更新，心里一直想有时间继续更新下去，今天又更新了一篇，基于PyQT5实现多线程、界面化、YOLOv8对象检测、实例分割、姿态评估的推理。

基本设计思路

这个系列我好久没有更新了，今天更新一篇PyQT5中如何实现YOLOv8 + ONNRUNTIME推理界面化与多线程支持。首先需要实现三个类分别完成YOLOv8的对象检测、实例分割、姿态评估模型推理。然后在实现界面类，构建如图：

推理类型通过单选按钮实现选择，支持对象检测、实例分割、姿态评估。参数设置选择模型文件与标签文件地址作为输入，同时选择置信度，置信度之在0~1之间。

推理按钮开发推理演示，支持视频与图像文件，开始推理会单独开启一个推理线程实现推理，推理结果通过信号发送到指定的槽函数处理之后更新界面，通过信号与槽机制实现界面线程跟推理线程相互独立与数据共享。

界面代码实现

界面部分通过一个QWidget实现一个面板，通过垂直与水平布局实现界面组件的布局管理，相关的代码实现如下：

class YOLOv8InferPanel(QtWidgets.QWidget):
    def __init__(self, parent=None):
        super().__init__(parent)

        # 文本标签
        self.rbtn0 = QtWidgets.QRadioButton("对象检测")
        self.rbtn1 = QtWidgets.QRadioButton("实例分割")
        self.rbtn3 = QtWidgets.QRadioButton("姿态评估")
        self.rbtn0.setChecked(True)

        hbox_layout1 = QtWidgets.QHBoxLayout()
        hbox_layout1.addWidget(self.rbtn0)
        hbox_layout1.addWidget(self.rbtn1)
        hbox_layout1.addWidget(self.rbtn3)

        panel3 = QtWidgets.QGroupBox("推理类型")
        panel3.setLayout(hbox_layout1)

        # 输入文本框
        self.image_file_edit = QtWidgets.QLineEdit()
        self.image_file_edit.setMinimumWidth(100)
        self.image_file_edit.setEnabled(False)
        fileBtn = QtWidgets.QPushButton("图像文件")
        self.weight_file_path = QtWidgets.QLineEdit()
        self.weight_file_path.setMinimumWidth(100)
        self.weight_file_path.setEnabled(False)
        modelBtn = QtWidgets.QPushButton("模型文件")

        self.label_file_path = QtWidgets.QLineEdit()
        self.label_file_path.setMinimumWidth(100)
        self.label_file_path.setEnabled(False)
        labelBtn = QtWidgets.QPushButton("标签文件")

        self.conf_spinbox = QtWidgets.QDoubleSpinBox()
        self.conf_spinbox.setRange(0, 1)
        self.conf_spinbox.setSingleStep(0.01)
        self.conf_spinbox.setValue(0.25)

        grid_layout2 = QtWidgets.QGridLayout()
        grid_layout2.addWidget(fileBtn, 0, 0)
        grid_layout2.addWidget(self.image_file_edit, 0, 1)
        grid_layout2.addWidget(modelBtn, 0, 2)
        grid_layout2.addWidget(self.weight_file_path, 0, 3)

        grid_layout2.addWidget(labelBtn, 1, 0)
        grid_layout2.addWidget(self.label_file_path, 1, 1)
        grid_layout2.addWidget(QtWidgets.QLabel("置信:"), 1, 2)
        grid_layout2.addWidget(self.conf_spinbox, 1, 3)

        panel2 = QtWidgets.QGroupBox("参数设置")
        panel2.setLayout(grid_layout2)

        # 输入文本框
        self.label = QtWidgets.QLabel()
        self.label.setMinimumSize(1280, 720)
        pixmap = QtGui.QPixmap("images/wp.jpg")
        pix = pixmap.scaled(QtCore.QSize(1280, 720), QtCore.Qt.KeepAspectRatio)
        self.label.setPixmap(pix)
        self.label.setAlignment(QtCore.Qt.AlignCenter)
        self.label.setStyleSheet("background-color:black; color: green")

        self.startBtn = QtWidgets.QPushButton("开始推理")
        self.stopBtn = QtWidgets.QPushButton("停止")
        self.startBtn.setStyleSheet("background-color:cyan; color: black")
        self.stopBtn.setStyleSheet("background-color:gray; color: white")
        self.stopBtn.setEnabled(False)
        hbox_layout = QtWidgets.QHBoxLayout()
        hbox_layout.addStretch(1)
        hbox_layout.addWidget(self.startBtn)
        hbox_layout.addWidget(self.stopBtn)
        panel1 = QtWidgets.QWidget()
        panel1.setLayout(hbox_layout)

        # 添加到布局管理器中
        vbox_layout = QtWidgets.QVBoxLayout()
        vbox_layout.addWidget(panel3)
        vbox_layout.addWidget(panel2)
        vbox_layout.addWidget(panel1)
        vbox_layout.addWidget(self.label)
        vbox_layout.addStretch(1)

        # 面板容器
        self.setLayout(vbox_layout)

        # setup listener
        modelBtn.clicked.connect(self.on_weight_select)
        fileBtn.clicked.connect(self.on_update_image)
        labelBtn.clicked.connect(self.on_label_select)
        self.startBtn.clicked.connect(self.on_yolov8_infer)
        self.work_thread = None

推理线程

基于QThread继承实现run方法，完成推理线程构建，根据传入的参数不同，初始化不同的推理类型（对象检测、实例分割、姿态评估），推理线程实现代码如下：

class InferenceThread(QtCore.QThread):
    fire_stats_signal = QtCore.pyqtSignal(dict)

    def __init__(self, settings):
        super(InferenceThread, self).__init__()
        self.settings = settings
        self.detector = None
        if self.settings.model_type == 0:
            self.detector = YOLOv8ORTDetector(settings)
        if self.settings.model_type == 1:
            self.detector = YOLOv8ORTSegment(settings)
        if self.settings.model_type == 2:
            self.detector = YOLOv8ORTPose(settings)
        self.input_image = settings.input_image

    def run(self):
        if self.detector is None:
            return
        if self.input_image.endswith(".mp4"):
            cap = cv.VideoCapture(self.input_image)
            while True:
                ret, frame = cap.read()
                if ret is True:
                    self.detector.infer_image(frame)
                    self.fire_stats_signal.emit({"result": frame})
                else:
                    break
        else:
            frame = cv.imread(self.input_image)
            self.detector.infer_image(frame)
            self.fire_stats_signal.emit({"result": frame})
        self.fire_stats_signal.emit({"done": "done"})
        return

应用程序演示

最终调用应用程序代码，实现启动与运行的界面如下：

# 初始化APP实例
import platform
app = QtWidgets.QApplication(sys.argv)
if 'Windows' == platform.system():
    app.setStyle('Windows')
# 初始化桌面容器
main_win = QtWidgets.QMainWindow()
# 设置APP窗口名称
main_win.setWindowTitle("YOLOv8多线程推理应用演示-2号高手")
# 初始化内容面板
content_panel = YOLOv8InferPanel()
# 设置窗口大小
main_win.setMinimumSize(1340, 960)
main_win.setCentralWidget(content_panel)

# 请求显示
main_win.show()
# 加载窗口并启动App
app.exec()