1. 项目概述
水面漂浮垃圾识别系统是一个基于Python深度学习的计算机视觉项目,旨在通过图像识别技术自动检测和分类水面上的各类垃圾。这个系统对于环境保护和水域治理具有重要应用价值,能够帮助相关部门快速定位污染源,提高清洁效率。
我在开发这个系统时,主要采用了YOLOv5目标检测框架作为核心算法,配合OpenCV进行图像预处理,最终实现了一个准确率超过85%的实用模型。整个项目从数据采集到模型部署耗时约3个月,期间遇到了不少技术挑战,也积累了一些宝贵的实战经验。
2. 技术选型与架构设计
2.1 核心算法选择
在算法选型阶段,我对比了Faster R-CNN、SSD和YOLO系列等多个目标检测模型。最终选择YOLOv5主要基于以下几个考虑:
- 检测速度:YOLO系列以实时检测著称,v5版本在保持精度的同时进一步优化了速度
- 模型大小:YOLOv5提供了从n(最小)到x(最大)多个预训练模型,可根据硬件条件灵活选择
- 易用性:PyTorch框架下的实现,社区支持完善,便于调试和二次开发
- 精度表现:在COCO数据集上,YOLOv5s的mAP@0.5达到56.8%,满足项目需求
2.2 系统架构设计
整个系统采用模块化设计,主要分为以下几个组件:
code复制数据采集模块 → 数据预处理模块 → 模型训练模块 → 推理检测模块 → 结果可视化模块
每个模块都设计了清晰的接口,方便单独测试和替换。例如数据预处理模块支持多种图像增强方式,可以通过配置文件灵活切换。
3. 数据集构建与处理
3.1 数据采集策略
构建高质量的数据集是项目成功的关键。我采用了多源数据采集方案:
- 公开数据集:使用已有的Water垃圾数据集作为基础(约2000张)
- 网络爬取:从环保网站爬取相关图片(约500张)
- 实地拍摄:在本地河流、湖泊拍摄真实场景(约300张)
- 数据增强:通过旋转、裁剪、色彩变换等方式扩充至10000+张
注意:实地拍摄时要注意安全,最好两人同行,并选择天气良好的白天进行
3.2 数据标注规范
使用LabelImg工具进行标注时,我们制定了严格的标注规范:
- 垃圾完全在水面上的才标注
- 被遮挡超过30%的不标注
- 每个标注框尽量贴近物体边缘
- 分类包括:塑料瓶、塑料袋、泡沫、木材、其他
标注文件采用YOLO格式保存,每个图像对应一个.txt文件,内容示例:
code复制0 0.45 0.32 0.12 0.08 # 类别 x_center y_center width height
1 0.67 0.21 0.15 0.10
4. 模型训练与优化
4.1 训练环境配置
硬件配置:
- GPU: NVIDIA RTX 3090 (24GB显存)
- CPU: AMD Ryzen 9 5950X
- 内存: 64GB DDR4
软件环境:
- Python 3.8
- PyTorch 1.9.0
- CUDA 11.1
- cuDNN 8.0.5
安装依赖:
bash复制git clone https://github.com/ultralytics/yolov5
cd yolov5
pip install -r requirements.txt
4.2 训练参数调优
经过多次实验,最终采用的训练参数配置:
yaml复制# yolov5s.yaml
nc: 5 # 类别数
depth_multiple: 0.33
width_multiple: 0.50
# hyp.scratch.yaml
lr0: 0.01
lrf: 0.2
momentum: 0.937
weight_decay: 0.0005
warmup_epochs: 3.0
warmup_momentum: 0.8
warmup_bias_lr: 0.1
启动训练命令:
bash复制python train.py --img 640 --batch 16 --epochs 100 --data water.yaml --cfg yolov5s.yaml --weights yolov5s.pt
4.3 训练过程监控
使用TensorBoard监控训练过程:
bash复制tensorboard --logdir runs/train
重点关注以下指标:
- 损失函数变化:box_loss, obj_loss, cls_loss
- 精度指标:mAP@0.5, mAP@0.5:0.95
- 学习率变化:确保正常衰减
5. 模型部署与应用
5.1 模型导出与优化
训练完成后,将模型导出为ONNX格式以便跨平台部署:
bash复制python export.py --weights runs/train/exp/weights/best.pt --include onnx
使用TensorRT进一步优化:
bash复制trtexec --onnx=best.onnx --saveEngine=best.engine --fp16
5.2 实时检测实现
基于OpenCV的实时检测代码框架:
python复制import cv2
import torch
model = torch.hub.load('ultralytics/yolov5', 'custom', path='best.pt')
cap = cv2.VideoCapture(0) # 0为默认摄像头
while True:
ret, frame = cap.read()
if not ret:
break
results = model(frame)
rendered = results.render()[0]
cv2.imshow('Water Garbage Detection', rendered)
if cv2.waitKey(1) == ord('q'):
break
cap.release()
cv2.destroyAllWindows()
5.3 性能优化技巧
- 多线程处理:使用生产者-消费者模式分离图像采集和推理
- 批处理:积累多帧后批量推理,提高GPU利用率
- 量化压缩:采用FP16或INT8量化减小模型体积
- 硬件加速:利用TensorRT或OpenVINO优化推理速度
6. 常见问题与解决方案
6.1 模型识别精度低
可能原因:
- 数据量不足或质量差
- 类别不平衡
- 标注不准确
解决方案:
- 增加数据增强方式(如MixUp, Mosaic)
- 使用类别权重调整损失函数
- 重新检查并修正标注错误
6.2 推理速度慢
优化方法:
- 降低输入分辨率(如从640x640降至416x416)
- 使用更小的模型版本(YOLOv5n)
- 启用TensorRT加速
6.3 误检率高
改进措施:
- 增加负样本(无垃圾的水面图像)
- 调整置信度阈值
- 添加后处理过滤规则(如大小、长宽比)
7. 项目扩展方向
在实际应用中,可以考虑以下几个扩展方向:
- 多模态检测:结合红外或声呐数据提高检测鲁棒性
- 垃圾追踪:实现跨帧目标跟踪,统计垃圾运动轨迹
- 垃圾量估算:基于检测结果估算垃圾体积和重量
- 移动端部署:优化模型在边缘设备(如无人机)上的运行
我在项目中尝试了无人机搭载方案,使用Jetson Xavier NX作为计算单元,实现了空中巡检功能。这个过程中发现的主要挑战是阳光反射造成的误检,通过增加偏振滤镜得到了一定改善。