如何實現(xiàn)YOLOv8 + ONNRUNTIME推理界面化與多線程支持

2023年一月份跟二月份創(chuàng)建了一個PyQT5人工智能軟件開發(fā)系列的文章系列，過去的兩個月都沒怎么更新，心里一直想有時間繼續(xù)更新下去，今天又更新了一篇，基于PyQT5實現(xiàn)多線程、界面化、YOLOv8對象檢測、實例分割、姿態(tài)評估的推理。

基本設(shè)計思路

這個系列我好久沒有更新了，今天更新一篇PyQT5中如何實現(xiàn)YOLOv8 + ONNRUNTIME推理界面化與多線程支持。首先需要實現(xiàn)三個類分別完成YOLOv8的對象檢測、實例分割、姿態(tài)評估模型推理。然后在實現(xiàn)界面類，構(gòu)建如圖：

推理類型通過單選按鈕實現(xiàn)選擇，支持對象檢測、實例分割、姿態(tài)評估。參數(shù)設(shè)置選擇模型文件與標簽文件地址作為輸入，同時選擇置信度，置信度之在0~1之間。 推理按鈕開發(fā)推理演示，支持視頻與圖像文件，開始推理會單獨開啟一個推理線程實現(xiàn)推理，推理結(jié)果通過信號發(fā)送到指定的槽函數(shù)處理之后更新界面，通過信號與槽機制實現(xiàn)界面線程跟推理線程相互獨立與數(shù)據(jù)共享。

界面代碼實現(xiàn)

界面部分通過一個QWidget實現(xiàn)一個面板，通過垂直與水平布局實現(xiàn)界面組件的布局管理，相關(guān)的代碼實現(xiàn)如下：

classYOLOv8InferPanel(QtWidgets.QWidget):
def__init__(self,parent=None):
super().__init__(parent)

#文本標簽
self.rbtn0=QtWidgets.QRadioButton("對象檢測")
self.rbtn1=QtWidgets.QRadioButton("實例分割")
self.rbtn3=QtWidgets.QRadioButton("姿態(tài)評估")
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("參數(shù)設(shè)置")
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)

#setuplistener
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繼承實現(xiàn)run方法，完成推理線程構(gòu)建，根據(jù)傳入的參數(shù)不同，初始化不同的推理類型（對象檢測、實例分割、姿態(tài)評估），推理線程實現(xiàn)代碼如下：

classInferenceThread(QtCore.QThread):
fire_stats_signal=QtCore.pyqtSignal(dict)

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

defrun(self):
ifself.detectorisNone:
return
ifself.input_image.endswith(".mp4"):
cap=cv.VideoCapture(self.input_image)
whileTrue:
ret,frame=cap.read()
ifretisTrue:
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

應(yīng)用程序演示

最終調(diào)用應(yīng)用程序代碼，實現(xiàn)啟動與運行的界面如下：

#初始化APP實例
importplatform
app=QtWidgets.QApplication(sys.argv)
if'Windows'==platform.system():
app.setStyle('Windows')
#初始化桌面容器
main_win=QtWidgets.QMainWindow()
#設(shè)置APP窗口名稱
main_win.setWindowTitle("YOLOv8多線程推理應(yīng)用演示-2號高手")
#初始化內(nèi)容面板
content_panel=YOLOv8InferPanel()
#設(shè)置窗口大小
main_win.setMinimumSize(1340,960)
main_win.setCentralWidget(content_panel)

#請求顯示
main_win.show()
#加載窗口并啟動App
app.exec()