1.确保Ultralytics源码无误且环境正常配置
环境异常请先搜索环境配置,这里不展开讲
建议使用Anaconda3+pycharm
2.界面实现
2.1 相关库
pip install PyQt5
2.2 导入包
我这里导入的是以下这些,同官方源码兼容
import sys
from PyQt5.QtCore import Qt
from PyQt5.QtWidgets import QApplication, QMainWindow, QLabel, QVBoxLayout, QWidget, QPushButton, QHBoxLayout, QMessageBox, QFileDialog
from PyQt5.QtGui import QImage, QPixmap
import cv2
from ultralytics import YOLO
2.3 功能实现
这里只实现了简单的权重选择和图片检测功能并添加了对应的按钮
# 添加模型选择按钮
self.load_model_button = QPushButton("模型选择")
self.load_model_button.clicked.connect(self.load_model)
hbox_buttons.addWidget(self.load_model_button)
# 添加图片检测按钮
self.image_detect_button = QPushButton("图片检测")
self.image_detect_button.clicked.connect(self.detect_image)
hbox_buttons.addWidget(self.image_detect_button)
# 添加退出按钮
self.exit_button = QPushButton("退出")
self.exit_button.clicked.connect(self.exit_application)
hbox_buttons.addWidget(self.exit_button)
#加载模型文件
def load_model(self):
model_path, _ = QFileDialog.getOpenFileName(None, "选择模型文件", "", "模型文件 (*.pt)")
if model_path:
self.model = YOLO(model_path)
if self.model:
return True
else:
return False
#加载待检测图片文件
def detect_image(self):
image_path, _ = QFileDialog.getOpenFileName(None, "选择图片文件", "", "图片文件 (*.jpg *.jpeg *.png)")
if image_path:
results = self.worker.detect_image(image_path)
if results:
annotated_image = results[0].plot()
height, width, channel = annotated_image.shape
bytesPerLine = 3 * width
qimage = QImage(annotated_image.data, width, height, bytesPerLine, QImage.Format_BGR888)
pixmap = QPixmap.fromImage(qimage)
self.label.setPixmap(pixmap.scaled(self.label.size(), Qt.KeepAspectRatio))
else:
QMessageBox.critical(self, "错误", "请选择图片文件")
def exit_application(self):
# 终止程序运行
sys.exit()
模型加载前图片是不允许被检测的,所以把按钮先禁用,等模型加载后再启用
self.image_detect_button.setEnabled(False)
def load_model(self):
if self.worker.load_model():
self.image_detect_button.setEnabled(True)
2.4 界面展示
这样一个简单的ui界面就写好了
2.5 完整代码
import sys
from PyQt5.QtCore import Qt
from PyQt5.QtWidgets import QApplication, QMainWindow, QLabel, QVBoxLayout, QWidget, QPushButton, QHBoxLayout, QMessageBox, QFileDialog
from PyQt5.QtGui import QImage, QPixmap
import cv2
from ultralytics import YOLO
class Worker:
def __init__(self):
self.model = None
def load_model(self):
model_path, _ = QFileDialog.getOpenFileName(None, "选择模型文件", "", "模型文件 (*.pt)")
if model_path:
self.model = YOLO(model_path)
if self.model:
return True
else:
return False
def detect_image(self, image_path):
image = cv2.imread(image_path)
if image is not None:
results = self.model.predict(image)
return results
else:
return None
class MainWindow(QMainWindow):
def __init__(self):
super().__init__()
self.setWindowTitle("Numb753")
self.setGeometry(100, 100, 800, 600)
self.label = QLabel()
self.label.setAlignment(Qt.AlignCenter)
layout = QVBoxLayout()
layout.addWidget(self.label)
central_widget = QWidget()
central_widget.setLayout(layout)
self.setCentralWidget(central_widget)
self.worker = Worker()
# 创建按钮布局
hbox_buttons = QHBoxLayout()
layout.addLayout(hbox_buttons)
# 添加模型选择按钮
self.load_model_button = QPushButton("模型选择")
self.load_model_button.clicked.connect(self.load_model)
hbox_buttons.addWidget(self.load_model_button)
# 添加图片检测按钮
self.image_detect_button = QPushButton("图片检测")
self.image_detect_button.clicked.connect(self.detect_image)
self.image_detect_button.setEnabled(False)
hbox_buttons.addWidget(self.image_detect_button)
# 添加退出按钮
self.exit_button = QPushButton("退出")
self.exit_button.clicked.connect(self.exit_application)
hbox_buttons.addWidget(self.exit_button)
def detect_image(self):
image_path, _ = QFileDialog.getOpenFileName(None, "选择图片文件", "", "图片文件 (*.jpg *.jpeg *.png)")
if image_path:
results = self.worker.detect_image(image_path)
if results:
annotated_image = results[0].plot()
height, width, channel = annotated_image.shape
bytesPerLine = 3 * width
qimage = QImage(annotated_image.data, width, height, bytesPerLine, QImage.Format_BGR888)
pixmap = QPixmap.fromImage(qimage)
self.label.setPixmap(pixmap.scaled(self.label.size(), Qt.KeepAspectRatio))
else:
QMessageBox.critical(self, "错误", "请选择图片文件")
def load_model(self):
if self.worker.load_model():
self.image_detect_button.setEnabled(True)
def exit_application(self):
# 终止程序运行
sys.exit()
if __name__ == '__main__':
app = QApplication(sys.argv)
window = MainWindow()
window.show()
sys.exit(app.exec_())
完整代码附上,其他需求可以私聊,遇到报错可以评论交流~