ARTICLE DETAIL

资讯详情

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

水下垃圾检测实战:YOLOv5数据集构建、训练调参与边缘部署全流程

水下垃圾检测实战:YOLOv5数据集构建、训练调参与边缘部署全流程 简介这份资源面向从事计算机视觉、环境监测或水下机器人研究的学生与工程师提供一套可直接复现的YOLOv5水下垃圾检测方案解决真实海洋场景中垃圾目标识别与分类问题。压缩包共147个文件约213.98MB包含31个Python脚本、42个YAML配置、3个pt权重文件以及jpg图片、xml标注、png曲线图、sh脚本与Dockerfile等覆盖训练、推理与部署环节。资源内含训练好的水下垃圾检测权重、PR曲线与loss曲线并附带数千张经labelimg标注的真实场景图片标签同时提供VOC与YOLO两种格式类别涵盖metal、wood、plastic、rubber、cloth等场景丰富多样。目前已有1929人学习下载读者可借此快速完成数据准备、模型训练、指标分析与结果复现适合作为课程设计、科研实验或工程落地的参考模板。1. 水下垃圾检测为什么值得用 YOLOv5 做一遍河道巡检船回传的视频里塑料瓶、泡沫箱、废弃渔网混在浮萍和水草之间人眼盯十分钟就疲劳漏检率直线上升。水下垃圾检测要解决的就是这件事把水面和水下成像设备拍到的画面自动框出垃圾目标并分类。YOLOv5 在这个场景里被反复提起不是因为它最新而是因为它够稳——单阶段检测、工程化脚本齐全、对中小数据集友好配合一份标注好的数据集和训练好的模型能在普通显卡上跑通完整链路。适合谁做环保监测的嵌入式工程师、想入门目标检测的学生、需要快速验证水下视觉方案的算法同学。这篇笔记按「数据集怎么组织 → 模型怎么训 → 参数怎么调 → 坑在哪」的顺序讲目标是让你拿到一份标注好的水下垃圾数据集后能独立训出自己的权重并部署到边缘设备上。2. 水下垃圾数据集从标注格式到 YOLOv5 目录结构2.1 标注好的数据集长什么样一份能直接喂给 YOLOv5 的水下垃圾数据集通常包含 images 和 labels 两个平行目录图片是 jpg 或 png标签是 txt。每行标签格式为class_id x_center y_center width height坐标全部归一化到 0~1。水下场景的类别一般不会太多常见的是塑料瓶、塑料袋、渔网、泡沫、金属罐、其他垃圾这几类类别数直接决定后面模型输出通道数。拿到数据集先别急着训用下面这段脚本统计类别分布和标注框尺寸心里有数再动手import os from collections import Counter label_dir datasets/underwater/labels/train cls_counter Counter() wh_list [] for name in os.listdir(label_dir): if not name.endswith(.txt): continue with open(os.path.join(label_dir, name)) as f: for line in f: parts line.strip().split() if len(parts) ! 5: continue c, x, y, w, h parts cls_counter[int(c)] 1 wh_list.append((float(w), float(h))) print(类别分布:, cls_counter) if wh_list: avg_w sum(w for w, _ in wh_list) / len(wh_list) avg_h sum(h for _, h in wh_list) / len(wh_list) print(f平均宽高: {avg_w:.4f}, {avg_h:.4f})逻辑说明遍历标签目录统计每个类别的框数量同时算平均宽高。参数说明label_dir换成你自己的训练标签路径如果发现某个类别只有几十个框说明数据不平衡后面训练要开类别权重或者做增强。平均宽高偏小比如小于 0.05意味着目标很小需要调大输入分辨率。2.2 转成 YOLOv5 要求的目录与 data.yamlYOLOv5 不强制目录名但社区惯例是images/train、images/val、labels/train、labels/val。如果你的数据集是 VOC 的 xml 或者 COCO 的 json需要先转。下面是一个把 VOC 转 YOLO 的片段import xml.etree.ElementTree as ET import os def voc_to_yolo(xml_path, out_txt, classes): tree ET.parse(xml_path) root tree.getroot() size root.find(size) w int(size.find(width).text) h int(size.find(height).text) lines [] for obj in root.iter(object): cls obj.find(name).text if cls not in classes: continue cls_id classes.index(cls) bbox obj.find(bndbox) x1 float(bbox.find(xmin).text) y1 float(bbox.find(ymin).text) x2 float(bbox.find(xmax).text) y2 float(bbox.find(ymax).text) xc (x1 x2) / 2 / w yc (y1 y2) / 2 / h bw (x2 - x1) / w bh (y2 - y1) / h lines.append(f{cls_id} {xc:.6f} {yc:.6f} {bw:.6f} {bh:.6f}) with open(out_txt, w) as f: f.write(\n.join(lines)) classes [bottle, bag, net, foam, can, other] voc_to_yolo(annotations/001.xml, labels/train/001.txt, classes)逻辑说明解析 xml 里的宽高和 bbox归一化后写成 YOLO 格式。参数说明classes列表顺序必须和后面 data.yaml 里的 names 完全一致否则类别会错位。转换完抽查几个 txt确认没有负数或大于 1 的值。data.yaml 是训练入口配置文件内容如下path: ../datasets/underwater train: images/train val: images/val nc: 6 names: [bottle, bag, net, foam, can, other]nc是类别数names顺序和上面 classes 一致。路径建议用相对路径换机器时少改一处。2.3 数据增强在水下场景的取舍水下图像普遍偏绿偏蓝、对比度低、有悬浮颗粒。YOLOv5 默认开启 mosaic、HSV 增强、随机翻转。我的经验是HSV 的 h 通道增强可以保留s 和 v 适当调大能模拟不同浑浊度mosaic 对小目标有帮助但如果你的垃圾目标本来就大mosaic 拼接后目标被裁切反而伤性能可以调低mosaic概率。翻转方面水平翻转安全垂直翻转要谨慎——水下拍摄有重力方向倒置的垃圾不符合物理常识可能引入噪声。3. 用 YOLOv5 训练水下垃圾模型命令、参数与显存控制3.1 环境配置与最小训练命令环境用 conda 建一个干净环境装 PyTorch 和依赖。版本不必追新能对上 CUDA 即可。常见做法是conda create -n yolo_uwater python3.9 -y conda activate yolo_uwater pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118 git clone https://github.com/ultralytics/yolov5 cd yolov5 pip install -r requirements.txt逻辑说明前三行建环境装 PyTorch后三行拉源码装依赖。参数说明CUDA 版本按你显卡驱动选30 系卡用 cu118 比较稳。装完跑python detect.py --weights yolov5s.pt --source data/images验证环境能出图就说明基础链路通了。最小训练命令python train.py \ --data ../datasets/underwater/data.yaml \ --weights yolov5s.pt \ --img 640 \ --batch 16 \ --epochs 100 \ --device 0 \ --project runs/train \ --name uwater_v1逻辑说明--data指向配置文件--weights用预训练权重做迁移学习--img是输入分辨率--batch是批大小--epochs是训练轮数。参数说明显存 8G 时--batch 16加--img 640通常能跑如果 OOM先把 batch 降到 8再考虑降 img。--device 0指定第一块 GPUCPU 训练把 device 改成 cpu但速度会慢很多。3.2 超参数怎么调学习率、锚框与损失权重YOLOv5 的超参数集中在data/hyp.scratch.yaml。水下垃圾数据集通常不大我一般会改这几项参数默认值建议值理由lr00.010.005~0.01小数据集学习率低一点更稳lrf0.010.01余弦退火终点保持默认warmup_epochs33~5小数据集多暖几轮box0.050.05框回归损失权重不动cls0.50.5~1.0类别不平衡时调大mosaic1.00.5~1.0目标大就调低锚框方面YOLOv5 自带自动锚框计算训练前会跑 k-means。如果你的垃圾目标尺寸和 COCO 差异大可以手动跑一次python utils/autoanchor.py --data ../datasets/underwater/data.yaml --img 640逻辑说明脚本会统计标注框分布重新聚类出 9 个锚框。参数说明如果输出的 best possible recall 低于 0.98说明锚框和你的数据不匹配把新锚框写进模型配置或 hyp 文件。3.3 训练过程看什么loss 曲线与 mAP训练启动后终端会打印每轮的 box_loss、obj_loss、cls_loss 和 mAP0.5。重点看三个信号box_loss 持续下降说明框回归在学obj_loss 震荡不降可能是学习率太大或标注有噪声mAP0.5 在 50 轮后还不涨检查数据增强是否过强。YOLOv5 会在runs/train/uwater_v1下生成 results.csv 和混淆矩阵用 pandas 读一下就能画曲线import pandas as pd df pd.read_csv(runs/train/uwater_v1/results.csv) df.columns df.columns.str.strip() print(df[[epoch, train/box_loss, metrics/mAP_0.5]].tail(10))逻辑说明读训练日志看最后 10 轮的 loss 和 mAP。参数说明列名可能因版本略有差异用df.columns先确认。如果 mAP 卡在 0.3 以下优先怀疑标注质量而不是模型。4. 推理、验证与部署把训练好的模型用起来4.1 用验证集跑一遍指标训练完先别急着上设备用val.py在验证集上跑一遍python val.py \ --data ../datasets/underwater/data.yaml \ --weights runs/train/uwater_v1/weights/best.pt \ --img 640 \ --batch 16 \ --task val逻辑说明加载 best.pt在 val 集上算 mAP、precision、recall。参数说明--task val只做验证不保存预测图想看每张图的预测结果加--save-txt。如果 precision 高但 recall 低说明模型保守可以调低置信度阈值再测。4.2 单张图片和视频推理推理命令很直接python detect.py \ --weights runs/train/uwater_v1/weights/best.pt \ --source ../datasets/underwater/images/val \ --img 640 \ --conf 0.25 \ --iou 0.45 \ --save-txt逻辑说明--source可以是单张图、目录或视频文件--conf是置信度阈值--iou是 NMS 的 IoU 阈值。参数说明水下垃圾目标如果重叠严重把 iou 调到 0.5~0.6 减少误删如果漏检多conf 降到 0.15~0.2 试试。--save-txt会输出每张图的检测框坐标方便后续统计。4.3 部署到边缘设备的注意点树莓派 5 上部署自己训练的 YOLOv5 模型是热词里常出现的需求。实际做法是先把模型导出 ONNX再用 onnxruntime 或 OpenCV DNN 推理python export.py --weights runs/train/uwater_v1/weights/best.pt --include onnx --img 640逻辑说明导出 ONNX 格式去掉训练相关节点。参数说明--img要和训练时一致否则精度掉。树莓派上建议用 yolov5n 或 yolov5s输入降到 416 或 320帧率能到 5~10 FPS。如果要用 NCNN加--include ncnn但需要自己编译。提示边缘设备内存有限导出时开--dynamic会增大模型体积固定 batch 和尺寸更省资源。5. 水下垃圾检测的避坑与排查清单5.1 标注框越界导致训练报错现象训练启动几轮后报AssertionError: negative coordinates或直接 NaN。原因VOC 转 YOLO 时某些框的 xmax 超过图片宽度归一化后大于 1。解决转换脚本里加裁剪x1 max(0, min(x1, w))四个坐标都做一遍再归一化。5.2 类别不平衡让模型只认多数类现象mAP 看着还行但混淆矩阵里少数类全被预测成多数类。原因水下垃圾里塑料瓶最多渔网可能只有几十个框。解决在 hyp 文件里调大cls权重或者用--weights做类别加权采样更直接的办法是给少数类做离线增强复制几份再改亮度对比度。5.3 显存不够导致的 OOM现象CUDA out of memory训练中断。原因batch 或 img 太大或者验证时同时加载了模型和缓存。解决先把 batch 减半再降 img 到 512 或 416如果还不行加--cache ram改成--cache disk或者关掉--multi-scale。5.4 验证集 mAP 高但实际画面漏检现象val 集指标 0.8实际视频里垃圾框不出来。原因验证集和实际场景分布不一致比如训练图都是清水实际是浑水。解决从实际视频里抽帧标一小批加入训练集做一次微调同时检查推理时的 conf 阈值是不是设太高。5.5 导出 ONNX 后精度下降现象PyTorch 权重检测正常ONNX 推理框偏移或漏检。原因导出时 img 尺寸和推理时不一致或者 NMS 在 ONNX 里实现有差异。解决导出和推理都用同一个 img 尺寸ONNX 推理时自己写 NMSIoU 阈值和 detect.py 对齐。6. 把模型训稳的一个笨办法先过拟合再放开如果你拿到的水下垃圾数据集只有几百张别一上来就开满增强训 300 轮。我自己的习惯是先用 50 张图、关掉 mosaic 和翻转把模型训到训练集 mAP 接近 1.0确认标注和代码链路没问题然后再把全量数据放进去开增强从头训。这个「先过拟合」的笨办法能帮你快速区分是数据问题还是模型问题。验证模型是否真的学到东西可以做一个遮挡测试把验证集里某张图的垃圾区域涂黑看模型是否还会在那个位置出框。如果涂黑后仍然出框说明模型可能依赖背景纹理而不是目标本身这时候要考虑加更多不同背景的负样本。最后分享一个我踩过的坑有次训完模型指标很好部署到巡检船上却频繁误检水面反光。后来把反光区域裁出来当负样本加进去误检才降下来。水下场景的干扰源和陆地完全不同数据集里一定要留一部分「没有垃圾但容易误检」的图。希望帮到你。本文还有配套的精品资源点击获取
返回列表