支持不同模型检测,实现类别统计
源码:
# from ultralytics import YOLO
# model=YOLO(r"C:\Users\Administrator\Desktop\mod\best.onnx",task='detect')
# results=model.predict(r"C:\Users\Administrator\Desktop\图片\huibanbings (4).JPG",save=True,task='detect')
### 作者@ :2744391598@qq.com 2023 !!
### 作者@ :2744391598@qq.com 2023 !!
### 作者@ :2744391598@qq.com 2023 !!
import sys
import cv2
import ultralytics
from PySide6.QtCore import Qt
from PySide6.QtWidgets import QApplication, QWidget, QVBoxLayout, QLabel, QFileDialog, QPushButton
from PySide6.QtGui import QImage, QPixmap
class FileDialogExample(QWidget):
def __init__(self):
super().__init__()
self.init_ui()
def init_ui(self):
self.setWindowTitle('图像检测系统—当前模型:未选择')
self.setGeometry(400, 100, 900, 730)# 设置窗体大小和位置
vbox = QVBoxLayout()
select_model_button = QPushButton('选择模型') # 添加选择模型按钮
select_model_button.clicked.connect(self.select_model) # 绑定选择模型按钮的点击事件
vbox.addWidget(select_model_button) # 将选择模型按钮添加到布局中
self.setLayout(vbox) # 设置窗体布局
self.image_button = QPushButton('选择检测图像', self)
self.image_button.clicked.connect(self.select_image)
vbox.addWidget(self.image_button)
self.video_button = QPushButton('选择检测视频', self)
self.video_button.clicked.connect(self.select_video)
vbox.addWidget(self.video_button)
self.cam_button=QPushButton('打开摄像头',self)
self.cam_button.clicked.connect(self.open_cam)
vbox.addWidget(self.cam_button)
self.stop_button=QPushButton('关闭摄像头',self)
self.stop_button.clicked.connect(self.stop_cam)
vbox.addWidget(self.stop_button)
# 添加一个 QLabel 用于显示预测结果
self.result_label = QLabel(self)
vbox.addWidget(self.result_label)
self.setLayout(vbox)
def select_model(self):
options = QFileDialog.Options()
options |= QFileDialog.DontUseNativeDialog
file_name, _ = QFileDialog.getOpenFileName(self, 'Select Model', '', 'Model Files (*.onnx *.pt)',
options=options)
if file_name is None or file_name == "":
print("model none")
return
else:
self.model_path = file_name
self.model = ultralytics.YOLO(self.model_path, task='detect')
self.setWindowTitle(f'万能图像检测系统-当前模型:{file_name}') # 更新当前模型的标签文本
def open_cam(self):
# 创建一个 VideoCapture 对象,用于调用摄像头
self.cap = cv2.VideoCapture(0)
# 检查摄像头是否被正确打开
if not self.cap.isOpened():
raise IOError("无法打开摄像头")
self.cam_button.setEnabled(False)
self.stop_button.setEnabled(True)
while self.cap.isOpened():
ret, frame = self.cap.read()
# 在窗口中显示帧
cv2.imshow('摄像头', frame)
results = self.model.predict(source=frame, save=False, task='detect',show=True)
# 处理 GUI 事件,以确保能够接收到 stop_button 的点击事件
QApplication.processEvents()
# 释放摄像头对象和销毁窗口
self.cap.release()
cv2.destroyAllWindows()
self.cam_button.setEnabled(True)
self.stop_button.setEnabled(False)
def stop_cam(self):
if self.cap:
self.cap.release()
cv2.destroyAllWindows()
self.cam_button.setEnabled(True)
self.stop_button.setEnabled(False)
# 设置停止按钮为不可用
self.stop_button.setEnabled(False)
def select_video(self):
options = QFileDialog.Options()
options |= QFileDialog.DontUseNativeDialog
file_name,_=QFileDialog.getOpenFileName(self,'Select Video','','Video (*.mp4)',options=options)
if file_name is None:
print("video none")
return
else:
print("请耐心等待")
self.videp_path=file_name
results = self.model.predict(source=self.videp_path, save=True, task='detect',show=False)
# print("检测完成,结果已保存:runs\detect...")
def select_image(self):
options = QFileDialog.Options()
options |= QFileDialog.DontUseNativeDialog
file_name, _ = QFileDialog.getOpenFileName(self, 'Select Image', '', 'Images (*.jpg *.png)', options=options)
if file_name is None:
print("image none")
return
else:
self.image_path = file_name
img = cv2.imread(self.image_path)
results = self.model.predict(self.image_path, save=True, task='detect',save_txt=True,show=False)
res = results[0].plot()
# 在结果图像上绘制物体数量信息
font = cv2.FONT_HERSHEY_SIMPLEX
font_scale = 1
font_color = (255, 255, 255)
thickness = 2
cv2.putText(res, f"NUM{len(results[0])}", (50, 50), font, font_scale, font_color, thickness)
# 将图像转换为 QPixmap
height, width, channel = res.shape
bytes_per_line = 3 * width
q_img = QImage(res.data, width, height, bytes_per_line, QImage.Format_RGB888)
pixmap = QPixmap.fromImage(q_img)
# 调整图像尺寸以适应 QLabel
scaled_pixmap = pixmap.scaled(self.result_label.size(), Qt.KeepAspectRatio, Qt.SmoothTransformation)
# 更新 QLabel 显示的图像
self.result_label.setPixmap(scaled_pixmap)
if __name__ == '__main__':
app = QApplication(sys.argv)
window = FileDialogExample()
window.show()
sys.exit(app.exec())
#yolo detect predict model=D:\1.27PycharmProjects\PycharmProjects\ultralytics-main\jg_detect.pt source='screen' device=0 save=True show=False