水下垃圾检测系统实战：基于YOLOv5与PyQt5的完整实现

简介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上跑到100FPS即使是CPU也能有可用的推理速度。而相比YOLOv8YOLOv5的部署生态更简单改起来也方便尤其要做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.8PyTorch 1.8CUDA 10.2以上即可如果是NVIDIA显卡装好CUDA和cuDNN注意PyTorch版本和CUDA版本要对应克隆YOLOv5仓库执行pip install -r requirements.txt训练前需要准备好data.yaml文件内容是指定训练验证路径和类别名。这个文件格式很简单但路径一定要写对最好是绝对路径避免相对路径在不同的工作目录下报错。关键训练参数我做了几个对比最终没有用默认参数因为默认参数很多时候不是最优的。参数默认值我的配置说明img-size640640水下小目标多不盲目降低输入尺寸batch-size1632显存允许的情况下尽量大能稳定训练epochs300200看模型收敛情况我200轮已经收敛patience10050验证集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、权重文件和混淆矩阵等可视化结果。看训练过程主要关注两条曲线mAP0.5和val_loss。我这次训练了200轮大约在120轮时mAP0.5已经到0.78之后提升很慢到190轮达到0.82。val_loss在80轮左右就已经趋于平缓后面基本在一个小范围内抖动说明模型已经收敛。刚开始跑的时候我遇到一个典型问题训练集loss下降但验证集loss不降反升也就是过拟合。当时数据集总量还不到2000张模型是从零开始训练的分类数又多。解决办法有两个方向使用预训练权重把--weights指定成yolov5s.pt用COCO预训练模型做初始化迁移学习后收敛更快泛化也更稳。降低模型复杂度从YOLOv5s换成YOLOv5n参数量少一半适合小数据集。如果检测结果里出现很多重复框或者误检可以调节后处理参数。也就是推理时的conf_thres和iou_thres。我建议默认conf_thres0.25可以改成0.3iou_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, pathbest.pt, force_reloadFalse) # 推理单张图片 results model(img, size640) # 获取渲染后的图像带框 rendered_img results.render()[0]这段代码是我在实际项目中用的精简版注意torch.hub.load第一次会联网拉取YOLOv5仓库之后设置了force_reloadFalse就不会每次都拉。如果你希望完全离线也可以直接把yolov5仓库克隆到本地用torch.hub.load(./yolov5, custom, pathbest.pt, sourcelocal)。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的mAP0.5达到了0.82YOLOv5n是0.76。但在CPU推理速度上YOLOv5n比s快了将近40%。所以最终给PyQt界面用的模型我选的是YOLOv5s换成YOLOv5n也是可行的看你的机器配置和使用场景。模型参数量mAP0.5推理耗时(CPU)推理耗时(GPU)YOLOv5n1.9M0.76~80ms~20msYOLOv5s7.2M0.82~130ms~25ms6.2 界面推理链路优化界面端为了保持流畅我做了三个优化输入图像缩放推理前把图像resize到640x640而不是直接用原始分辨率。模型对输入尺寸敏感太小的目标如果原图很大直接resize可能会丢失细节所以这一步在YOLOv5内部已经处理好不需要额外干预。帧率控制视频流检测时每隔一帧推理一次这样在机器人巡游场景中既能保持连续检测又能降低CPU占用。结果缓存对于静态图片检测如果用户没有改变阈值参数同一张图片不重复推理直接显示上一次的结果。这个优化在看图和调参配合时特别有用。实际跑下来在i5-12400 16GB内存 无独立显卡的电脑上视频检测可以达到每秒7~8帧基本满足实时监控需求。如果换成有RTX 3050以上的显卡流畅度就非常舒服了。7. 我对这套系统的几点心得整个项目做完最大的体会是YOLOv5本身极其成熟真正的难点在于数据的质量和整个工程链路的完整性。你要花很多时间去清洗数据、检查标注、调参数而这些工作并不是算法层面的难它要求你非常耐心。如果你也想复现这个项目我建议你按下面的顺序来第一步别直接训练。先把数据集的格式、类别、数量核对清楚做一次数据可视化把每一张图像的标注框画出来看看。第二步用默认参数训练20轮确认能跑通再回来调参。不要一上来就改一堆超参数不然根本不知道是数据问题还是参数问题。第三步确认模型效果达标后再开始写PyQt界面。界面就封装推理逻辑不值得你花太多时间在美化上。这套系统后续还能扩展的东西很多比如加入水体清晰度增强预处理、用ONNX Runtime替代PyTorch做推理、或者把界面里的检测结果自动保存成Excel报表。对我来说当前这版已经能解决水下垃圾检测从模型到交互展示的全流程需求拿来当毕设主体或者实际项目原型都够用。本文还有配套的精品资源点击获取