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

点击上方蓝字关注我们

微信公众号：OpenCV学堂

关注获取更多计算机视觉与深度学习知识

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

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()