ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

水下垃圾检测实战:YOLOv5从训练到树莓派部署全流程

水下垃圾检测实战:YOLOv5从训练到树莓派部署全流程 简介这份资源面向计算机视觉学习者与水下垃圾检测方向的开发者提供一套可直接复现的YOLOv5水下垃圾检测完整方案帮助解决从数据准备到模型训练、评估的全流程问题。压缩包共147个文件约213.98MB包含31个Python脚本、42个yaml配置、3个pt权重文件以及jpg图片、xml标注、csv训练日志、png曲线图等覆盖训练、推理、评估各环节所需材料。资源内含训练好的水下垃圾检测权重并附PR曲线、loss曲线等结果文件便于直接验证效果。数据集为VOC格式由labelimg标注的真实场景高质量图片同时提供VOC与yolo两种标签格式类别涵盖metal、wood、plastic、rubber、cloth等多种海洋垃圾目标场景丰富。目前已有1929人学习下载适合希望快速上手水下目标检测、开展课程设计或科研实验的读者参考使用。1. 水下垃圾检测为什么值得用 YOLOv5 做一遍水下垃圾检测这件事真正下过水的人都知道难点从来不在能不能检测到而在检测得稳不稳。水下图像普遍存在色偏、低对比度、悬浮颗粒遮挡加上垃圾目标本身形态多变——塑料袋会飘、瓶子会反光、渔网会缠绕——同一类目标在不同帧里长得完全不一样。YOLOv5 之所以在这个场景里被反复拿出来用核心原因是它在小目标召回和推理速度之间取得了比较好的平衡而且训练好的模型加标注好的数据集这套组合能让后来者直接跳过最耗时的数据采集阶段把精力放在调优和部署上。这套方案适合谁如果你手上有水下机器人、ROV 或者固定式水下监控设备想快速验证自动识别垃圾这条路走不走得通那用现成的 YOLOv5 权重加标注数据集做迁移训练是最短路径。如果你是从零开始做水下视觉项目这套东西也能帮你把 baseline 立起来后面再换 backbone 或者加注意力模块都有对照。下面我会按数据检查、环境配置、训练调参、部署推理、踩坑排查的顺序把每个环节的具体操作和参数含义讲清楚。2. 拿到标注数据集先别急着训练三个检查动作2.1 用脚本统计类别分布和标注框尺寸拿到标注好的数据集第一件事不是train.py而是先看清楚数据长什么样。常见做法是写一个统计脚本把每个类别的实例数、标注框的宽高分布、以及每张图的目标数量跑一遍。这一步能帮你判断是否存在严重的长尾问题——比如塑料瓶有 3000 个实例渔网只有 80 个那训练时就得考虑加权或者过采样。import os import xml.etree.ElementTree as ET from collections import Counter import matplotlib.pyplot as plt # 假设标注是 VOC 格式的 XML如果是 YOLO txt 格式改一下解析逻辑即可 label_dir datasets/labels/train class_counter Counter() box_sizes [] for xml_file in os.listdir(label_dir): if not xml_file.endswith(.xml): continue tree ET.parse(os.path.join(label_dir, xml_file)) root tree.getroot() for obj in root.findall(object): cls_name obj.find(name).text class_counter[cls_name] 1 bbox obj.find(bndbox) w float(bbox.find(xmax).text) - float(bbox.find(xmin).text) h float(bbox.find(ymax).text) - float(bbox.find(ymin).text) box_sizes.append((w, h)) print(类别分布:, class_counter) # 画一下宽高散点看有没有大量极小框 ws, hs zip(*box_sizes) plt.scatter(ws, hs, s1, alpha0.3) plt.xlabel(width) plt.ylabel(height) plt.savefig(box_distribution.png)这段脚本的逻辑很直接遍历标注文件累计类别计数同时收集每个框的宽高。参数上唯一需要注意的是label_dir要指向训练集标注目录验证集和测试集最好也各跑一遍确认分布一致。如果发现某个类别实例数低于总实例数的 5%训练时就要在data.yaml里考虑用copy_paste或者mixup做增强否则模型对这个类基本学不动。2.2 检查图像与标注是否一一对应水下数据集常见的一个坑是图像和标注文件名对不上——比如图像叫frame_001.jpg标注叫frame_001.xml但实际内容对应的是frame_002。这种错位在训练时不会报错但会让模型学出完全错误的映射。我一般会写一个快速校验脚本import os img_dir datasets/images/train lbl_dir datasets/labels/train img_names {os.path.splitext(f)[0] for f in os.listdir(img_dir)} lbl_names {os.path.splitext(f)[0] for f in os.listdir(lbl_dir)} only_img img_names - lbl_names only_lbl lbl_names - img_names print(f有图无标注: {len(only_img)} 个) print(f有标注无图: {len(only_lbl)} 个) if only_img: print(示例:, list(only_img)[:5]) if only_lbl: print(示例:, list(only_lbl)[:5])逻辑说明用集合差集找出不匹配的文件。参数上没什么可调的但要注意如果数据集里混了.DS_Store或者Thumbs.db这类系统文件先过滤掉再统计。有图无标注的样本可以直接删掉或者补标有标注无图的说明标注文件是多余的留着会让 dataloader 报错。2.3 把 VOC 格式转成 YOLO 格式的边界处理YOLOv5 默认吃的是 YOLO txt 格式每行class_id x_center y_center width height全部归一化到 0 到 1。如果你拿到的标注是 VOC XML转换时最容易翻车的地方是坐标越界——水下图像有时候会有标注框超出图像边界的情况直接归一化会得到大于 1 的值训练时 loss 会炸。import os import xml.etree.ElementTree as ET def voc_to_yolo(xml_path, img_w, img_h, class_map): tree ET.parse(xml_path) root tree.getroot() lines [] for obj in root.findall(object): cls_name obj.find(name).text if cls_name not in class_map: continue cls_id class_map[cls_name] bbox obj.find(bndbox) xmin max(0, float(bbox.find(xmin).text)) ymin max(0, float(bbox.find(ymin).text)) xmax min(img_w, float(bbox.find(xmax).text)) ymax min(img_h, float(bbox.find(ymax).text)) # 过滤掉宽高为 0 的无效框 if xmax xmin or ymax ymin: continue x_center (xmin xmax) / 2.0 / img_w y_center (ymin ymax) / 2.0 / img_h w (xmax - xmin) / img_w h (ymax - ymin) / img_h lines.append(f{cls_id} {x_center:.6f} {y_center:.6f} {w:.6f} {h:.6f}) return lines关键参数是img_w和img_h必须和实际图像尺寸一致不能想当然用 640×640。转换完之后建议随机抽 20 张图用cv2.rectangle画出来肉眼看一下确认框的位置没偏。这个步骤花不了十分钟但能省掉后面几小时的排查时间。3. 环境配置与训练参数从 conda 到超参文件3.1 用 conda 建环境并锁定版本YOLOv5 的依赖不算复杂但版本不对照样出玄学问题。我一般用 conda 建一个干净环境Python 版本选 3.8 或 3.9这两个版本和 PyTorch 的兼容性最稳。conda create -n yolo5_underwater python3.9 -y conda activate yolo5_underwater # 安装 PyTorch根据你的 CUDA 版本选对应命令 # 这里以 CUDA 11.8 为例 pip install torch2.0.1 torchvision0.15.2 --index-url https://download.pytorch.org/whl/cu118 # 克隆 YOLOv5 仓库并安装依赖 git clone https://github.com/ultralytics/yolov5.git cd yolov5 pip install -r requirements.txt逻辑说明先建独立环境避免和系统里的其他包冲突然后装 PyTorch最后装 YOLOv5 的依赖。参数上唯一要确认的是 CUDA 版本用nvidia-smi看一下驱动支持的 CUDA 版本再去 PyTorch 官网找对应的安装命令。如果显存低于 8GB训练时 batch size 要往下调后面会讲。3.2 data.yaml 和超参数文件的写法YOLOv5 的数据配置集中在data.yaml里路径、类别数、类别名都在这里定义。水下垃圾检测一般类别不会太多常见的是塑料瓶、塑料袋、渔网、金属罐、其他垃圾这几类。# data_underwater.yaml path: ./datasets/underwater # 数据集根目录 train: images/train val: images/val test: images/test nc: 5 # 类别数 names: 0: plastic_bottle 1: plastic_bag 2: fishing_net 3: metal_can 4: other_trash参数说明path是根目录train、val、test是相对路径。nc必须和names的长度一致否则训练时分类头维度对不上会直接报错。类别名不要用中文YOLOv5 内部处理时对非 ASCII 字符支持不好容易在日志里乱码。超参数文件hyp.scratch-low.yaml是给从零训练用的如果你用预训练权重做迁移学习建议用hyp.finetune.yaml或者自己改一版。水下场景我一般会把hsv_h调到 0.015、hsv_s调到 0.7、hsv_v调到 0.4因为水下色偏严重颜色增强幅度大一点反而有帮助。mosaic保持 1.0但mixup调到 0.1 就够了太高会让水下模糊目标更难学。3.3 启动训练与关键参数解读训练命令本身不复杂但参数怎么设直接决定你能不能跑出可用的模型。python train.py \ --img 640 \ --batch 16 \ --epochs 150 \ --data data_underwater.yaml \ --weights yolov5s.pt \ --cfg models/yolov5s.yaml \ --hyp data/hyps/hyp.finetune.yaml \ --name underwater_v1 \ --cache逐参数说明--img 640是输入分辨率水下小目标多的话可以提到 800 甚至 1024但显存占用会成倍增加。--batch 16在 8GB 显存上跑 640 分辨率差不多是上限如果 OOM 就降到 8。--epochs 150是经验值水下数据集通常不大100 到 200 轮之间看 mAP 曲线决定要不要早停。--weights yolov5s.pt加载预训练权重这是迁移学习的关键不要从零开始。--cache把图像缓存到内存能加快训练速度但如果数据集超过内存容量就别加这个参数。训练过程中重点看三个指标box_loss是否稳定下降、mAP0.5是否在涨、precision和recall是否平衡。如果box_loss震荡严重大概率是学习率太高或者 batch size 太小可以试着把lr0从 0.01 降到 0.005。4. 推理部署从本地测试到树莓派落地4.1 用 detect.py 做本地推理验证训练完之后第一件事是用detect.py跑几张测试图确认模型不是只会在验证集上背答案。python detect.py \ --weights runs/train/underwater_v1/weights/best.pt \ --source datasets/underwater/images/test \ --img 640 \ --conf-thres 0.35 \ --iou-thres 0.45 \ --save-txt \ --project runs/detect \ --name test_v1参数说明--conf-thres 0.35是置信度阈值水下场景我一般设得比陆地上低一点因为水下目标模糊模型给出的置信度普遍偏低设太高会漏检。--iou-thres 0.45是 NMS 的 IoU 阈值如果同一类目标密集出现比如一堆瓶子堆在一起可以适当调高到 0.5 减少误抑制。--save-txt会把检测结果存成 txt方便后面做定量评估。4.2 在树莓派 5 上部署的模型转换树莓派 5 的算力比前代强不少但直接跑 PyTorch 模型还是吃力。常见做法是把模型导出成 ONNX 或者 NCNN 格式再用对应的推理引擎跑。# 导出 ONNX python export.py \ --weights runs/train/underwater_v1/weights/best.pt \ --include onnx \ --img 640 \ --batch 1 # 导出 NCNN树莓派上性能更好 python export.py \ --weights runs/train/underwater_v1/weights/best.pt \ --include ncnn \ --img 640 \ --batch 1逻辑说明export.py会把 PyTorch 模型转成中间格式。--batch 1是因为树莓派上一般一次只推理一张图batch 设大了反而浪费内存。导出 NCNN 之后把生成的model.ncnn.param和model.ncnn.bin拷到树莓派上用 NCNN 的 Python 或者 C 接口加载。实测在树莓派 5 上YOLOv5s 的 NCNN 模型跑 640 分辨率大概能到 5 到 8 FPS对于水下机器人这种低速场景够用了。4.3 低显存运行模型的三个实用技巧如果你手头的设备显存有限比如只有 4GB 的 Jetson Nano 或者老款显卡下面三个方法能帮你把模型跑起来。第一个是降低输入分辨率从 640 降到 416 甚至 320显存占用能减少一半以上代价是小目标召回会下降。第二个是用--half开启 FP16 推理显存占用直接砍半速度还能提升 20% 左右但要注意有些老显卡对 FP16 支持不好开了反而更慢。第三个是裁剪模型把 YOLOv5s 换成 YOLOv5n参数量从 7.2M 降到 1.9M精度会掉几个点但至少能跑起来。# FP16 推理示例 python detect.py \ --weights best.pt \ --source test_images \ --img 416 \ --half \ --device 0参数说明--half开启半精度--device 0指定用第一块 GPU。如果报错说设备不支持 FP16去掉--half就行。--img 416要和导出时的分辨率一致否则精度会异常。5. 水下场景避坑与排查五条血泪经验5.1 现象训练 loss 正常下降但 mAP 始终在 0.1 以下原因最常见的是标注格式不对。YOLO 格式要求坐标归一化到 0 到 1如果转换时忘了除以图像宽高或者除错了尺寸模型学到的就是一堆错位框。另一个可能是data.yaml里的nc和实际类别数不一致分类头输出维度对不上但 YOLOv5 在某些版本里不会直接报错而是静默地把多余类别忽略掉。解决先用 2.1 节的统计脚本确认标注框的宽高都在 0 到 1 之间再检查data.yaml的nc和names长度是否一致。如果都没问题拿一张训练图跑detect.py看模型输出的框和真实标注差多少差得离谱就是格式问题。5.2 现象验证集 mAP 很高但实际推理时漏检严重原因验证集和实际场景的分布不一致。水下数据集如果是在特定水域、特定光照条件下采集的模型会过拟合到那个域。实际部署时换了水域或者水深变了图像色偏和浊度都不一样模型就懵了。解决在训练时加入更强的颜色增强把hsv_h、hsv_s、hsv_v的幅度调大。另外可以在验证集里混入一些实际部署场景的图哪怕没有标注用detect.py跑一遍看置信度分布如果普遍低于 0.2说明域偏移严重需要补数据重新训练。5.3 现象树莓派上推理速度只有 1 到 2 FPS原因大概率是没用 NCNN 或者 ONNX 加速直接跑的 PyTorch 模型。树莓派的 CPU 跑 PyTorch 的卷积层效率很低而且默认用的是 FP32计算量翻倍。解决按 4.2 节导出 NCNN 模型然后用 NCNN 的推理接口。如果还是慢把输入分辨率降到 320并且开启 NCNN 的多线程选项。另外确认树莓派没在跑其他占 CPU 的进程htop看一下负载。5.4 现象检测框大量重叠同一个目标被框了好几次原因NMS 的 IoU 阈值设得太高或者模型对同一类目标的置信度普遍偏高。水下垃圾有时候形态模糊模型会在同一个位置输出多个框每个框的置信度都差不多。解决把--iou-thres从默认的 0.45 降到 0.3 到 0.35 之间让 NMS 更激进地抑制重叠框。如果降了之后漏检变多说明模型本身对目标的定位不够准需要回去检查标注质量看是不是同一个目标被标了多次。5.5 现象训练到一半突然报 CUDA out of memory原因YOLOv5 在训练过程中会动态调整一些缓存如果--cache开了但内存不够或者 dataloader 的 worker 数太多显存会在某个 epoch 突然爆掉。解决去掉--cache把--workers从默认的 8 降到 4 或 2--batch再降一档。如果还不行用--img 512代替 640。另外可以在训练命令前加CUDA_VISIBLE_DEVICES0确保只用一块卡避免多卡显存分配不均。6. 把 mAP 再往上推两个点的三个微调技巧训练完第一版模型之后大部分人会发现 mAP0.5 卡在 0.6 到 0.7 之间上不去。这时候别急着换模型结构先把下面三个技巧试一遍通常能再挤出两三个点。第一个是调整锚框。YOLOv5 默认的锚框是在 COCO 上聚类出来的水下垃圾的宽高比和 COCO 里的目标差别很大。用utils/autoanchor.py在你自己的数据集上重新跑一遍 k-means把生成的锚框写进模型配置里。具体操作是在训练命令里加--noautoanchor关掉自动锚框然后手动把anchors参数替换成聚类结果。这个步骤听起来麻烦但对小目标召回提升很明显。第二个是测试时增强TTA。推理时把图像翻转、缩放几次分别跑模型再把结果融合能稳定提升一个点左右。YOLOv5 的detect.py里加--augment就能开启。python detect.py \ --weights best.pt \ --source test_images \ --augment \ --conf-thres 0.3参数说明--augment会做多尺度加翻转的 TTA推理时间大概变成原来的三倍。如果对实时性要求高这个技巧就不适合。--conf-thres在 TTA 模式下可以适当降低因为融合之后误检会被抑制掉一部分。第三个是难例挖掘。把验证集里模型漏检或者误检的图挑出来人工补标或者修正标注然后把这些图加到训练集里重新训一轮。这个做法最笨但最有效我自己的经验是每轮难例挖掘能带来一到两个点的提升做两三轮之后模型基本就稳定了。最后一个习惯每次训练完把results.csv存下来把关键参数和最终 mAP 记在一个表格里。水下场景的调参很依赖经验没有记录的话过两周你就不记得哪个参数组合跑出过最好的结果。这个习惯帮我省了很多重复试错的时间希望帮到你。本文还有配套的精品资源点击获取
返回列表