拓十年匠心定制 · 商业建站与技术教学双线并行 咨询热线:400-886-1026 service@lmnt.cn
ARTICLE DETAIL

资讯详情

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

YOLO11工地安全帽与反光衣检测:3000张图从标签转换到训练避坑

YOLO11工地安全帽与反光衣检测:3000张图从标签转换到训练避坑

简介:面向智慧工地安全监控场景的目标检测数据集资源,适合需要训练工人安全设备佩戴检测模型的算法工程师与研究人员。数据集包含3000张真实工地图片,覆盖建筑工地、道路施工、室内装修、工人密集作业及设备遮挡等多样场景,标注了Fall-Detected、Hardhat、Mask、Safety Vest、Person等14个类别,并同时提供VOC、COCO、YOLO三种格式标签,可直接用于YOLO等主流算法训练。资源包共1个文件,为PDF格式,整体大小8.59MB;由于图片数据较大,PDF内附数据集基本情况、缩略图、labelimg标注截图及百度网盘获取方式。目前已吸引426人学习浏览。配套资料还提供YOLO11一键训练脚本,支持GPU、CPU、Mac(M芯片)多平台运行,并附博主训练结果日志,便于快速验证与对比模型效果,适合作为智慧工地安全合规检测项目的数据补充与训练参考。

1. 目标检测里的硬骨头:工地工人安全设备佩戴检测,3000张图怎么榨出可用模型

工地工人安全帽和反光衣的佩戴检测,是目标检测里一类非常“吃场景”的任务:框架可以选 YOLO11,模型还是那套 backbone 加检测头,可数据一旦从 COCO 换成工地现场,精度就会掉得让你怀疑代码写错了。真正决定模型能不能在工地上落地的,从来不是算法花样,而是三件事:数据集够不够贴场景、VOC/COCO/YOLO 三种标签格式能不能对齐、训练脚本在 GPU(GPUs)/CPU/Mac 三平台上能不能顺利跑完。这篇笔记就顺着这个标题,把 3000 张图的工地安全设备佩戴检测数据集,从标签格式、YOLO11 一键训练到踩坑排查完整过一遍,给想直接照做的人一条能走通的路。

2. 看懂三种标签格式:VOC、COCO、YOLO 的坐标系统与转换路径

2.1 VOC 的 XML:每个框坐标都存在 bndbox 里,从源头理解标注

标题里写了“对应 VOC/COCO/YOLO 三种格式标签”,意思是同一批图,每张图同时存在三种格式标注。其中 VOC 是很好理解的入门格式,它的标签是 XML 文件,文件名和图片名一一对应。打开一个 VOC 格式的安全帽标注文件,看到的不是像素图,而是一段树状结构:根节点 annotation 下面有 filename、size,再往下是多个 object 节点,每个 object 代表画面里的一个目标。

object 节点里最关键的是 name 和 bndbox 两个子字段。name 是类别名,比如 helmet、vest、person;bndbox 里是四个整数,xmin、ymin、xmax、ymax,分别代表框的左上角和右下角在图像像素坐标系里的位置。这个坐标系的原点在图片左上角,x 向右增大,y 向下增大,单位是像素,不是归一化值。很多新手第一次写 VOC 转 YOLO 脚本时,最容易犯的错就是把这四个像素值直接当坐标用,忘了 YOLO 要的是归一化后的中心点坐标。

判断一个 VOC 数据标注是否合格,我一般会做三件事:第一,检查 bndbox 的 xmin 是否严格小于 xmax,ymin 是否严格小于 ymax,出现等于号说明标签被画成了点或线;第二,检查四个值是否都在图片宽度和高度范围内,越界说明标注软件曾对图片做过缩放但没同步标签;第三,检查同一个 XML 里有没有重叠度过高的两个同类框,这往往是标注员手误造成的重复框。数据集的干净程度,直接决定后面一键训练能不能顺利跑起来。

2.2 COCO 的 JSON 与 YOLO 的 txt:同一张图的三份“身份证”

COCO 格式的标签不是每张图一个文件,而是把整个数据集的标注塞进一个大的 JSON 文件里。JSON 里有三个关键数组:images 记录每张图的 id、file_name、width、height;annotations 记录每个框的 image_id、category_id、bbox;categories 记录类别 id 和名称的映射。COCO 的 bbox 写的是 [x, y, width, height],也就是左上角横坐标、左上角纵坐标、框宽度、框高度,单位和 VOC 一样都是像素。

YOLO 格式则完全反过来。YOLO 训练要的是每张图片同名的一个 txt 文件,txt 里每行对应一个框,格式是“类别id 中心点x 中心点y 宽度 高度”,而且这四个数值全部是归一化后的,范围在 0 到 1 之间。归一化的计算方法很简单:中心点x 等于 (xmin + xmax) / 2 再除以图片宽度,宽度等于 (xmax - xmin) 除以图片宽度,y 方向同理。这里有个容易踩坑的点:归一化值可能非常小,比如一个远处工人头顶的安全帽框,宽度可能只有图片宽度的 0.01 甚至更低,保存 txt 时如果只保留两位小数,这个框就变成 0.00 了,训练时会被当作无效框直接丢弃。

三种格式的差别,一句话概括:VOC 和 COCO 都用像素坐标,区别只在记录方式;YOLO 用归一化坐标,而且按“中心点加宽高”记录。标题里的数据集之所以要同时给三种格式,好处在迁移性:COCO 格式方便你在 mmdetection 或 Detectron2 上做对比实验,VOC 格式方便传统工具链读取,YOLO 格式直接喂给 YOLO11 训练,互不依赖,省得临时转换出问题。

2.3 标签转换与一致性校验脚本:从 XML 到 txt,先对齐再开训

拿到三种格式的标签后,第一件事不是训练,而是校验。最常见的工作流是:先用 VOC 做直观检查,确认框没问题;再转出 YOLO txt 作为训练输入;最后把 COCO 当作交叉验证的“第三把尺子”。下面这段是把 VOC XML 转成 YOLO txt 的参考脚本,适合 3000 张图这种中小规模数据集,一条命令就能跑完:

import os import xml.etree.ElementTree as ET def voc_xml_to_yolo_txt(xml_path, img_w, img_h, class_map): """ 将单个VOC XML转为YOLO txt 参数: xml_path - XML文件路径; img_w/img_h - 图片宽高; class_map - 类别名到id的映射 """ tree = ET.parse(xml_path) root = tree.getroot() lines = [] for obj in root.iter('object'): name = obj.find('name').text if name not in class_map: continue # 跳过未在映射中出现的类别,防止类别id错位 cid = class_map[name] bndbox = obj.find('bndbox') xmin = float(bndbox.find('xmin').text) ymin = float(bndbox.find('ymin').text) xmax = float(bndbox.find('xmax').text) ymax = float(bndbox.find('ymax').text) # 归一化:中心点坐标除以宽高,得到0-1之间的值 cx = ((xmin + xmax) / 2) / img_w cy = ((ymin + ymax) / 2) / img_h w = (xmax - xmin) / img_w h = (ymax - ymin) / img_h # 兜底:裁剪越界值,防止坐标小于0或大于1 cx = min(max(cx, 0.0), 1.0) cy = min(max(cy, 0.0), 1.0) w = min(max(w, 0.0), 1.0) h = min(max(h, 0.0), 1.0) lines.append(f"{cid} {cx:.6f} {cy:.6f} {w:.6f} {h:.6f}") return "\n".join(lines) # 使用示例:遍历数据集目录,逐张转换 class_map = {"person": 0, "helmet": 1, "vest": 2} dataset_dir = "helmet_dataset" for img_name in os.listdir(os.path.join(dataset_dir, "images")): base = os.path.splitext(img_name)[0] xml_path = os.path.join(dataset_dir, "voc", base + ".xml") txt_path = os.path.join(dataset_dir, "yolo", base + ".txt") if not os.path.exists(xml_path): continue # 注意:这里的宽高应从原图读取,不要用XML里的数值,防止XML存的是缩略图尺寸 img_w, img_h = 1920, 1080 result = voc_xml_to_yolo_txt(xml_path, img_w, img_h, class_map) with open(txt_path, "w", encoding="utf-8") as f: f.write(result)

脚本逻辑不复杂,但有两个细节容易被忽略。第一,XML 根节点里的 size 字段虽然也记录宽度和高度,但有些标注工具在导出时会写错,我遇到过 XML 里写 1920×1080、实际原图是 3840×2160 的情况,所以稳妥做法是用 OpenCV 读一遍原图拿真实宽高。第二,类别映射一定要固定,person 是 0、helmet 是 1、vest 是 2,如果排序不同,训练出来的类别名就会错位,看起来模型在测安全帽,实际输出的置信度是反光衣的。这个坑在三种格式都提供的项目里尤其常见,因为 COCO 的 category_id 和 YOLO 的类别 id 是两套数字,必须人工核对。

转换完之后,写一个十行以内的校验脚本,检查每张图的 YOLO txt 内容是否干净:文件中是不是有空行、每一行是不是刚好五个数值、有没有任何值不在 0 到 1 区间、空 txt 是否对应一张没有标注的原图。这些检查跑完,数据才算真正可以进训练管线。

3. YOLO11 一键训练脚本拆解:GPU(GPUs)/CPU/Mac 三平台跑通路径

3.1 环境配置的差异:CUDA、CPU 与 Mac 的 MPS 各走一条路

标题里的“一键训练脚本”如果真能做到一键,那一定把平台差异封装在了底层。YOLO11 基于 Ultralytics 框架,官方包安装很简单,一行pip install ultralytics就够,但背后的 PyTorch 版本要按平台选。GPU(GPUs) 机器要手动装 CUDA 版 PyTorch,常见的做法是去 PyTorch 官网按 CUDA 版本生成安装命令;纯 CPU 机器直接装默认版本即可;Mac 的 Apple Silicon 芯片则要用支持 MPS 后端的 PyTorch,安装后还要确认 torch 能识别到 MPS。

在写训练脚本之前,我会先跑一个三平台通用的设备自检,把当前环境到底能用什么设备炸出来,避免训练到一半才发现 device 参数写错:

import torch def detect_device(): """ 自动检测当前平台可用的训练设备 优先级: NVIDIA GPU > Mac MPS > CPU """ if torch.cuda.is_available(): device = "cuda" count = torch.cuda.device_count() print(f"检测到GPU设备{count}张: {torch.cuda.get_device_name(0)}") elif torch.backends.mps.is_available(): device = "mps" print("检测到Apple MPS设备") else: device = "cpu" print("未发现加速设备,使用CPU训练") return device device = detect_device()

这段脚本的价值在于“黑匣子”透明化。很多一键训练脚本翻车,不是训练代码错了,而是环境不对:比如 Windows 上装了 CPU 版 PyTorch,代码里写 device="cuda",启动后直接报错找不到显卡驱动;比如 Mac 上安装的是 Intel 版 Python,MPS 根本不生效,训练速度比 CPU 还慢。把这些检测逻辑放进训练脚本最前面,比任何文档说明都管用。

GPU 多卡场景还要注意:Ultralytics 的 device 参数可以传 "0,1" 这种字符串来启用多卡训练,但数据集很小的时候多卡收益不明显,3000 张图单卡 16 的 batch 也只要几百个 step,多卡反而因为通信开销拖慢速度。我一般只在单卡显存小于 8G 时才考虑多卡或梯度累积。

3.2 一键训练脚本的四个步骤:校验、划分、训练、验证

所谓一键训练,是把“校验数据 + 划分数据集 + 设定超参 + 执行训练 + 跑验证”串成一条流水线。下面是一个可参考的最小训练入口,我习惯把它写成train.py直接运行:

import os import yaml from ultralytics import YOLO # 第一步:数据校验,检查图片和标签数量是否一致 img_dir = "helmet_dataset/images" label_dir = "helmet_dataset/yolo" imgs = {os.path.splitext(f)[0] for f in os.listdir(img_dir)} labels = {os.path.splitext(f)[0] for f in os.listdir(label_dir)} missing = imgs - labels if missing: raise RuntimeError(f"有图片缺失标签: {list(missing)[:5]}") # 第二步:写YOLO训练用的YAML配置 data_cfg = { "path": "helmet_dataset", "train": "images/train", "val": "images/val", "nc": 3, "names": ["person", "helmet", "vest"], } with open("helmet_dataset.yaml", "w", encoding="utf-8") as f: yaml.dump(data_cfg, f, allow_unicode=True) # 第三步:选择模型并开始训练 device = detect_device() model = YOLO("yolo11n.pt") # 如果没有本地权重文件,会自动联网下载 results = model.train( data="helmet_dataset.yaml", epochs=120, imgsz=640, batch=16, device=device, patience=30, project="runs/helmet", name="exp", lr0=0.01, ) # 第四步:用best.pt在验证集上出指标 best_model = YOLO("runs/helmet/exp/weights/best.pt") metrics = best_model.val(data="helmet_dataset.yaml") print(f"mAP@0.5: {metrics.box.map50:.4f}, mAP@0.5:0.95: {metrics.box.map:.4f}")

这段脚本里的关键参数值得逐一说清。epochs=120是训练轮数,3000 张图配数据增强,一般到 80 到 120 轮就会收敛,轮数再多就开始过拟合工地背景;patience=30是早停轮数,连续 30 轮验证集 mAP 不提升就自动停止,省时间;imgsz=640是训练输入分辨率,工地安全帽在画面里通常是小目标,如果原图是 1080p 及以上,可以后续提高到 800,但训练显存和耗时也会跟着涨;lr0=0.01是初始学习率,YOLO11 默认值对这个规模的数据集基本不用动。

另外,首次运行会自动下载yolo11n.pt预训练权重,如果现场机器不能联网,就得提前在其他机器上下载好后拷贝到本地,如果你的部署场景是工地项目部那种网络不发达的办公室,这一步必须在脚本里做文件存在性判断,否则一键训练会卡在下载阶段。

3.3 训练产物与目标检测评价指标:怎么判断这个模型靠谱

训练完成后,runs/helmet/exp/目录下会生成 weights、混淆矩阵、PR 曲线、验证结果图等一堆东西。最重要两个权重文件是best.pt和last.pt,前者是验证集 mAP 最高的权重,后者是最后一轮结束时的权重。项目交付时我只会用 best.pt,但保留 last.pt 作为续训的后悔药。

评价指标看两个数就够起步。第一个是 mAP@0.5,指 IoU 阈值 0.5 时所有类别的平均 AP,它反映的是框“大概准不准”;第二个是 mAP@0.5:0.95,是不同 IoU 阈值下的平均,它更严格,反映框“贴不贴边”。对工地安全设备检测来说,mAP@0.5 能达到 0.85 以上才算能用,mAP@0.5:0.95 如果能到 0.5 以上,说明框质量不错。还要单独看 recall,反光衣检测最怕的就是漏检,宁可多框一个背景也不能漏掉一个没穿反光衣的工人,所以 recall 比 precision 优先保证。

4. 3000 张图怎么喂出可用模型:数据逻辑与训练超参调整

4.1 数据质量优先于数量:类别均衡、场景多样、负样本一个都不能少

3000 张图在目标检测里属于中小规模,按每张图 2 到 3 个工人来算,总标注框数大概在 6000 到 9000 个之间,不算多。但工地安全设备佩戴检测有个特殊性:核心目标不是检测工人,而是检测“没戴安全帽、没穿反光衣”的隐患。这要求数据里必须有一定比例的负样本,也就是工人没戴安全帽、没穿反光衣的图。如果数据集中所有人都是规范佩戴状态,模型学到的是“有人就有安全帽”,到了真实工地遇到没戴的情况就直接漏检,这种模型在验收时会被当场打回去。

开训之前建议做一个类别分布统计,看每个类别出现在多少张图里,而不是看框数量。反光衣往往是最容易被忽视的类别,因为很多工地照片里反光衣反光强烈,颜色和背景融为一体,标注员容易漏标。如果发现某类出现在不足 20% 的图片里,后续训练它即使有 loss 权重也很难学出来,要么补充数据,要么在增强阶段针对该类做专门处理。

4.2 训练关键参数:imgsz、epochs、batch 与数据增强的取舍

YOLO11 训练参数的默认值是基于 COCO 这种大而全的数据集调出来的,套到工地场景必须改几个。第一个是 imgsz。我见过有人直接用默认的 640 训,结果 1080p 原图里远处一个工人只有 40×80 像素,安全帽只有 10×10 像素,640 分辨率下网络根本看不清,mAP 自然上不去。对这类小目标场景,通常的做法是先降到 480 快速跑通流程,再用 imgsz=800 正式训练,或者直接切割原图为 640 瓦片训练、推理时再拼回全图。

第二个是 batch。显存够的情况下 batch 越大越稳,但 3000 张图没必要追求大 batch,16 到 32 已经足够。显存不足时优先降 batch,不要先降 imgsz,因为分辨率对小目标的影响比 batch 更大。第三个是数据增强开关。YOLO11 默认开 mosaic 增强,四张图拼一张训练,能显著提升复杂场景的鲁棒性,但对小目标不太友好,因为拼接后目标变得更小;如果发现小目标漏检严重,可以试试关掉 mosaic,改成只做随机翻转、色彩抖动和仿射变换。

下面是一份针对工地数据的增强配置参考,可以直接塞进训练 yaml 里:

# helmet_dataset_aug.yaml path: helmet_dataset train: images/train val: images/val nc: 3 names: 0: person 1: helmet 2: vest # 增强项:关闭部分mosaic,保留色彩与几何增强 mosaic: 0.5 mixup: 0.2 fliplr: 0.5 hsv_h: 0.02 hsv_s: 0.6 hsv_v: 0.4 scale: 0.5

mosaic 设成 0.5 表示一半的 epoch 用拼接、一半正常训练,这样既保留增强多样性,又不至于让小目标被切烂。hsv 三个参数控制色相、饱和度、亮度的扰动范围,安全帽大多是红黄橙色,这些颜色在工地光照下变化很大,适当调大饱和度扰动能让模型更抗光照变化。scale 控制随机缩放范围,0.5 的含义是允许目标被缩放到原大小的 0.5 到 1.5 倍之间,这个值对远近尺度的泛化很有帮助。

4.3 CPU 和 Mac 的省时路径:小模型、冻结权重、分级训练

如果手头只有 CPU 或 Mac,就不要傻乎乎地按 GPU 的参数来。CPU 训练 120 轮、imgsz=640 可能要跑十几个小时甚至更久,而 Mac 的 MPS 虽然比 CPU 快不少,但遇到某些算子在 float16 下会产生 NaN 梯度,训练会突然崩掉。我在这两个平台上的实际做法是三步走。

第一步,换小模型。YOLO11 分为 n、s、m、l、x 五档,CPU 和 Mac 优先用 n 档,它的参数量只有最小档,训练速度快好几倍,精度损失对安全帽这种大目标来说在可接受范围内。第二步,冻结 backbone 训练前 30 轮。Ultralytics 的 train 方法支持 freeze 参数,比如freeze=10表示冻结前 10 层,先让检测头学会基础框定位,再解冻全部层精调,这样既省算力又能避免刚开局就震荡。第三步,如果 Mac 上遇到 loss 变成 NaN,检查是不是混合精度的问题,直接在 train 参数里加amp=False,用单精度训练,速度稍慢但过程不翻车。

5. 避坑:工地数据集从标签到训练的 5 个常见翻车点

5.1 图片被 EXIF 旋转导致标签错位

现象:训练 loss 一直不降,可视化训练样本时发现安全帽的框和实际位置差了 90 度,有些图横着、有些竖着,标注框看起来是旋转过的。

原因:手机或部分相机拍摄的照片内部带有 EXIF 旋转标记,OpenCV 默认不处理这个标记,按像素原始排列读图;而标注工具用 PIL 读图时自动应用了旋转。结果就是标注工具的坐标系和训练读图的坐标系不一致,标签全部错位。

解决:在数据预处理阶段把所有图片重写成统一方向,并清除 EXIF 信息。用 Python 的 PIL 读图后调用ImageOps.exif_transpose,再另存为新 jpg,转存完成后用校验脚本重新检查同样的图片在 OpenCV 下读出的宽高是否和之前一致。这个坑在三种格式数据集里尤其隐蔽,因为 XML 标签里存的是旋转后的坐标,转出来的 YOLO txt 也跟着错。

5.2 验证集泄漏:训练和验证来自同一工地场景

现象:训练完看指标,mAP@0.5 高达 0.95,现场一把子数据一测,掉到 0.6 以下,差别大到没法解释。

原因:随机划分训练集和验证集时,同一个工地同一个安全通道的连续帧图片被踢进了两边,模型在训练时见过几乎一样的背景,验证时自然“考得好”。但这种记忆对没见过的工地场景完全不适用。

解决:不要按文件序号随机划分数据集,按场景“工地段落”划分。常见做法是:把同一时间、同一位置采集的图片依次放在同一个场景文件夹下,用文件夹名做分层采样,保证一个场景内的图片只出现在训练集或验证集其中一边。划分完后打印验证集文件名列表,人工扫一眼有没有看起来像同一个场景的。

5.3 反光衣检测不出来:类别不平衡与混淆矩阵定位

现象:person 和 helmet 的 AP 都不错,vest 的 AP 接近 0,验证集图片里穿反光衣的工人一个都没框出来。

原因:反光衣目标在数据集中占比太低,加上反光材料在强光下变成白色亮斑,和浅色背景几乎融为一体,模型没有足够的正样本学习区分。

解决:先用混淆矩阵确认是“漏检”还是“错检”。如果是漏检,反光衣类别数量再不够,不要直接加 loss 权重,优先做的事情是收集更多反光衣正样本,哪怕是裁剪图片局部拼贴,也比硬调权重有用。如果数据实在凑不齐,可以改两阶段思路:先用通用 person 检测把所有人找出来,再用第二级分类判断是否穿反光衣,用 3000 张图里的有限样本反而更容易训。

5.4 训练中途退出:显存溢出与 last.pt 续训

现象:训练到第 47 轮时显存溢出报错退出,或者断网、断电导致进程被杀,重新开始又得从第 0 轮跑。

原因:batch 或 imgsz 超过显存承载上限,或训练环境本身不稳定,进程没有设置自动续训。

解决:报错先看日志里是不是 CUDA out of memory,是就降 batch 到 8 或 4,再不行降 imgsz 到 480。中途退出后用 Ultralytics 的续训功能,直接加载last.pt再执行一次同样的 train 命令,把 resume=True 加上,它会自动从保存的 epoch 继续,不需要重头跑。注意续训时不能再传不同的 data 路径,否则会重新初始化数据管道,等于白续。

5.5 三种格式数量对不上:XML、JSON、txt 文件互相矛盾

现象:训练时发现一部分图片没有对应 txt,排查发现 VOC 的 XML 完整、COCO 的 JSON 也完整,但 YOLO 的 txt 缺了几十张,或者某几张图的 txt 里框的数量和 XML 不一致。

原因:三种格式不是同一时间导出的,中间某一次转换脚本漏了文件,或者一张图里有目标但被转换代码continue跳过了。

解决:写一个全量一致性脚本,统计图片、XML、txt 三者的 base 文件名集合是否完全一致,再逐图比对 XML 中 object 的数量和对应 txt 的行数。这一步必须放在训练之前完成,否则训练时不会报错,但指标会莫名其妙地少几个点,查起来非常棘手。我就吃过这个亏,最后一行行对标签才发现少了 200 多个框,当时内心只想说这数据集水太深了。

6. 部署前最后一个动作:在陌生工地场景上做一次推理验收

6.1 只盯 mAP 不够,还要过一遍边缘场景清单

训练指标再好,都不如在没见过的工地场景上跑一遍。常见做法是准备一段 1 分钟的真实施工视频,包含白天强光、阴天、工人蹲着焊接、远处塔吊上的工人、多人重叠走动的画面。把这段视频从训练集和验证集里完全摘出来,当作“压测集”。压测时不看 mAP,直接看两件事:有没有漏检,以及安全帽小目标到底能不能在视频帧里被稳定框出。按下表逐项打勾,比跑一百个指标图都实在:

验收项通过标准实测手段
远距离小目标10×10 像素安全帽能检出 60% 以上降采样视频或裁剪小区域测试
人员密集重叠重叠框不丢失且不误检用多人并行施工片段测试
光照剧烈变化逆光、阴影下不出现大面积漏检分别截取正光/逆光帧比较
未佩戴负样本没戴安全帽的工人能被 person 框出且不判成 helmet统计负样本帧的误检率
真实场景推理速度满足现场摄像头帧率要求导出 ONNX 后测单帧耗时

6.2 导出 ONNX 并跑通一次实时推理

模型训练完不要直接用best.pt交付,工地现场的机器大概率没有 PyTorch 环境,装齐依赖不现实。我最后一次部署时就是把模型导出成 ONNX,让现场一台无 GPU 的普通台式机也能跑起来。导出命令很简单:

yolo export model=runs/helmet/exp/weights/best.pt format=onnx imgsz=640 opset=12

导出后用 onnxruntime 加载模型做一次最小推理验证,确认导出的模型和 PyTorch 版本输出差异在可接受范围:

import cv2 import onnxruntime as ort import numpy as np session = ort.InferenceSession("best.onnx") input_name = session.get_inputs()[0].name image = cv2.imread("val_sample.jpg") image = cv2.resize(image, (640, 640)) image = image[:, :, ::-1] # BGR转RGB image = image.astype(np.float32) / 255.0 input_data = image.transpose(2, 0, 1)[None, ...] outputs = session.run(None, {input_name: input_data}) # outputs[0]包含检测框、类别和置信度,按YOLO后处理解析

我自己的习惯是导出后还会再用 ONNX 跑一次推理视频,抽几帧和 PyTorch 的结果并排放在一起对比框的位置和置信度,确保没有因为算子不支持导致输出退化。这两年做过的工地安全设备检测项目,凡是验收顺利的,都在这一步花了功夫;凡是临时抱佛脚直接交best.pt的,基本都在现场装环境时翻车。把最后一道验收做完整,这个项目才算真正立得住,希望帮到你。

本文还有配套的精品资源,点击获取

返回列表