该项目的重点是利用计算机视觉技术,使用在Raspberry Pi 4B上实现的YOLOv5s对象检测模型实时检测和分类塑料垃圾。塑料废物检测将废物分为 PET、HDPE、PP 和不可回收物。
- 树莓派 4B :树莓派 4B 是我们项目的硬件平台。其紧凑的尺寸、低功耗和 GPIO(通用输入/输出)功能使其成为边缘计算和物联网应用的理想选择。
- 实时检测:我们的实施可以对网络摄像头捕获的塑料垃圾进行实时检测和分类。
存储库结构
- 数据集准备
- 模型训练
- 检波
数据集准备
“在 Raspberry Pi 4B 上使用 YOLOv5s 进行塑料垃圾检测”项目利用了由 6000 张图像组成的自定义数据集。这些图像被捕获以涵盖环境中常见的各种类型的塑料废物,包括塑料瓶、塑料袋、塑料杯、电缆、肥皂瓶和聚苯乙烯泡沫塑料。
为了有效地训练YOLOv5s模型,使用Roboflow平台对数据集进行了注释。Roboflow 提供用户友好的界面和注释工具,可简化注释过程,节省时间和精力。
注释包括边界框,用于指定图像中每个塑料废物对象的位置和大小。这些注释对于训练模型以准确检测和分类塑料废物至关重要。
模型训练
“在 Raspberry Pi 4B 上使用 YOLOv5s 进行塑料废物检测”模型已使用使用 PyTorch 框架实现的 YOLOv5s 架构进行训练。YOLOv5s 的选择是基于模型精度和大小之间的平衡,使其适合部署在 Raspberry Pi 等资源受限的设备上。
数据集比例
用于训练模型的数据集分为三个主要子集:训练、验证和测试。这些子集的比例如下:
- 训练:70% 的数据集用于训练 YOLOv5s 模型。该子集用于教导模型准确检测和分类塑料废物对象
- 验证:预留 20% 的数据集用于验证目的。该子集有助于评估模型在训练过程中的性能,从而实现微调和参数优化。
- 测试:数据集的其余 10% 专门用于测试训练模型的性能。该子集提供了对模型泛化到看不见的塑料废物场景的能力的独立评估。
培训流程
要使用指定的参数训练“在 Raspberry Pi 4B 上使用 YOLOv5s 进行塑料废物检测”模型,您可以按照以下步骤操作:
- 加载必要的框架和库:
%pip install -qr requirements.txt # install dependencies
%pip install -q roboflow
import torch
import os
from IPython.display import Image, clear_output # to display images
print(f"Setup complete. Using torch {torch.__version__} ({torch.cuda.get_device_properties(0).name if torch.cuda.is_available() else 'CPU'})")
2. Assamble 数据集
# set up environment
os.environ["DATASET_DIRECTORY"] = "/content/datasets"
from roboflow import Roboflow
rf = Roboflow(api_key="Your Roboflow API KEY")
project = rf.workspace("workspace").project("project name")
dataset = project.version(1).download("yolov5")
3.训练模型
训练过程包括将带注释的数据集提供给 YOLOv5s 模型并优化其参数以最大程度地减少检测错误。此过程遍历多个 epoch,每个 epoch 代表整个训练数据集的完整传递。
在训练过程中,该模型通过调整其内部重量来学习识别各种塑料废物。使用的损失函数指导模型最小化预测边界框与地面实况注释之间的差异。
!python train.py --img 416 --batch 16 --epochs 200 --data {dataset.location}/data.yaml --weights yolov5s.pt --cache
4.模型评估
为了评估训练模型的性能,同时使用了验证和测试子集。验证子集有助于在训练期间监控模型的进度,而测试子集则提供对其在看不见的数据上的性能的公正评估。
可以使用各种评估指标来衡量模型的准确性,例如精确度、召回率和平均精度 (mAP)。这些指标有助于量化模型正确检测和分类塑料废物的能力。
%load_ext tensorboard
%tensorboard --logdir runs
评估完成后,您可以在 /runs/train/exp 中下载名为 best.pt