简介:YOLOv5水下垃圾检测完整方案,面向计算机视觉学习者和海洋环保领域开发者,解决水下场景目标识别与分类难题。资源包含训练好的模型权重及PR曲线、loss曲线等过程文件,可直接推理使用;配套VOC格式水下垃圾数据集,内含数千张经LabelImg标注的真实海洋场景图片,覆盖金属、木材、塑料、橡胶、布料等类别,并同时提供VOC与YOLO两种标签格式,便于接入不同训练框架。压缩包共2000个文件,以txt标签文件为主,另有3个Python脚本和3个PDF环境配置/使用说明,整体约278.55MB,结构清晰易于检索。已有339人学习下载。PyQt可视化界面可实时呈现检测效果,适合快速搭建演示系统。结合CSDN配套博文可完整复现训练与测试流程,是入门水下目标检测、开展算法对比或完成毕业设计的实用资料。 水下垃圾检测这两年越来越受关注,不只是环保话题,也是很多毕设、竞赛和工程项目的热门方向。我做的这套系统用的是YOLOv5做检测核心,配套一份标注好的水下垃圾数据集,外加一个基于PyQt5的桌面可视化界面,可以加载训练好的权重对图片、视频或摄像头画面实时推理。整套东西跑下来,从数据集标注到模型训练再到界面封装,链路完整,很适合拿来当毕设框架或者进一步做工程化改造的起点。
这篇文章我把整个项目拆开讲一遍,包括数据怎么标、模型怎么训、界面怎么写,以及我踩过的坑和排查思路。关键是所有环节都有可复现的细节,不是只讲概念。
1. 项目概述与整体思路
1.1 为什么选YOLOv5做水下垃圾检测
水下环境相比于陆地场景有几个很麻烦的地方:光照不均匀、水体浑浊导致对比度低、垃圾种类多且形态差异大,比如塑料袋、玻璃瓶、金属罐、渔网等。在这种背景下做目标检测,模型的实时性和鲁棒性都是硬指标。
YOLOv5在工业界和学术界都很成熟,它在速度和精度之间拿捏得比较平衡。实际测下来,用GFLOPs更小的YOLOv5s也能在GPU上跑到100+FPS,即使是CPU也能有可用的推理速度。而相比YOLOv8,YOLOv5的部署生态更简单,改起来也方便,尤其要做PyQt桌面端的时候,模型导出和推理接口都很顺。另外,网上关于YOLOv5的教程和调参经验非常多,遇到问题随便搜都能找到解决方案,这对做毕设或者短期项目特别友好。
1.2 项目整体架构与功能拆解
整个项目的技术链路分为三大块:
- 数据集层:包括原始水下图像采集、垃圾类别标注、数据增强、数据集划分。
- 模型层:基于YOLOv5进行训练、验证、测试,导出最佳权重。
- 应用层:PyQt5构建可视化界面,实现本地模型加载、图像/视频/摄像头检测、结果实时展示与导出。
从功能上看,界面需要支持三种输入源:静态图片、视频文件和USB摄像头。用户加载训练好的.pt权重文件后,点击检测按钮即可看到带边界框和类别标签的输出。界面还应该支持调整置信度阈值和IoU阈值,方便在不同场景下切换使用。
做这套系统有一个很核心的思路:先把模型跑通,再做界面。不要在前期纠结界面漂不漂亮,先保证检测结果可靠,界面只是壳。我自己早期反着来,界面写了三天,模型没训好,最后调试时全在刷报错,效率极低。
2. 数据集准备:从零构建水下垃圾数据集
2.1 数据集来源与标注规范
水下垃圾数据集的来源主要有三个渠道:公开数据集、自行拍摄、网络爬取。如果要快速跑通,建议先用公开数据集,比如TrashCan、DeepTrash、SUIM等,但要注意这些数据集的标注格式和类别定义不完全一致。我的做法是以公开数据集为基础,再补充一部分自采数据,统一转成YOLOv5需要的YOLO txt格式。
YOLO格式的标注文件是每个图像对应一个同名txt,每一行内容是:class x_center y_center width height,其中坐标是归一化后的值,范围0到1。无论你用什么标注工具,最后都要转成这个格式。
我用的标注工具是LabelImg,它可以直接输出Pascal VOC格式的XML,再用脚本转成YOLO txt。LabelImg的操作很简单:框选目标、选择类别、保存。但项目里有个容易忽略的点,所有图像的标注框不能越界。明明目标在图像边缘,手工标的时候框可能会超出图像边界,这个在训练时会报错或者导致损失异常,需要写脚本把超出的部分裁剪回边界内。
类别方面我最终定的是6类:plastic_bag、plastic_bottle、glass_bottle、can、net、other。类别的确定不是随便搞的,而是根据实际场景中出现的频率和形态差异来的。如果你只有几类垃圾,类别定义太细会导致每个类样本太少,训练效果差;太粗又会把不同外观的物体混在一起,模型难以收敛。这个需要根据你自己的数据分布来权衡。
2.2 数据增强与数据集划分
水下图像普遍存在蓝色或绿色色偏、亮度不均、模糊等问题。为了让模型泛化能力更强,我的做法是在训练前做在线数据增强。YOLOv5内置了Mosaic、MixUp、HSV扰动、随机翻转等方法,我只需要在hyp.scratch.yaml里打开对应参数就行。
建议开启的增强项包括:
- Mosaic增强:把4张图拼成1张,有助于检测小目标和遮挡目标。
- HSV扰动:随机调整色调、饱和度、亮度,模拟不同水质下的颜色变化。
- 随机水平翻转:增加样本多样性。
- 轻微旋转和缩放:但角度不要超过±10度,否则边界框会明显变形。
数据集划分上,我按train : val : test = 8 : 1 : 1的比例切分。切分时要注意保证同一场景下的图像不要同时出现在训练集和验证集,不然模型在验证集上的指标会虚高,看起来80%的mAP,实际换一组视频就掉到50%。这个坑我踩过,当时为了省事直接随机切,后来发现验证集里很多图和训练集是连续帧,等于变相“背题”了。
3. YOLOv5模型训练与优化
3.1 环境配置与关键参数
YOLOv5的训练环境其实不复杂,我这里以最常用的PyTorch路线为例:
- Python 3.8+
- PyTorch 1.8+,CUDA 10.2以上即可
- 如果是NVIDIA显卡,装好CUDA和cuDNN,注意PyTorch版本和CUDA版本要对应
- 克隆YOLOv5仓库,执行
pip install -r requirements.txt
训练前需要准备好data.yaml文件,内容是指定训练验证路径和类别名。这个文件格式很简单,但路径一定要写对,最好是绝对路径,避免相对路径在不同的工作目录下报错。
关键训练参数我做了几个对比,最终没有用默认参数,因为默认参数很多时候不是最优的。
| 参数 | 默认值 | 我的配置 | 说明 |
|---|---|---|---|
img-size | 640 | 640 | 水下小目标多,不盲目降低输入尺寸 |
batch-size | 16 | 32 | 显存允许的情况下尽量大,能稳定训练 |
epochs | 300 | 200 | 看模型收敛情况,我200轮已经收敛 |
patience | 100 | 50 | 验证集mAP停止提升早停轮数 |
hyp | 默认 | 调低lr | 初始学习率0.01降到0.005,防止loss震荡 |
训练命令我放到下面,注意把路径换成你自己的:
python train.py --img 640 --batch 32 --epochs 200 --data data/trash.yaml --weights yolov5s.pt --cache --device 0--cache参数会把图像提前缓存到内存,训练速度会有明显提升。如果你的内存不够大,建议去掉这个参数,否则内存可能直接被吃满。
3.2 训练结果分析与模型调优
训练结束后,在runs/train/exp目录下会生成results.csv、权重文件和混淆矩阵等可视化结果。看训练过程主要关注两条曲线:mAP@0.5和val_loss。
我这次训练了200轮,大约在120轮时mAP@0.5已经到0.78,之后提升很慢,到190轮达到0.82。val_loss在80轮左右就已经趋于平缓,后面基本在一个小范围内抖动,说明模型已经收敛。
刚开始跑的时候我遇到一个典型问题:训练集loss下降但验证集loss不降反升,也就是过拟合。当时数据集总量还不到2000张,模型是从零开始训练的,分类数又多。解决办法有两个方向:
- 使用预训练权重:把
--weights指定成yolov5s.pt,用COCO预训练模型做初始化,迁移学习后收敛更快,泛化也更稳。 - 降低模型复杂度:从YOLOv5s换成YOLOv5n,参数量少一半,适合小数据集。
如果检测结果里出现很多重复框或者误检,可以调节后处理参数。也就是推理时的conf_thres和iou_thres。我建议默认conf_thres=0.25可以改成0.3,iou_thres保持0.45。在PyQt界面里我把这两个参数做成滑块,方便用户直接调节。
4. PyQt可视化界面开发
4.1 界面功能设计与布局
PyQt5做桌面应用已经很成熟了,它可以和OpenCV的读取图像方式无缝配合。我的界面布局分三个区域:左侧是功能控制面板,中间是图像显示区域,底部是日志输出区。
功能面板上包含:
- 模型加载按钮:选择本地的
.pt权重文件。 - 文件选择按钮:支持打开图片或视频。
- 摄像头开关:调用USB摄像头开始实时检测。
- 置信度阈值滑块:动态调整检测框的最低置信度。
- 检测/停止按钮:控制单次检测或连续视频检测。
界面代码里最关键的一段是加载模型和推理。这里我建议直接使用YOLOv5仓库里自带的detect.py逻辑,但不要直接调命令行,而是用Python API,这样能和PyQt方便地交互。
import torch from pathlib import Path # 加载YOLOv5模型 model = torch.hub.load('ultralytics/yolov5', 'custom', path='best.pt', force_reload=False) # 推理单张图片 results = model(img, size=640) # 获取渲染后的图像(带框) rendered_img = results.render()[0]这段代码是我在实际项目中用的精简版,注意torch.hub.load第一次会联网拉取YOLOv5仓库,之后设置了force_reload=False就不会每次都拉。如果你希望完全离线,也可以直接把yolov5仓库克隆到本地,用torch.hub.load('./yolov5', 'custom', path='best.pt', source='local')。
4.2 模型加载与推理逻辑
PyQt界面里,模型加载和推理不能在主线程里做。否则处理视频帧时界面会卡住,鼠标都动不了。正确的思路是:主线程跑UI,用QThread开一个工作线程做模型的加载和推理,通过信号把结果传回主线程刷新画面。
我踩过的一个真实教训是,最开始偷懒直接在while循环里调model(image),虽然帧率也不低,但只要模型推理耗时超过100ms,界面就会出现明显的掉帧和卡顿。后来改成QThread后,流畅度提升很明显。
视频帧读取用OpenCV的cv2.VideoCapture,每一帧传递给YOLOv5推理。这里要注意OpenCV读取的是BGR格式,而YOLOv5的推理接口内部会做颜色转换,所以你直接传BGR的numpy数组就行,不要自己手动转,否则颜色会异常。
摄像头实时检测还有一个容易忽略的地方:摄像头的帧率可能比模型推理速度高。如果你不控制帧率,CPU会一直被占用。我的做法是加一个时间判断,每两帧或者每隔0.05秒处理一次,避免系统资源被吃满。
5. 常见问题与排查技巧
5.1 训练阶段典型问题
训练的时候最容易出现以下几类报错和异常:
显存不足(CUDA out of memory)
如果batch size设为32显存不够,可以先降到16或者8。还有一个办法是开启梯度累积,YOLOv5里可以通过设置--batch-size配合--nbs参数来模拟更大batch的训练。比如你想等效64的batch,但显存只能跑16,那就设--batch-size 16 --nbs 64。
训练loss为NaN
这个大概率是学习率过大或者数据标注里有空文件。检查一下每张图片对应的txt文件是否为空,如果有空文件,训练时就会出问题。建议写个脚本遍历一下数据集目录,找出那些没有标注框的图片,把它们从数据集中剔除。
验证集mAP一直为0
先检查数据集的标签类别编号是否从0开始连续递增。比如你有6类,类别id不能出现0、1、2、3、4、6这种跳跃,否则模型在计算mAP时会发生维度不匹配。另外检查data.yaml中的类别顺序和标注工具导出顺序是否一致。
5.2 PyQt集成阶段典型问题
模型加载时界面假死
原因基本是我前面说的没有开线程。解决办法就是把模型加载和推理都放到QThread里,不要用time.sleep阻塞主线程。另外模型权重文件如果比较大(YOLOv5x能有好几百MB),加载耗时几秒是正常的,建议在加载前禁用界面按钮,加载完成后再启用。
摄像头打不开
这个问题很多时候是摄像头索引号不对。笔记本自带摄像头一般是0,外接USB摄像头可能是1或者2,可以写一个循环尝试打开0、1、2三个索引。如果是在虚拟机上运行,还需要把主机的摄像头设备连接到虚拟机,这个容易忽略。
检测结果显示不全或崩溃
检查推理结果转成图像时是否有results.render()[0]操作,这个返回的是numpy数组。在传给它之前必须确保图像的副本没有被其他线程占用。我用PyQt显示图像时习惯先把numpy数组转成QImage,示例代码如下:
from PyQt5.QtGui import QImage, QPixmap h, w, ch = frame.shape bytes_per_line = ch * w q_img = QImage(frame.data, w, h, bytes_per_line, QImage.Format_RGB888).rgbSwapped() self.label.setPixmap(QPixmap.fromImage(q_img))这里的rgbSwapped()很关键,因为OpenCV是BGR顺序,而QImage需要RGB,不调用这个函数的话画面颜色会偏蓝。
6. 实测效果与性能调优记录
6.1 不同模型规模对比
我这次训练跑了两组模型做对比,一组是YOLOv5s,一组是YOLOv5n。在同一个测试集上,YOLOv5s的mAP@0.5达到了0.82,YOLOv5n是0.76。但在CPU推理速度上,YOLOv5n比s快了将近40%。所以最终给PyQt界面用的模型我选的是YOLOv5s,换成YOLOv5n也是可行的,看你的机器配置和使用场景。
| 模型 | 参数量 | mAP@0.5 | 推理耗时(CPU) | 推理耗时(GPU) |
|---|---|---|---|---|
| YOLOv5n | 1.9M | 0.76 | ~80ms | ~20ms |
| YOLOv5s | 7.2M | 0.82 | ~130ms | ~25ms |
6.2 界面推理链路优化
界面端为了保持流畅,我做了三个优化:
- 输入图像缩放:推理前把图像resize到640x640,而不是直接用原始分辨率。模型对输入尺寸敏感,太小的目标如果原图很大,直接resize可能会丢失细节,所以这一步在YOLOv5内部已经处理好,不需要额外干预。
- 帧率控制:视频流检测时每隔一帧推理一次,这样在机器人巡游场景中既能保持连续检测,又能降低CPU占用。
- 结果缓存:对于静态图片检测,如果用户没有改变阈值参数,同一张图片不重复推理,直接显示上一次的结果。这个优化在看图和调参配合时特别有用。
实际跑下来,在i5-12400 + 16GB内存 + 无独立显卡的电脑上,视频检测可以达到每秒7~8帧,基本满足实时监控需求。如果换成有RTX 3050以上的显卡,流畅度就非常舒服了。
7. 我对这套系统的几点心得
整个项目做完,最大的体会是:YOLOv5本身极其成熟,真正的难点在于数据的质量和整个工程链路的完整性。你要花很多时间去清洗数据、检查标注、调参数,而这些工作并不是算法层面的难,它要求你非常耐心。
如果你也想复现这个项目,我建议你按下面的顺序来:
- 第一步,别直接训练。先把数据集的格式、类别、数量核对清楚,做一次数据可视化,把每一张图像的标注框画出来看看。
- 第二步,用默认参数训练20轮,确认能跑通,再回来调参。不要一上来就改一堆超参数,不然根本不知道是数据问题还是参数问题。
- 第三步,确认模型效果达标后,再开始写PyQt界面。界面就封装推理逻辑,不值得你花太多时间在美化上。
这套系统后续还能扩展的东西很多,比如加入水体清晰度增强预处理、用ONNX Runtime替代PyTorch做推理、或者把界面里的检测结果自动保存成Excel报表。对我来说,当前这版已经能解决水下垃圾检测从模型到交互展示的全流程需求,拿来当毕设主体或者实际项目原型都够用。
本文还有配套的精品资源,点击获取