使用深度学习目标检测算法YOLOv5进行物体检测 通过训练自己的垃圾数据集 建立实时垃圾分类系统
文章目录
- 使用深度学习目标检测算法YOLOv5进行物体检测 通过训练自己的垃圾数据集 建立实时垃圾分类系统
- 1. 数据准备
- 标注工具
- 2. 模型训练
- 安装YOLOv5
- 准备配置文件 来训练
- 训练模型
- 3. UI设计
- 安装PyQt5
- 创建主窗口
- 4. 实时检测
- 总结
以下文字及代码仅供参考
Python实时垃圾分类系统
1
构建一个实时垃圾分类系统,使用YOLOv5进行物体检测,并结合PyQt5来设计用户界面。以下是一个完整的示例代码,
从数据准备、模型训练、UI设计和实时检测,仅供参考。
1. 数据准备
首先,同学,你得有自己的垃圾数据集吧 。
需要一个包含不同类别垃圾的数据集。你可以从公开数据集中获取,或者自己采集并标注数据。对不对
标注工具
- 使用LabelImg等工具进行数据标注。
2. 模型训练
从官网自己下载就是了。
安装YOLOv5
gitclone https://github.com/ultralytics/yolov5.gitcdyolov5 pipinstall-rrequirements.txt准备配置文件 来训练
在yolov5目录下创建data.yaml:
train:./images/trainval:./images/valtest:./images/testnc:4# number of classesnames:['paper','plastic','glass','metal']训练模型
python train.py--img640--batch16--epochs100--datadata.yaml--cfgcfg/training/yolov5s.yaml--weightsyolov5s.pt--nameyolov5s_results3. UI设计
使用PyQt5设计用户界面。
安装PyQt5
pipinstallPyQt5创建主窗口
importsysfromPyQt5.QtWidgetsimportQApplication,QMainWindow,QLabel,QPushButton,QVBoxLayout,QWidget,QFileDialog,QMessageBoxfromPyQt5.QtGuiimportQPixmap,QImagefromPyQt5.QtCoreimportQTimerimportcv2importnumpyasnpfromultralyticsimportYOLOclassGarbageClassification(QMainWindow):def__init__(self):super().__init__()self.initUI()self.model=YOLO('best.pt')# 加载训练好的模型definitUI(self):self.setWindowTitle("垃圾分类系统")self.setGeometry(100,100,800,600)self.label=QLabel(self)self.label.setText("请选择图片或打开摄像头")self.label.setAlignment(Qt.AlignCenter)self.button_load=QPushButton("选择图片",self)self.button_load.clicked.connect(self.load_image)self.button_camera=QPushButton("打开摄像头",self)self.button_camera.clicked.connect(self.open_camera)self.button_close_camera=QPushButton("关闭摄像头",self)self.button_close_camera.clicked.connect(self.close_camera)layout=QVBoxLayout()layout.addWidget(self.label)layout.addWidget(self.button_load)layout.addWidget(self.button_camera)layout.addWidget(self.button_close_camera)container=QWidget()container.setLayout(layout)self.setCentralWidget(container)self.cap=Noneself.timer=Nonedefload_image(self):options=QFileDialog.Options()file_name,_=QFileDialog.getOpenFileName(self,"选择图片文件","","Image Files (*.jpg *.png)",options=options)iffile_name:self.image_path=file_name self.detect_garbage()defdetect_garbage(self):image=cv2.imread(self.image_path)results=self.model(image)annotated_frame=results[0].plot()height,width,channel=annotated_frame.shape bytes_per_line=3*width q_img=QImage(annotated_frame.data,width,height,bytes_per_line,QImage.Format_RGB888).rgbSwapped()self.label.setPixmap(QPixmap.fromImage(q_img))defopen_camera(self):self.cap=cv2.VideoCapture(0)self.timer=QTimer(self)self.timer.timeout.connect(self.update_frame)self.timer.start(30)# 每秒更新30帧defclose_camera(self):ifself.capisnotNone:self.timer.stop()self.cap.release()self.cap=Nonedefupdate_frame(self):ret,frame=self.cap.read()ifnotret:returnresults=self.model(frame)annotated_frame=results[0].plot()height,width,channel=annotated_frame.shape bytes_per_line=3*width q_img=QImage(annotated_frame.data,width,height,bytes_per_line,QImage.Format_RGB888).rgbSwapped()self.label.setPixmap(QPixmap.fromImage(q_img))if__name__=="__main__":app=QApplication(sys.argv)window=GarbageClassification()window.show()sys.exit(app.exec_())4. 实时检测
在上述代码中,同学 :实现了以下功能:
- 选择图片文件。
- 打开摄像头。
- 关闭摄像头。
- 实时显示检测结果。
总结
仅供参考。基于YOLOv5的实时垃圾分类系统。这个系统包括数据准备、模型训练、UI设计和实时检测等功能。