水下生物目标检测实战：YOLO工程与PyTorch训练推理全流程解析

简介面向水下生物目标检测场景这份基于Python与PyTorch的深度学习资源包整合了YOLO模型训练与推理所需的数据集、脚本及预训练权重适合有一定深度学习基础、希望快速上手目标检测项目的开发者。资源共1830个文件压缩包112.1MB其中910张JPG水下生物图像、448个XML标注和453个TXT标签构成可直接使用的YOLO格式数据集3个Python脚本覆盖数据划分、模型训练与PyQt可视化识别另含训练好的pt模型权重与运行日志。已有111人学习。按01、02、03顺序运行脚本即可完成从数据集格式转换、模型训练到界面加载图片检测的完整流程若仅需调用已训练模型直接运行03pyqt.py即可弹出可视化界面进行实时识别。对于想深入理解YOLO训练流程或快速搭建水下检测演示系统的读者是一份可操作、结构清晰的参考方案。1. 水下生物目标检测这份 YOLO 工程能直接跑通训练和可视化识别做水下生物检测最烦的不是模型选型而是数据集和自己标注格式对不上。网上公开的水下数据集不少但大多是 VOC 或 COCO 格式要落到 YOLO 训练流程里得先写转换脚本一个字段对不上就全盘翻车。这份基于 Python 深度学习的工程把路铺好了——数据集已经整理好、标注是 YOLO 格式环境装完按顺序跑三个 py 文件就能完成从数据划分到可视化检测的完整流程。它用的是 PyTorch 安装环境检测部分带 PyQt5 界面既适合刚入门目标检测的学生跑通一个完整项目也适合需要快速出演示效果的从业者。整个工程不是黑匣子三个脚本分别负责数据划分、模型训练、界面推理每一段都能单独调试这也意味着你可以只替换数据集把整个流程复用到其他检测场景。2. 环境配置与目录结构先看懂每个文件是干什么的2.1 requirement.txt 之外的硬性环境要求项目摘要里说得很直接代码基于 Python PyTorch 环境环境需要自己配。常见做法是先装 PyTorch 再装其他依赖因为 requirement.txt 里如果直接写 torch 的版本号默认会装 CPU 版训练速度会让你怀疑人生。建议先到 PyTorch 官网按 CUDA 版本生成安装命令再执行 requirement.txt 的剩余依赖# 先确认 CUDA 版本再装 PyTorch以 CUDA 11.8 为例 nvcc --version pip install torch2.0.1 torchvision0.15.2 --index-url https://download.pytorch.org/whl/cu118 # 装剩余依赖 pip install -r requirement.txt代码里涉及图像读取和 PyQt 界面显示opencv-python 和 PyQt5 这两个包版本冲突是常见问题。如果运行 03pyqt.py 时报 Qt 相关错误直接强制重装 PyQt5 的匹配版本即可。Python 版本建议 3.8 到 3.10太新的版本部分依赖包还没有预编译轮子会现场编译 C 扩展耗时长且容易失败。2.2 工程目录里的产物文件各是什么拿到压缩包解压后除了三个 py 文件和数据集文件夹还能看到一串带 events、cache 后缀的文件。这些都是训练过程中自动生成的看懂它们能帮你判断训练是否正常。labels.cache 是数据加载时生成的标注缓存下次启动训练时不需要重新解析所有 txt 标注results.csv 记录每个 epoch 的 loss、mAP、precision 等指标val_batch0_labels.jpg 到 val_batch2_pred.jpg 是验证集的可视化结果左边是标注框右边是预测框打开看一眼就知道模型学到什么程度了。如果你打算重新训练保险做法是删掉 labels.cache因为数据集一旦改动缓存里的信息就和实际标注不一致训练时会报 shape mismatch 之类的问题。type 命令在命令行依次运行 01、02、03 三个文件即可如果只做推理验证直接跑 03pyqt.py 就行它默认加载训练好的模型权重。3. 数据划分与 YOLO 格式转换01 脚本的完整逻辑拆解3.1 从原始标注到 YOLO txt 的转换细节运行python 01划分数据集.py后脚本做三件事把标注文件转成 YOLO 格式的 txt、生成 train.txt 和 val.txt、生成 data.yaml。YOLO 格式的标注是归一化后的坐标每行内容为类别id x_center y_center width height四个坐标值都在 0 到 1 之间。原始标注如果来自 LabelImg 的 VOC XML 格式需要从 XML 里提取 object 的 name 和 bndbox 坐标再除以图片宽高做归一化import xml.etree.ElementTree as ET def voc_to_yolo(xml_file, img_w, img_h, class_list): tree ET.parse(xml_file) root tree.getroot() yolo_lines [] for obj in root.iter(object): name obj.find(name).text cls_id class_list.index(name) # 类别不在列表里会直接报错 box obj.find(bndbox) xmin float(box.find(xmin).text) ymin float(box.find(ymin).text) xmax float(box.find(xmax).text) ymax float(box.find(ymax).text) x_center (xmin xmax) / 2 / img_w y_center (ymin ymax) / 2 / img_h w (xmax - xmin) / img_w h (ymax - ymin) / img_h yolo_lines.append(f{cls_id} {x_center:.6f} {y_center:.6f} {w:.6f} {h:.6f}) return yolo_lines这段转换逻辑里最容易出错的是class_list.index(name)如果图片里有类别没在 class_list 里定义程序会抛 ValueError 直接中断。我一般会在脚本开头加上一个类别检查先把整个数据集里出现的所有 name 打印出来再手动对齐 class_list避免训到一半才发现漏了类别。另外注意坐标归一化用的是图片的宽高拿错宽高会导致标注框偏移训练时 loss 怎么调都降不下去。3.2 train.txt、val.txt 和 data.yaml 的生成逻辑数据划分的比例没有固定标准常见做法是 8 : 2 或 9 : 1。脚本里会随机打乱图片路径列表按比例切分后分别写入 train.txt 和 val.txt每行是一个图片的绝对路径。data.yaml 内容是 YOLO 训练时的数据配置入口包含 train/val 路径、类别数和类别名train: /path/to/your/dataset/train.txt val: /path/to/your/dataset/val.txt nc: 5 names: [holothurian, echinus, scallop, starfish, water_weeds]这里 nc 的值必须和 names 列表长度一致否则训练会报错。如果数据集里的类别在 names 里找不到属于哪个类别的目标都会被忽略或当成背景处理这部分得仔细核对。data.yaml 的路径建议写绝对路径相对路径在 Windows 和 Linux 下的解析方式不同换个环境就要改一次。3.3 划分后如何验证数据没问题训练前先做一次数据校验能省不少时间。我一般会拿 01 脚本生成的训练集路径随机读取十几张图片和对应的 txt 标注把框画在原图上保存肉眼确认标注框贴不贴目标、方向有没有反、类别对不对。这个习惯救过我很多次之前导出的标注里 xmin 和 ymin 全部是 0看了可视化才反应过来是 XML 解析时字段名写错了。4. 训练与可视化推理从 02train.py 到 03pyqt.py 的完整链路4.1 02train.py 训练参数怎么调下载的工程里 02train.py 的默认参数一般对应脚本里的默认值核心参数是 epochs、batch-size、img-size。batch-size 的第一优先级是显存大小6GB 显存跑 640x640 输入batch-size 设在 8 到 16 之间比较稳超出显存直接 OOM 报错。img-size 决定输入分辨率640 是速度和精度的平衡点水下生物有不少小目标分辨率太低会把海参、海胆这类目标直接糊掉预算够就上 640 以上。训练启动后观察终端输出重点关注 loss 下降趋势和 mAP 曲线。如果 loss 在最初的几个 epoch 没有明显下降先检查学习率YOLO 系列默认学习率是 0.01数据量小或 batch-size 小的时候可以调到 0.001。如果 GIOU loss 在训练后期持续震荡说明学习率偏大或数据里有大量错标框此时优先检查数据而不是调参数。训练正常结束后weights 目录下会有 best.pt 和 last.ptbest.pt 是验证集上 mAP 最高的权重界面推理用的就是它。results.csv 里逐行记录每个 epoch 的指标打开看一眼 precision 和 recall 的变化曲线能判断模型有没有过拟合或欠拟合。4.2 03pyqt.py 界面推理的工作过程03pyqt.py 做的事情本质上是加载模型权重、读入图片、前向推理、画框显示。PyQt 界面上的加载图片按钮调用 QFileDialog 选择图片文件检测按钮把图片传给模型推理再把结果转换回 QImage 显示在窗口上。对应核心逻辑如下from PyQt5.QtWidgets import QFileDialog, QLabel from PyQt5.QtGui import QImage, QPixmap import cv2 import torch class DetectWindow: def __init__(self, model_path): self.model torch.hub.load(ultralytics/yolov5, custom, pathmodel_path) self.model.conf 0.25 # 置信度阈值低于这个值的检测框会被过滤 self.model.iou 0.45 # NMS 的 IoU 阈值重叠框去重 def load_image(self): fname, _ QFileDialog.getOpenFileName(None, 选择图片, , 图片文件 (*.jpg *.png)) self.img cv2.imread(fname) self.img cv2.cvtColor(self.img, cv2.COLOR_BGR2RGB) def detect(self): results self.model(self.img) rendered results.render()[0] h, w, ch rendered.shape bytes_per_line ch * w qimg QImage(rendered.data, w, h, bytes_per_line, QImage.Format_RGB888) self.label.setPixmap(QPixmap.fromImage(qimg))置信度阈值 conf 的取值直接影响识别结果的误检和漏检。阈值设太低0.1 以下会框出一堆背景杂物设太高0.5 以上则小目标和水下模糊目标容易被漏掉。水下生物数据集里背景复杂度高我会先把阈值设成 0.3 跑一遍看结果再往两个方向微调。PyQt 界面里如果要加一个置信度滑条改动也不大核心就一行self.model.conf slider.value() / 100。4.3 训练产物和界面推理的关系跑通整个流程后你会发现02train.py 训练出的 best.pt 就是 03pyqt.py 的输入。如果训练中断过或没跑完weights 目录下是空的界面推理会报模型文件不存在的错误。摘要里提到直接运行 03pyqt.py 即可调用已训练好的模型说明压缩包自带权重如果想重新训练又不破坏原有权重记得先把 weights 目录备份一份。5. 避坑与常见问题排查四个高频翻车现场5.1 训练到一半显存爆了现象02train.py 运行后终端报CUDA out of memory程序直接退出显卡看着没跑多少 epoch 就撑不住了。 原因batch-size 设太大输入分辨率高导致中间特征图占满显存或者同时跑着其他程序占用显存。 解决按显存大小设置 batch-size6GB 显存从 batch-size 8 起步往下调直至不报错。也可以打开 YOLOv5 源码中的--cache参数把数据缓存到内存减少显存占用或使用梯度累积功能让显存占用折半。5.2 labels.cache 文件导致的报错现象改动数据集后直接运行 02train.py报assertion error: shape mismatch或数据加载卡住不动。 原因工程里残留了之前训练生成的 labels.cache 缓存新数据集标注内容和缓存不一致加载逻辑读取到旧的标注信息。 解决删掉数据集目录下的 labels.cache 文件再重新训练。这是个很玄学的问题因为报错信息不会直说缓存过期如果你排查了很久找不到原因多半是这里。5.3 中文路径导致图片读取失败现象03pyqt.py 加载图片时界面无反应控制台报Could not find a writer for the specified extension或 cv2.imread 返回 None。 原因OpenCV 的 imread 函数对中文路径支持不好Windows 系统下尤其明显文件路径里有中文就会读取失败。 解决把整个工程路径和数据路径都改成纯英文图片文件名也不要带中文。我习惯把工程放在D:\projects\underwater_detect\这类纯英文目录下省去一堆潜在问题。5.4 PyQt 界面启动卡死或闪退现象运行 03pyqt.py 后界面弹不出来或点击检测按钮时界面卡住变成白屏无响应。 原因检测逻辑是同步执行的图片推理期间主线程被阻塞UI 无法刷新或者 PyQt5 和 OpenCV 的 DLL 冲突导致加载失败。 解决推理逻辑放到 QThread 线程里运行时用 QTimer 定时刷新界面状态。如果是 DLL 冲突重新安装匹配版本的 PyQt5或者把 cv2 的显示逻辑改成 PyQt 的 QImage 转换方式。6. 进阶用法把单张图片检测扩展成批量推理和实时检测跑通 PyQt 界面后整个工程的能力边界就清楚了——它解决的是单张图片的交互式检测。但实际项目里遇到的往往是批量和视频流场景比如水下机器人传回一段录像需要逐帧分析或者摄像头实时画面需要持续检测。把 03pyqt.py 里的模型推理部分抽出来稍微改造就能扩展成批量推理脚本。核心逻辑是把读图片的循环改成遍历目录下的所有文件检测结果统一输出到指定文件夹import torch import cv2 import glob import os model torch.hub.load(ultralytics/yolov5, custom, pathweights/best.pt) model.conf 0.3 img_dir test_images/ out_dir test_results/ os.makedirs(out_dir, exist_okTrue) for img_path in glob.glob(os.path.join(img_dir, *.jpg)): img cv2.imread(img_path) results model(img) # results.xyxy[0] 是每个检测框的坐标、置信度和类别id rendered results.render()[0] fname os.path.basename(img_path) cv2.imwrite(os.path.join(out_dir, fname), rendered) print(fdone: {fname}, detections: {len(results.xyxy[0])})模型推理的结果保存在results.xyxy[0]它是一个 N 行 6 列的 Tensor每行对应一个检测框前四列是 x1、y1、x2、y2 坐标第五列是置信度第六列是类别索引。批量跑完后统计所有图片的检测框数量可以大致判断模型在不同光照、不同水质条件下的稳定性。这个脚本在工程里可以直接加入新的 py 文件不影响原有的三个文件。验证模型效果时还有一个习惯值得养成把检测结果按置信度分桶统计。比如设定 0.5、0.7、0.9 三档分别统计每一档的检测框数量和分布。如果大部分检测框落在低置信度区间说明模型的区分度不够需要补充训练数据或调高 conf 阈值如果高置信度区间存在大量误检则说明数据里存在相似干扰物。这个分析能帮你判断模型能不能上生产环境而不只是「看着差不多」。我那之后每次拿到新的检测工程都会先跑一遍批量推理脚本把所有结果图拼成一张大图快速翻一遍再决定是调参数还是补数据基本不会直接进训练环节。这个流程走下来翻车概率小了很多希望帮到你。本文还有配套的精品资源点击获取