简介:本资源是面向人工智能与计算机视觉初学者及项目实践者的YOLOv5舰船目标检测专用数据集,聚焦海洋监控、智能航海等实际场景中的小目标、复杂背景下的舰船识别任务。压缩包共1648个文件,含549张JPG格式船舶图像、549份VOC标准XML标注文件(含精确边界框坐标与类别标签)及550个对应TXT格式标签,总大小56.16MB;XML与TXT双格式支持灵活适配YOLOv5训练流程,便于快速完成数据格式转换与预处理。目前已有1479人学习下载,资源结构规范、来源清晰(源自VOCtrainval2012子集),附带典型样本如2011_000238.jpg等,可直接用于模型训练、验证与可视化分析。读者可获得完整可用的舰船检测数据基础、标准化标注体系及YOLO格式转换参考依据,显著降低目标检测项目启动门槛。
1. 项目概述:从一份数据集开始,聊聊舰船检测的实战门道
最近在整理硬盘时,翻到了一个名为“boat-舰船检测数据集.rar”的压缩包,这让我想起了几年前做的一个海上目标监测项目。当时为了训练一个能准确识别各类舰船的模型,可没少在数据上折腾。今天,我就以这个数据集为引子,结合当下依然热门的YOLOv5框架,和大家深入聊聊如何从零开始,搞定一个舰船检测任务。无论你是刚接触计算机视觉的新手,还是想优化现有流程的老兵,希望这篇从数据准备到模型调优的完整复盘,能给你带来一些实实在在的参考。
这个“boat-舰船检测数据集”本身,就是一个非常典型的垂直领域目标检测数据集。它的核心价值在于,为我们提供了一个专注于“舰船”这一特定类别的图像集合,通常包含了各种型号、不同尺度、在各种海况和光照条件下的舰船图片,并且每张图片都带有精确的边界框标注。对于想研究海事监控、港口管理、海洋资源调查或者国防相关应用的开发者来说,这样一个高质量、标注规范的数据集是至关重要的起点。接下来,我会围绕如何使用YOLOv5来“消化”这个数据集,构建一个鲁棒的舰船检测器,展开详细的步骤解析和经验分享。
2. 核心需求解析与数据集深度剖析
在动手写一行代码之前,我们必须先想清楚:我们要用这个数据集解决什么问题?以及,这个数据集本身质量如何?这两个问题是所有后续工作的基石。
2.1 舰船检测的核心应用场景与挑战
舰船检测绝非一个“为检测而检测”的玩具项目,其背后对应着强烈的现实需求。最主要的应用场景包括:
- 海事交通监控与安全管理:在港口、航道、近海区域,自动识别和跟踪船只,用于流量统计、碰撞预警、非法闯入监测等。
- 海洋权益维护与态势感知:通过卫星遥感或无人机航拍图像,广域监测特定海域的船只活动情况。
- 渔业资源管理与非法捕捞监管:识别渔船类型,监控禁渔区内的作业活动。
- 搜救任务辅助:在广阔的海域中快速定位失事或需要援助的船只。
这些场景给模型带来了独特的挑战:
- 尺度变化剧烈:同一张图片里,近处的渔船可能占据大半画面,而远处的货轮只是几个像素点。模型必须同时具备识别大目标和小目标的能力。
- 背景复杂:海面并非一成不变的蓝色。它会有波浪、泡沫、云影、岛屿、海岸线等干扰,尤其在恶劣天气下,海天界线模糊,对模型的分割能力要求极高。
- 目标姿态多样:船只可能呈现正面、侧面、斜向等多种角度,部分被遮挡(如被浪花、其他船只遮挡)的情况也很常见。
- 类别内差异大:“舰船”是一个大类,下面可能包含邮轮、货船、油轮、帆船、快艇、军舰等,它们的外观、长宽比差异巨大。
理解了这些挑战,我们就能在数据准备和模型训练阶段做出更有针对性的设计。
2.2. “boat-舰船检测数据集”的预处理与质量评估实战
拿到一个.rar格式的数据集压缩包,第一步绝不是直接解压扔给模型。一个严谨的预处理流程能避免后续无数坑。
步骤一:解压与结构探查首先,解压“boat-舰船检测数据集.rar”。一个规范的YOLO格式数据集通常包含以下目录:
boat_dataset/ ├── images/ │ ├── train/ # 训练集图片 │ └── val/ # 验证集图片 └── labels/ ├── train/ # 训练集标签(.txt文件,与图片同名) └── val/ # 验证集标签如果解压后结构混乱,你需要手动按上述结构整理。关键点:每个.txt标签文件对应一张同名的图片,内容格式为[class_id] [x_center] [y_center] [width] [height],其中坐标和宽高都是相对于图片宽度和高度的归一化值(0-1之间)。
步骤二:数据质量诊断脚本我强烈建议在训练前运行一个简单的诊断脚本,检查数据健康度。以下是一个Python示例,使用OpenCV和PIL:
import os import cv2 from PIL import Image import matplotlib.pyplot as plt def dataset_diagnosis(data_dir): img_dir = os.path.join(data_dir, 'images', 'train') label_dir = os.path.join(data_dir, 'labels', 'train') img_files = [f for f in os.listdir(img_dir) if f.endswith(('.jpg', '.png', '.jpeg'))] print(f"Total training images: {len(img_files)}") # 检查图片能否正常打开,并统计尺寸分布 sizes = [] broken_imgs = [] for img_name in img_files[:100]: # 抽样检查前100张 img_path = os.path.join(img_dir, img_name) try: with Image.open(img_path) as img: sizes.append(img.size) # (width, height) except Exception as e: broken_imgs.append(img_name) # 检查标签文件是否存在且格式正确 missing_labels = [] for img_name in img_files: label_name = os.path.splitext(img_name)[0] + '.txt' label_path = os.path.join(label_dir, label_name) if not os.path.exists(label_path): missing_labels.append(img_name) else: # 可选:检查标签内容是否合法(数值在0-1之间,类别ID有效等) pass # 输出诊断报告 if sizes: avg_w = sum(s[0] for s in sizes) / len(sizes) avg_h = sum(s[1] for s in sizes) / len(sizes) print(f"Average image size: {avg_w:.0f}x{avg_h:.0f}") if broken_imgs: print(f"Warning: {len(broken_imgs)} images are broken.") if missing_labels: print(f"Critical: {len(missing_labels)} images have missing labels.") return sizes, broken_imgs, missing_labels # 使用 data_path = "./boat_dataset" sizes, broken, missing = dataset_diagnosis(data_path)步骤三:可视化与统计分析使用YOLOv5自带的工具是最高效的。在YOLOv5项目根目录下运行:
python utils/plots.py --data data/boat.yaml --output plots/这个命令会生成一系列分析图,包括:
- 标签分布图:查看所有边界框的中心点分布。理想情况是均匀分布在整个图像上。如果中心点大量聚集在图像中心,可能意味着你的数据多是远景拍摄,需要补充一些近景特写数据。
- 边界框尺寸分布图:显示标注框的宽度和高度的分布。这对于了解数据集中目标的尺度范围至关重要。如果大量框的宽高都非常小(比如<0.05),说明小目标很多,在训练时需要调整模型结构或损失函数来加强对小目标的关注。
- 类别平衡图:如果你的数据集有多个舰船子类(如‘fishing_boat’, ‘cargo_ship’),这个图能帮你判断各类别样本数是否均衡。严重不均衡会导致模型偏向于样本多的类别。
实操心得:我曾在某个数据集上发现,超过60%的边界框高度小于图片高度的5%。直接训练后模型对小船漏检严重。后来通过复制-粘贴增强(Copy-Paste Augmentation),人工增加了一些小目标样本,才显著提升了召回率。所以,前期花半小时做数据分析,可能节省你后期数天的调参时间。
3. YOLOv5环境搭建与项目初始化避坑指南
工欲善其事,必先利其器。一个干净、可控的环境是成功训练的前提。很多人卡在第一步,问题往往出在环境冲突上。
3.1 基于Conda的Python环境隔离方案
我强烈推荐使用Conda来管理环境,它能完美解决不同项目间Python版本和包版本的冲突。
# 1. 创建并激活一个全新的环境,指定Python版本(YOLOv5推荐3.8+) conda create -n yolov5_boat python=3.8 conda activate yolov5_boat # 2. 安装PyTorch(核心!) # 先去PyTorch官网(https://pytorch.org/get-started/locally/)根据你的CUDA版本选择命令。 # 例如,如果你有CUDA 11.3,则安装命令可能如下: pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 torchaudio==0.12.1 --extra-index-url https://download.pytorch.org/whl/cu113 # 如果没有GPU或CUDA,则安装CPU版本: # pip install torch torchvision torchaudio # 3. 克隆YOLOv5官方仓库(建议使用稳定版本,如v6.1) git clone -b v6.1 https://github.com/ultralytics/yolov5.git cd yolov5 # 4. 安装YOLOv5的依赖包 pip install -r requirements.txt注意事项:
requirements.txt里的opencv-python有时会安装很慢或出错。可以尝试使用国内镜像源,或者先单独安装一个较小的opencv-python-headless。另外,确保你的pip已经升级到最新版。
3.2 数据集配置文件(YAML)的编写艺术
YOLOv5通过一个.yaml文件来定义数据集路径和类别信息。这是连接你的数据和模型的桥梁,必须准确无误。
在yolov5/data/目录下,创建一个名为boat.yaml的文件,内容如下:
# boat.yaml # 训练和验证图像的路径(相对路径或绝对路径) train: ../boat_dataset/images/train/ val: ../boat_dataset/images/val/ # 类别数量 nc: 1 # 假设我们的数据集只有‘ship’这一个类别。如果是多类别,改为相应数字。 # 类别名称列表 names: ['ship'] # 如果多类别,例如:['fishing_boat', 'cargo_ship', 'yacht'] # 可选:下载地址/说明(对于开源数据集) # download: https://your-dataset-url.com关键点解析:
train和val路径:支持绝对路径和相对路径。相对路径是相对于你运行训练命令时所在的目录(通常是yolov5/)。我更喜欢用相对路径,便于项目迁移。nc(number of classes):必须与你的标签文件里最大的类别ID对应(类别ID从0开始)。如果你的标签里类别ID是0,1,2,那么nc就是3。names:列表顺序必须与类别ID严格对应。即names[0]对应ID为0的类别。
验证配置是否正确:在yolov5/目录下,运行一个简单的测试命令,加载几张图片看看:
python train.py --data data/boat.yaml --weights yolov5s.pt --epochs 1 --img 640 --batch-size 2这个命令会尝试训练1个epoch,如果数据路径或格式有误,通常会在开始时报错。用--epochs 1和很小的--batch-size是为了快速失败,避免浪费时间去下载预训练权重或进行长时间的错误训练。
4. 模型训练:超参数调优与训练策略详解
环境就绪,数据备好,终于来到核心的训练环节。YOLOv5的训练命令看似简单,但背后每个参数都大有乾坤。
4.1 训练命令全参数解读与基准实验
一个完整的训练命令示例如下:
python train.py \ --weights yolov5s.pt \ # 初始权重(使用预训练模型迁移学习) --data data/boat.yaml \ # 数据集配置文件 --epochs 100 \ # 训练总轮数 --img-size 640 \ # 输入图像尺寸(长边缩放到此,短边按比例缩放) --batch-size 16 \ # 批次大小(根据GPU内存调整) --device 0 \ # 使用GPU 0,如果是CPU则用 --device cpu --workers 8 \ # 数据加载的线程数(通常设为CPU核心数) --name boat_exp1 \ # 本次实验的名称,用于保存结果 --exist-ok \ # 允许覆盖同名实验目录 --hyp data/hyps/hyp.scratch-low.yaml \ # 使用自定义的超参数文件 --seed 42 # 固定随机种子,确保实验可复现参数选择背后的逻辑:
--weights yolov5s.pt:从YOLOv5s的预训练权重开始。这是强烈推荐的做法,尤其是在数据集不大的情况下(比如几千张图)。预训练权重在COCO等大型通用数据集上学到的通用特征(边缘、纹理、形状),能极大地加速你的模型在特定领域(舰船)的收敛,并提升最终性能。yolov5s是“小”模型,速度快,适合快速原型验证。如果对精度要求高且算力充足,可以尝试yolov5m或yolov5l。--img-size 640:YOLOv5默认的输入尺寸。更大的尺寸(如1280)能检测到更小的目标,但会显著增加计算量和内存消耗,并可能降低训练速度。对于舰船检测,如果数据集中有很多远距离的小船,可以考虑尝试--img-size 1280,但需要同步调整batch-size。--batch-size 16:这是一个需要权衡的参数。较大的批次能使梯度估计更稳定,可能有助于收敛。但受限于GPU显存(VRAM)。你可以从16开始尝试,如果出现“CUDA out of memory”错误,逐步降低到8、4。同时,batch-size改变后,学习率通常也需要按线性规则调整。--hyp:超参数文件。YOLOv5内置了几个(hyp.scratch-low.yaml,hyp.scratch-high.yaml等)。scratch-low学习率较低,更适合微调(Fine-tuning);scratch-high学习率较高,更适合从头训练(Scratch)。对于舰船检测,我们使用预训练权重,所以通常从hyp.scratch-low.yaml开始。
4.2 学习率调度与早停策略实战配置
超参数文件中,学习率(lr0)和权重衰减(weight_decay)是最关键的。但手动调参效率低,YOLOv5集成了自动学习率调整和早停功能。
在hyp.scratch-low.yaml基础上修改:
# hyp.boat.yaml (自定义) lr0: 0.01 # 初始学习率 (对于微调,可以从0.001或0.0005开始尝试) lrf: 0.01 # 最终学习率因子 = lr0 * lrf (余弦退火) momentum: 0.937 weight_decay: 0.0005 warmup_epochs: 3.0 warmup_momentum: 0.8 warmup_bias_lr: 0.1 ...然后在训练命令中指定--hyp data/hyps/hyp.boat.yaml。
启用早停(Early Stopping): YOLOv5的train.py支持早停,但默认不开启。早停可以防止模型在验证集上性能不再提升时继续训练导致的过拟合。
python train.py ... --patience 50--patience 50表示如果验证集指标在连续50个epoch内没有提升,则停止训练。这个值可以根据总epoch数设置,一般设为总epoch数的1/3到1/2。
监控训练过程: 训练开始后,YOLOv5会在runs/train/boat_exp1/目录下生成大量有用的文件:
results.png:核心监控图表,包含训练损失、验证损失、精度(Precision)、召回率(Recall)、mAP@0.5、mAP@0.5:0.95随epoch的变化曲线。这是你判断训练状态是否健康的“仪表盘”。weights/:保存了best.pt(验证集上表现最好的权重)和last.pt(最后一个epoch的权重)。部署时务必使用best.pt。confusion_matrix.png:混淆矩阵,对于多类别任务非常有用,可以查看各类别间的误检情况。
实操心得:训练初期,重点关注
train/box_loss和val/box_loss是否在稳步下降。如果val_loss在几个epoch后开始上升,而train_loss持续下降,这是典型的过拟合信号。此时应该考虑:1. 增加数据增强强度;2. 增加正则化(如DropOut,但YOLO中不常用);3. 提前停止训练。对于舰船数据,我常发现增加随机旋转(degrees)和随机透视(perspective)增强,对提升模型在倾斜、侧视船只上的鲁棒性很有效。
5. 模型评估与性能优化深入解析
训练完成后,我们得到了一个模型权重文件(best.pt)。但这远不是终点,我们需要客观地评估它的性能,并找到优化方向。
5.1 使用验证集进行综合评估
YOLOv5提供了专门的验证脚本,它会加载你训练好的模型,在验证集上跑一遍,生成详细的评估报告。
python val.py \ --weights runs/train/boat_exp1/weights/best.pt \ --data data/boat.yaml \ --img 640 \ --batch-size 32 \ --task val \ --name boat_val \ --save-txt \ # 保存预测的标签文件,用于后续分析 --save-conf \ # 保存预测的置信度 --save-json \ # 保存COCO格式的JSON结果文件 --exist-ok运行后,你会在终端看到类似下面的输出,并在runs/val/boat_val/目录下生成结果:
Class Images Labels Precision Recall mAP@.5 mAP@.5:.95 all 500 2154 0.945 0.892 0.932 0.681 ship 500 2154 0.945 0.892 0.932 0.681指标解读:
- Precision(精度):模型预测为舰船的框里,有多少是真正的舰船。高精度意味着误报(把非船物体认成船)少。
- Recall(召回率):数据集中所有的真实舰船,有多少被模型找出来了。高召回率意味着漏检少。
- mAP@0.5:在IoU(交并比)阈值为0.5时的平均精度均值。这是目标检测最常用的核心指标,值越高越好。0.932是一个非常不错的结果。
- mAP@0.5:0.95:在IoU阈值从0.5到0.95(步长0.05)区间内,计算的平均mAP。这个指标更严格,因为它要求预测框与真实框的重合度更高。0.681说明模型对边界框的定位精度还有提升空间。
5.2 可视化分析与错误排查
数字指标是抽象的,我们需要可视化结果来发现具体问题。YOLOv5的验证脚本会生成预测结果图。
分析runs/val/boat_val/目录下的图片:
confusion_matrix.png:如果有多类别,看混淆矩阵。labels.jpg和labels_correlogram.jpg:标签分布图。*_batch*pred.jpg:随机一些批次的预测结果可视化。
重点观察:
- 漏检(False Negative):图片中有船,但模型没框出来。这通常发生在目标太小、太模糊、或被严重遮挡时。解决方案:在数据增强中增加小目标增强(如随机复制粘贴小目标),或使用更专注于小目标检测的模型变体(如YOLOv5的
P6模型,输入尺寸1280)。 - 误检(False Positive):模型把非船物体(如海浪尖、云朵形状、海岛)框成了船。解决方案:增加包含这些负样本(没有船)的图片到数据集中,或者在数据增强中增加更多样的背景干扰。
- 定位不准(Low IoU):框出来了,但框的位置或大小不准。这会影响mAP@0.5:0.95。解决方案:检查标注质量,有时是标注本身就不够精确。可以尝试调整损失函数中定位损失的权重(在
hyp.yaml中修改box_loss增益),但需谨慎。
5.3 针对舰船场景的专项优化技巧
基于上述分析,我们可以进行针对性优化:
1. 针对小目标漏检的优化:
- 修改模型结构:YOLOv5默认的检测头(Head)有三个尺度(P3, P4, P5),分别对应大、中、小目标。如果你的小目标特别多,可以尝试使用官方提供的
--model yolov5s6.yaml(即YOLOv5s-P6模型),它增加了一个P6/64尺度,专门用于检测更小的目标,但模型会变大变慢。 - 调整Anchor Boxes:YOLOv5默认使用K-means聚类在COCO数据集上生成的Anchor。对于舰船这种长宽比可能比较特殊的物体(比如货轮很长),可以在自己的数据集上重新聚类Anchor。使用YOLOv5提供的脚本:
运行后,它会评估当前Anchor与数据集的匹配度,并给出建议的新Anchor尺寸。你可以将这些新尺寸更新到模型配置文件(如python utils/autoanchor.py --data data/boat.yamlmodels/yolov5s.yaml)中,然后重新训练。
2. 针对复杂背景误检的优化:
- 数据增强策略调整:在
hyp.yaml文件中,增强参数集中在augment:部分。对于舰船,可以尝试:augment: hsv_h: 0.015 # 色调增强,模拟不同光照 hsv_s: 0.7 # 饱和度增强,让海水颜色更多变 hsv_v: 0.4 # 明度增强 degrees: 10.0 # 旋转角度,船只有各种角度 translate: 0.1 # 平移 scale: 0.5 # 缩放 shear: 5.0 # 剪切变换 perspective: 0.0005 # 透视变换,模拟不同视角 flipud: 0.0 # 上下翻转(对于舰船,上下翻转通常不合理,设为0) fliplr: 0.5 # 左右翻转(合理) mosaic: 1.0 # Mosaic增强,非常有效,但可能增加小目标 mixup: 0.0 # MixUp增强,可尝试0.1-0.2,但需谨慎,可能模糊特征注意:
mosaic增强会把四张图拼成一张,能极大地丰富背景,并让模型学习在不同位置、尺度上检测目标。但对于小目标,拼图后可能变得更小,如果小目标问题严重,可以尝试降低mosaic概率或关闭它(设为0)。
3. 模型集成与测试时增强(TTA)如果单个模型性能达到瓶颈,可以考虑:
- 模型集成:训练多个不同初始化或不同数据子集的模型,在推理时取它们的预测平均值或加权平均值。
- 测试时增强:在推理时,对输入图像进行多种变换(如翻转、缩放),将多个预测结果合并。YOLOv5的
detect.py或val.py支持--augment参数来启用TTA,通常会提升mAP,但会显著增加推理时间。python val.py --weights best.pt --data boat.yaml --augment
6. 模型部署与应用:从PyTorch到实际推理
模型训练评估好了,最终要落地应用。这里涉及模型导出、优化和编写推理代码。
6.1 模型导出为部署格式
YOLOv5的export.py脚本支持将PyTorch模型(.pt)导出为多种格式,以适应不同的部署环境。
导出为ONNX格式(推荐,通用性强):
python export.py \ --weights runs/train/boat_exp1/weights/best.pt \ --img 640 \ --batch 1 \ # 指定批处理大小,部署时常为1 --device cpu \ # 在CPU上执行导出 --include onnx \ # 导出为ONNX --opset 12 \ # ONNX算子集版本 --simplify # 启用ONNX简化,优化模型结构导出的best.onnx文件可以被OpenCV DNN、TensorRT、ONNX Runtime等多种推理引擎加载。
导出为TensorRT引擎(追求极致GPU速度):
python export.py \ --weights best.pt \ --img 640 640 \ # 输入尺寸(高,宽) --batch 1 \ --device 0 \ # 在GPU 0上导出 --include engine \ # 直接导出为TensorRT引擎文件(需要提前安装TensorRT) --half # 使用FP16精度,进一步提速减存注意:直接导出TensorRT引擎需要复杂的本地环境配置。更常见的做法是先导出ONNX,再用TensorRT的trtexec工具或Python API进行转换和优化。
6.2 编写Python推理脚本
这里给出一个使用导出的ONNX模型,通过ONNX Runtime进行推理的完整示例脚本:
import cv2 import numpy as np import onnxruntime as ort from pathlib import Path import time class YOLOv5ONNXInference: def __init__(self, onnx_model_path, conf_thresh=0.25, iou_thresh=0.45): """ 初始化ONNX推理会话 Args: onnx_model_path: .onnx模型文件路径 conf_thresh: 置信度阈值 iou_thresh: NMS的IoU阈值 """ self.conf_threshold = conf_thresh self.iou_threshold = iou_thresh # 创建ONNX Runtime会话 providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] if ort.get_device() == 'GPU' else ['CPUExecutionProvider'] self.session = ort.InferenceSession(onnx_model_path, providers=providers) # 获取模型输入输出信息 self.model_inputs = self.session.get_inputs() self.model_outputs = self.session.get_outputs() self.input_name = self.model_inputs[0].name self.input_shape = self.model_inputs[0].shape # 通常是(1, 3, 640, 640) self.input_height, self.input_width = self.input_shape[2], self.input_shape[3] # 类别名(需要根据你的数据集修改) self.class_names = ['ship'] def preprocess(self, image): """将输入图像预处理为模型需要的格式""" # 保持宽高比resize,并在边缘填充灰色 h, w = image.shape[:2] scale = min(self.input_height / h, self.input_width / w) new_h, new_w = int(h * scale), int(w * scale) resized_img = cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_LINEAR) # 创建画布并填充 canvas = np.full((self.input_height, self.input_width, 3), 114, dtype=np.uint8) top = (self.input_height - new_h) // 2 left = (self.input_width - new_w) // 2 canvas[top:top+new_h, left:left+new_w, :] = resized_img # 转换通道顺序 HWC -> CHW, BGR -> RGB, 归一化 canvas = canvas.transpose(2, 0, 1) # HWC to CHW canvas = canvas[::-1, :, :] # BGR to RGB canvas = canvas.astype(np.float32) / 255.0 # 归一化到 [0,1] # 添加批次维度并返回 blob = np.expand_dims(canvas, axis=0) return blob, (scale, left, top), (h, w) def postprocess(self, outputs, preprocess_info, orig_shape): """将模型输出后处理为检测框""" scale, pad_left, pad_top = preprocess_info orig_h, orig_w = orig_shape # outputs是一个列表,取第一个元素(不同版本输出可能不同) predictions = np.squeeze(outputs[0]) # 形状: (num_boxes, 85) # 过滤低置信度框 conf_mask = predictions[:, 4] > self.conf_threshold predictions = predictions[conf_mask] if len(predictions) == 0: return [] # 分离框坐标和类别分数 boxes = predictions[:, :4] scores = predictions[:, 4] class_probs = predictions[:, 5:] class_ids = np.argmax(class_probs, axis=1) max_class_scores = np.max(class_probs, axis=1) # 综合置信度 = 目标置信度 * 类别置信度 final_scores = scores * max_class_scores # 将框的格式从 (cx, cy, w, h) 转换为 (x1, y1, x2, y2) boxes[:, 0] = (boxes[:, 0] - boxes[:, 2] / 2) # x1 boxes[:, 1] = (boxes[:, 1] - boxes[:, 3] / 2) # y1 boxes[:, 2] = (boxes[:, 0] + boxes[:, 2]) # x2 boxes[:, 3] = (boxes[:, 1] + boxes[:, 3]) # y2 # 将框的坐标映射回原始图像尺寸 boxes[:, [0, 2]] = (boxes[:, [0, 2]] - pad_left) / scale boxes[:, [1, 3]] = (boxes[:, [1, 3]] - pad_top) / scale # 应用非极大值抑制 (NMS) keep_indices = self.nms(boxes, final_scores) results = [] for idx in keep_indices: x1, y1, x2, y2 = boxes[idx].astype(int) # 确保框在图像范围内 x1, y1 = max(0, x1), max(0, y1) x2, y2 = min(orig_w, x2), min(orig_h, y2) if x2 <= x1 or y2 <= y1: continue results.append({ 'bbox': [x1, y1, x2, y2], 'confidence': float(final_scores[idx]), 'class_id': int(class_ids[idx]), 'class_name': self.class_names[int(class_ids[idx])] }) return results def nms(self, boxes, scores): """简化的非极大值抑制实现""" x1 = boxes[:, 0] y1 = boxes[:, 1] x2 = boxes[:, 2] y2 = boxes[:, 3] areas = (x2 - x1 + 1) * (y2 - y1 + 1) order = scores.argsort()[::-1] keep = [] while order.size > 0: i = order[0] keep.append(i) xx1 = np.maximum(x1[i], x1[order[1:]]) yy1 = np.maximum(y1[i], y1[order[1:]]) xx2 = np.minimum(x2[i], x2[order[1:]]) yy2 = np.minimum(y2[i], y2[order[1:]]) w = np.maximum(0.0, xx2 - xx1 + 1) h = np.maximum(0.0, yy2 - yy1 + 1) inter = w * h ovr = inter / (areas[i] + areas[order[1:]] - inter) inds = np.where(ovr <= self.iou_threshold)[0] order = order[inds + 1] return keep def infer(self, image_path): """完整的推理流程""" # 读取图像 orig_img = cv2.imread(image_path) if orig_img is None: print(f"Error: Could not read image {image_path}") return [] # 预处理 start_time = time.time() blob, preprocess_info, orig_shape = self.preprocess(orig_img) preprocess_time = time.time() - start_time # 推理 start_time = time.time() outputs = self.session.run(None, {self.input_name: blob}) inference_time = time.time() - start_time # 后处理 start_time = time.time() detections = self.postprocess(outputs, preprocess_info, orig_shape) postprocess_time = time.time() - start_time print(f"Timing - Preprocess: {preprocess_time*1000:.2f}ms, " f"Inference: {inference_time*1000:.2f}ms, " f"Postprocess: {postprocess_time*1000:.2f}ms") return detections, orig_img # 使用示例 if __name__ == "__main__": # 初始化检测器 detector = YOLOv5ONNXInference(onnx_model_path='best.onnx', conf_thresh=0.3) # 对单张图片进行推理 test_image = 'test_ship.jpg' results, image = detector.infer(test_image) # 可视化结果 for det in results: x1, y1, x2, y2 = det['bbox'] label = f"{det['class_name']} {det['confidence']:.2f}" cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2) cv2.putText(image, label, (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2) cv2.imwrite('result.jpg', image) print(f"Detected {len(results)} ships.")这个脚本提供了完整的预处理、推理、后处理流程,并包含了计时功能,方便进行性能分析。你可以将其集成到你的视频流处理、Web服务或桌面应用中。
6.3 部署到边缘设备(如Jetson系列)的注意事项
如果你需要在NVIDIA Jetson这样的边缘设备上部署,流程会有所不同,核心是使用TensorRT来获得最佳性能。
- 环境准备:在Jetson上安装JetPack SDK,它包含了CUDA、cuDNN和TensorRT。
- 模型转换:在Jetson上,使用
trtexec工具将ONNX模型转换为TensorRT引擎。由于Jetson算力有限,通常使用FP16甚至INT8精度进行量化。/usr/src/tensorrt/bin/trtexec --onnx=best.onnx --saveEngine=best_fp16.engine --fp16 --workspace=1024 - 编写TensorRT推理代码:需要使用TensorRT的Python API或C++ API来加载
.engine文件并进行推理。这部分代码比ONNX Runtime复杂,需要处理内存分配、绑定、执行上下文等。
避坑指南:在边缘设备上,内存和算力是瓶颈。务必使用
--img-size 640甚至更小的输入尺寸(如416),并选择yolov5n(纳米)或yolov5s模型。在转换TensorRT引擎时,如果失败,尝试减小--workspace参数的值。实测在Jetson Nano上,YOLOv5s模型在640分辨率下,使用FP16精度,推理速度可以达到10-15 FPS,基本满足实时性要求。
7. 项目总结与未来优化方向
回顾整个从“boat-舰船检测数据集”到可部署模型的过程,其实是一个标准的机器学习项目流水线:数据准备、模型训练、评估优化、部署应用。其中,数据质量是天花板,模型调优是爬坡过程,而部署优化是让成果落地的最后一公里。
我个人在多次类似项目中最深的体会是:不要迷信模型和算法。很多时候,性能瓶颈不在模型本身,而在数据。花时间分析你的数据集,理解标注的噪声分布,设计针对性的数据增强策略,其投资回报率往往远高于无脑尝试更复杂的模型。例如,在这个舰船项目中,如果发现模型对雾天船只漏检严重,那么去收集或合成更多雾天场景的数据,比把YOLOv5换成YOLOv8可能更有效。
这个项目还可以从多个方向进行深化:
- 多类别细分:将单一的“ship”类别细分为“fishing_boat”, “container_ship”, “sailboat”, “warship”等,为更精细的应用提供支持。
- 引入跟踪算法:在视频流上,将检测器与跟踪算法(如DeepSORT, ByteTrack)结合,实现舰船的连续跟踪与轨迹分析。
- 半自动/主动学习:将模型部署到实际场景中,收集模型不确定的或预测错误的困难样本,人工复核后加入训练集,循环迭代,让模型在不断反馈中越变越强。
- 模型轻量化与量化:为了在更廉价的硬件或移动端部署,可以探索知识蒸馏、剪枝、量化(INT8)等技术,在精度损失可控的前提下,大幅提升推理速度。
最后,分享一个调试小技巧:当你对模型性能不满意时,不要只看整体的mAP。把验证集上置信度最高的假阳性(误检)和置信度最高的假阴性(漏检)样本挑出来,做成一个图册反复看。这些是模型“最自信的错误”和“最不该的错过”,它们往往直指数据或模型最根本的缺陷。解决这些问题,模型的性能通常会有立竿见影的提升。
本文还有配套的精品资源,点击获取