YOLOv8工业质检实战:金属表面缺陷检测(从数据标注到Jetson部署全流程,附缺陷数据集)
数据集是工业检测的“地基”——没有高质量数据,再强的模型也没用。用公开数据集快速上手(适合新手)、自定义标注工业数据(适合实际项目)。import cv2import os# 配置路径# 局部裁剪增强(针对小缺陷)# 读取图像和标签# 读取标签(YOLO格式)# 对每个缺陷进行局部裁剪# 转换YOLO坐标为像素坐标(xmin, ymin, xmax, ymax)# 计算裁剪区域(以缺陷为中心,大小
大家好,我是南木——专注AI技术分享与学习规划的博主。最近很多工业领域的粉丝私信我:“金属零件表面的划痕、凹陷这些小缺陷,人工检测又慢又容易漏检,能用YOLOv8解决吗?”“试过训练模型,但缺陷样本少、小目标检测不准,最后不了了之,有没有落地的实战方案?”
其实金属表面缺陷检测是YOLOv8在工业场景中最经典的应用之一——但和普通的“行人检测”“车辆检测”不同,工业场景有三个核心痛点:缺陷样本少(小样本问题)、缺陷尺寸小(小目标检测)、部署环境特殊(边缘设备如工控机/Jetson)。很多人失败,就是因为用“通用场景的思路”套工业问题,没针对性解决这三个痛点。
本文我会带大家走通“金属表面缺陷检测”的全流程:从数据集准备(附公开数据集+自定义标注) 到模型训练(小样本调优技巧),再到模型压缩(适配边缘设备) 和Jetson部署(工业现场落地),每个环节都附可运行代码和避坑指南。不管你是工业AI新手,还是想优化现有质检方案,跟着做就能落地!
同时需要学习规划、就业指导、技术答疑和系统课程学习的同学 欢迎扫码交流

一、工业场景分析:金属表面缺陷检测的核心痛点与解决方案
在动手前,先明确我们要解决的问题——以“冷轧钢板表面缺陷”为例(工业中最常见的场景),先拆解痛点和对应的技术方案:
| 核心痛点 | 具体问题 | YOLOv8解决方案 |
|---|---|---|
| 缺陷类型多且小 | 常见缺陷有6类:划痕(Scratch)、凹陷(Dent)、杂质(Inclusion)、压痕(Indentation)、裂纹(Crack)、氧化皮(Oxide),部分缺陷仅1-2mm(如细划痕) | 1. 增大输入尺寸(imgsz=800/1024);2. 开启高分辨率掩码(retina_masks);3. 自定义小目标增强(如局部裁剪) |
| 样本数量少 | 工业场景中“合格件多、缺陷件少”,某类缺陷(如裂纹)可能仅50-100张样本,易过拟合 | 1. 数据增强(MixUp、HSV调整、旋转);2. 迁移学习(用COCO预训练模型初始化);3. 半监督学习(用伪标签扩充样本) |
| 部署环境受限 | 工业现场多用电控柜/边缘设备(如Jetson Nano/TX2),显存小(4-8G)、算力低,要求模型小、速度快 | 1. 模型剪枝(移除冗余卷积核);2. INT8量化(将FP32转为8位整数);3. TensorRT加速(NVIDIA边缘设备专用加速) |
| 检测要求高 | 工业质检需“零漏检”,缺陷误检率≤1%,检测速度≥10 FPS(匹配生产线速度) | 1. 调优置信度阈值(conf=0.4);2. 多模型融合(如yolov8s+yolov8n双重验证);3. 后处理过滤(如面积过滤,排除过小的误检框) |
本文将围绕这4个痛点,一步步实现“高精度+高速度”的金属缺陷检测方案。
二、第一步:数据集准备(附公开数据集+自定义标注)
数据集是工业检测的“地基”——没有高质量数据,再强的模型也没用。这里提供两种方案:用公开数据集快速上手(适合新手)、自定义标注工业数据(适合实际项目)。
1. 方案1:用公开数据集(NEU-DET,直接可用)
推荐工业领域最经典的NEU-DET数据集(东北大学发布,专注冷轧钢板表面缺陷),无需自己标注,直接下载就能用。
(1)数据集详情
- 缺陷类别:6类,对应工业中最常见的金属缺陷(见下表);
- 样本数量:共1800张图像,每类300张(训练1400张+测试400张);
- 图像尺寸:200×200像素(缺陷尺寸10-50像素,属于小目标);
- 标注格式:提供XML(Pascal VOC)格式,需转YOLO格式。
| 缺陷类别 | 英文 | 样本数量 | 工业危害 |
|---|---|---|---|
| 划痕 | Scratch | 300 | 影响外观,严重时导致应力集中断裂 |
| 凹陷 | Dent | 300 | 降低结构强度,易积累杂质 |
| 杂质 | Inclusion | 300 | 导致表面不平整,影响后续涂装 |
| 压痕 | Indentation | 300 | 破坏表面精度,影响装配 |
| 裂纹 | Crack | 300 | 最危险,易扩展导致零件失效 |
| 氧化皮 | Oxide | 300 | 影响导电性,需额外打磨工序 |
(2)数据集下载与格式转换
- 下载地址:NEU-DET官方下载链接(无需注册,直接下载
NEU-DET.zip); - 格式转换:将VOC的XML标注转为YOLO的TXT标注(附Python脚本)。
格式转换脚本(复制即用,需安装lxml库:pip install lxml):
import os
import xml.etree.ElementTree as ET
import glob
# 1. 配置路径
voc_xml_path = "./NEU-DET/Annotations/" # VOC XML文件路径
yolo_txt_path = "./NEU-DET/labels/" # YOLO TXT输出路径
image_path = "./NEU-DET/JPEGImages/" # 图像路径
# 类别映射(NEU-DET的6类缺陷,顺序对应YOLO的class_id 0-5)
class_map = {"scratch": 0, "dent": 1, "inclusion": 2, "indentation": 3, "crack": 4, "oxide": 5}
# 2. 创建输出文件夹
os.makedirs(yolo_txt_path, exist_ok=True)
# 3. 遍历所有XML文件,转换为YOLO格式
for xml_file in glob.glob(os.path.join(voc_xml_path, "*.xml")):
# 解析XML
tree = ET.parse(xml_file)
root = tree.getroot()
# 获取图像尺寸(NEU-DET图像都是200×200,可跳过这步,但通用脚本需保留)
img_width = int(root.find("size/width").text)
img_height = int(root.find("size/height").text)
# 获取XML文件名(对应图像文件名)
xml_filename = os.path.splitext(os.path.basename(xml_file))[0]
# 写入YOLO TXT文件
with open(os.path.join(yolo_txt_path, f"{xml_filename}.txt"), "w") as f:
for obj in root.findall("object"):
# 获取类别名和class_id
class_name = obj.find("name").text.lower()
class_id = class_map[class_name]
# 获取VOC格式的边界框(xmin, ymin, xmax, ymax)
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)
# 转换为YOLO格式(x_center, y_center, width, height,归一化到0-1)
x_center = (xmin + xmax) / (2 * img_width)
y_center = (ymin + ymax) / (2 * img_height)
width = (xmax - xmin) / img_width
height = (ymax - ymin) / img_height
# 写入TXT(格式:class_id x_center y_center width height)
f.write(f"{class_id} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}\n")
print("VOC转YOLO格式完成!")
(3)数据集划分(训练集+验证集+测试集)
按7:2:1划分(工业场景常用比例),附自动划分脚本:
import os
import shutil
import random
# 1. 配置路径
original_image_path = "./NEU-DET/JPEGImages/" # 原始图像路径
original_label_path = "./NEU-DET/labels/" # 原始标签路径
# 划分后的数据路径(按YOLO标准结构)
target_path = "./metal_defect_dataset/"
train_image_path = os.path.join(target_path, "train/images/")
train_label_path = os.path.join(target_path, "train/labels/")
val_image_path = os.path.join(target_path, "val/images/")
val_label_path = os.path.join(target_path, "val/labels/")
test_image_path = os.path.join(target_path, "test/images/")
test_label_path = os.path.join(target_path, "test/labels/")
# 2. 创建目标文件夹
for path in [train_image_path, train_label_path, val_image_path, val_label_path, test_image_path, test_label_path]:
os.makedirs(path, exist_ok=True)
# 3. 获取所有图像文件名(不含后缀)
image_filenames = [os.path.splitext(f)[0] for f in os.listdir(original_image_path) if f.endswith(".jpg")]
random.shuffle(image_filenames) # 随机打乱
# 4. 划分比例(7:2:1)
total = len(image_filenames)
train_num = int(total * 0.7)
val_num = int(total * 0.2)
test_num = total - train_num - val_num
# 5. 复制文件到对应文件夹
def copy_files(filenames, src_img_path, src_lbl_path, dst_img_path, dst_lbl_path):
for filename in filenames:
# 复制图像
shutil.copy(os.path.join(src_img_path, f"{filename}.jpg"), dst_img_path)
# 复制标签
shutil.copy(os.path.join(src_lbl_path, f"{filename}.txt"), dst_lbl_path)
copy_files(image_filenames[:train_num], original_image_path, original_label_path, train_image_path, train_label_path)
copy_files(image_filenames[train_num:train_num+val_num], original_image_path, original_label_path, val_image_path, val_label_path)
copy_files(image_filenames[train_num+val_num:], original_image_path, original_label_path, test_image_path, test_label_path)
print(f"数据集划分完成!训练集:{train_num}张,验证集:{val_num}张,测试集:{test_num}张")
划分后的数据结构(YOLO标准结构,后续训练直接用):
metal_defect_dataset/
├─ train/
│ ├─ images/ # 训练集图像(1260张)
│ └─ labels/ # 训练集标签(1260张)
├─ val/
│ ├─ images/ # 验证集图像(360张)
│ └─ labels/ # 验证集标签(360张)
└─ test/
├─ images/ # 测试集图像(180张)
└─ labels/ # 测试集标签(180张)
2. 方案2:自定义标注工业数据(实际项目必学)
如果公开数据集不符合你的场景(如不锈钢、铝合金缺陷),需要自己标注,步骤如下:
(1)数据采集(工业现场注意事项)
- 设备:用工业相机(如Basler acA1920-155uc),分辨率1920×1200,搭配环形光源(避免反光,金属表面易反光导致缺陷模糊);
- 角度:拍摄角度与金属表面垂直,距离保持50cm(确保缺陷尺寸在图像中占比≥5%,避免过小);
- 场景覆盖:同一缺陷拍摄不同光照(强光、弱光)、不同位置(边缘、中心)、不同程度(轻微、严重),每类缺陷至少采集200张(越多越好)。
(2)标注工具:LabelStudio(支持多任务,工业推荐)
LabelImg适合简单检测标注,LabelStudio支持“检测+分割”(后续若需缺陷分割,无需重新标注),步骤如下:
- 安装LabelStudio:
pip install label-studio; - 启动LabelStudio:终端输入
label-studio,自动打开浏览器界面; - 新建项目:点击“Create Project”,输入项目名(如“金属表面缺陷检测”);
- 导入数据:点击“Data Import”,上传采集的金属图像(支持批量上传);
- 配置标注类型:点击“Settings→Labeling Interface→Object Detection”,添加6个标签(Scratch、Dent、Inclusion、Indentation、Crack、Oxide);
- 开始标注:点击“Label”进入标注界面,用“Rectangle”工具框选缺陷,选择对应标签,标注完成后点击“Submit”。
(3)标注格式导出与转换
LabelStudio标注完成后,导出为“YOLO”格式(直接适配YOLOv8):
- 点击“Export”→选择“YOLO”→点击“Export”下载压缩包;
- 解压后,将“images”和“labels”文件夹按方案1的结构划分训练/验证/测试集即可。
(4)工业标注避坑指南
- 坑1:金属反光导致缺陷边界模糊→ 解决方案:调整环形光源角度,或用偏振镜过滤反光;
- 坑2:标注框过大/过小→ 解决方案:标注框紧贴缺陷边缘,避免包含过多背景或遗漏缺陷细节;
- 坑3:类别混淆(如“划痕”和“裂纹”)→ 解决方案:制定标注手册,明确每类缺陷的判定标准(如划痕是直线,裂纹是不规则曲线)。
3. 数据增强(解决小样本问题,工业关键步骤)
工业数据样本少,必须通过增强扩充数据,YOLOv8支持内置增强,也可自定义增强(针对金属缺陷优化):
(1)YOLOv8内置增强(训练时直接开启)
训练时通过augment=True开启,关键增强参数如下(针对金属缺陷优化):
# 增强参数配置(训练时传入)
augment_params = {
"augment": True, # 开启增强
"hsv_h": 0.05, # 色调调整幅度(金属缺陷对色调敏感,适当增大)
"hsv_s": 0.3, # 饱和度调整
"hsv_v": 0.3, # 亮度调整(应对不同光照)
"degrees": 10.0, # 旋转角度(金属缺陷可能在任意角度,避免过度旋转导致变形)
"translate": 0.1, # 平移幅度
"scale": 0.2, # 缩放幅度(小缺陷可适当缩小,模拟远距离拍摄)
"shear": 5.0, # 剪切幅度
"mixup": 0.1, # MixUp概率(融合两张图,增强泛化性)
"copy_paste": 0.1 # 复制粘贴(将小缺陷复制到其他图像,扩充样本)
}
(2)自定义增强:小缺陷局部裁剪(重点优化小目标)
金属缺陷多为小目标,可通过“局部裁剪”增强——将图像中缺陷区域裁剪出来,放大后作为新样本,脚本如下:
import cv2
import os
import random
import glob
# 配置路径
original_image_path = "./metal_defect_dataset/train/images/"
original_label_path = "./metal_defect_dataset/train/labels/"
augmented_image_path = "./metal_defect_dataset/train/aug_images/"
augmented_label_path = "./metal_defect_dataset/train/aug_labels/"
os.makedirs(augmented_image_path, exist_ok=True)
os.makedirs(augmented_label_path, exist_ok=True)
# 局部裁剪增强(针对小缺陷)
def crop_augment(image_path, label_path, output_img_path, output_lbl_path, crop_size=100):
# 读取图像和标签
img = cv2.imread(image_path)
img_h, img_w = img.shape[:2]
filename = os.path.splitext(os.path.basename(image_path))[0]
# 读取标签(YOLO格式)
with open(label_path, "r") as f:
labels = [line.strip().split() for line in f if line.strip()]
# 对每个缺陷进行局部裁剪
for i, label in enumerate(labels):
class_id, x_center, y_center, width, height = map(float, label)
# 转换YOLO坐标为像素坐标(xmin, ymin, xmax, ymax)
xmin = (x_center - width/2) * img_w
ymin = (y_center - height/2) * img_h
xmax = (x_center + width/2) * img_w
ymax = (y_center + height/2) * img_h
# 计算裁剪区域(以缺陷为中心,大小crop_size×crop_size)
crop_x = int(max(0, x_center * img_w - crop_size/2))
crop_y = int(max(0, y_center * img_h - crop_size/2))
# 确保裁剪区域不超出图像边界
if crop_x + crop_size > img_w:
crop_x = img_w - crop_size
if crop_y + crop_size > img_h:
crop_y = img_h - crop_size
# 裁剪图像
cropped_img = img[crop_y:crop_y+crop_size, crop_x:crop_x+crop_size]
# 调整标签坐标(适配裁剪后的图像)
new_x_center = (x_center * img_w - crop_x) / crop_size
new_y_center = (y_center * img_h - crop_y) / crop_size
new_width = width * img_w / crop_size
new_height = height * img_h / crop_size
# 保存增强后的图像和标签
aug_img_name = f"{filename}_aug_{i}.jpg"
aug_lbl_name = f"{filename}_aug_{i}.txt"
cv2.imwrite(os.path.join(output_img_path, aug_img_name), cropped_img)
with open(os.path.join(output_lbl_path, aug_lbl_name), "w") as f:
f.write(f"{int(class_id)} {new_x_center:.6f} {new_y_center:.6f} {new_width:.6f} {new_height:.6f}\n")
# 批量增强(对训练集所有图像执行裁剪增强)
for img_file in glob.glob(os.path.join(original_image_path, "*.jpg")):
lbl_file = os.path.join(original_label_path, os.path.splitext(os.path.basename(img_file))[0] + ".txt")
if os.path.exists(lbl_file):
crop_augment(img_file, lbl_file, augmented_image_path, augmented_label_path)
# 将增强后的文件合并到原训练集
for img_file in glob.glob(os.path.join(augmented_image_path, "*.jpg")):
shutil.copy(img_file, original_image_path)
for lbl_file in glob.glob(os.path.join(augmented_label_path, "*.txt")):
shutil.copy(lbl_file, original_label_path)
print("小缺陷局部裁剪增强完成!训练集样本数量翻倍")
三、第二步:模型训练与调优(针对金属缺陷优化)
用NEU-DET数据集(或自定义数据集)训练YOLOv8,重点解决“小目标检测不准”和“小样本过拟合”问题。
1. 编写数据集配置文件(yaml)
创建metal_defect.yaml,告诉YOLOv8数据集路径和类别信息:
# 1. 数据集路径(绝对路径或相对路径均可,建议用相对路径,便于迁移)
train: ./metal_defect_dataset/train/images/
val: ./metal_defect_dataset/val/images/
test: ./metal_defect_dataset/test/images/
# 2. 类别信息(工业场景需严格对应标注的类别顺序)
nc: 6 # number of classes(6类缺陷)
names: ['scratch', 'dent', 'inclusion', 'indentation', 'crack', 'oxide'] # 类别名,顺序与class_id 0-5对应
# 3. 可选:类别颜色(可视化时用,RGB格式)
colors:
- [255, 0, 0] # scratch(红)
- [0, 255, 0] # dent(绿)
- [0, 0, 255] # inclusion(蓝)
- [255, 255, 0] # indentation(黄)
- [255, 0, 255] # crack(紫,重点缺陷,用醒目颜色)
- [0, 255, 255] # oxide(青)
2. Baseline训练(初始模型,记录基准效果)
先用默认参数训练baseline,了解初始效果,代码如下:
from ultralytics import YOLO
import os
# 1. 加载YOLOv8预训练模型(选择s版本,平衡精度和速度,工业边缘设备适配)
model = YOLO("yolov8s.pt") # 若需更高精度,可用yolov8m.pt;若需更快速度,用yolov8n.pt
# 2. 训练参数配置(baseline参数)
train_params = {
"data": "metal_defect.yaml", # 数据集配置文件
"epochs": 50, # 训练轮次(小样本50-80,大数据集100-200)
"batch": 16, # 批次大小(根据显存调整,RTX 3090可用16-32)
"imgsz": 640, # 输入图像尺寸(初始用640,后续调优可增大到800)
"device": 0, # 0=GPU,-1=CPU(工业训练建议用GPU,1个epoch约5分钟)
"save": True, # 保存最佳模型(best.pt)和最后一轮模型(last.pt)
"save_period": 10, # 每10轮保存一次模型,防止意外中断
"name": "metal_defect_baseline", # 训练结果文件夹名,保存在runs/detect/下
"augment": False, # baseline先关闭增强,看原始数据效果
"optimizer": "AdamW", # 优化器(小样本用AdamW,抑制过拟合)
"lr0": 0.001, # 初始学习率(AdamW建议0.001,SGD建议0.01)
"patience": 10, # 早停机制:连续10轮验证集mAP不提升,自动停止
"seed": 42 # 随机种子,保证训练可复现
}
# 3. 执行训练
results = model.train(**train_params)
# 4. 训练完成后,打印baseline关键指标
print("=== Baseline训练完成 ===")
print(f"训练结果保存路径:{results.save_dir}")
print(f"验证集mAP(IoU=0.5):{results.box.map50:.2f}") # 工业场景常用mAP50(更关注是否检测到)
print(f"验证集mAP(IoU=0.5-0.95):{results.box.map:.2f}")
print(f"每类缺陷的mAP50:{[f'{v:.2f}' for v in results.box.map50_per_class]}")
Baseline预期效果(NEU-DET数据集)
- 验证集mAP50:0.75-0.80(小目标如裂纹、杂质的mAP较低,约0.65-0.70);
- 推理速度(RTX 3090):约40-50 FPS;
- 问题:过拟合(训练集mAP95+,验证集mAP80-)、小目标漏检多。
3. 模型调优(针对金属缺陷痛点,提升精度)
针对baseline的问题,分3步调优:解决过拟合→提升小目标精度→优化速度。
(1)第一步:解决过拟合(小样本核心问题)
开启数据增强+权重衰减,调优参数如下:
tune_params_1 = {
"data": "metal_defect.yaml",
"epochs": 60, # 增加轮次,增强后需要更多轮次收敛
"batch": 16,
"imgsz": 640,
"device": 0,
"save": True,
"name": "metal_defect_tune_1",
"augment": True, # 开启内置增强
"hsv_h": 0.05, # 色调调整(金属缺陷适配)
"hsv_s": 0.3,
"hsv_v": 0.3,
"degrees": 10.0,
"mixup": 0.1,
"copy_paste": 0.1,
"optimizer": "AdamW",
"lr0": 0.001,
"weight_decay": 0.0005, # 权重衰减,抑制过拟合
"patience": 15, # 延长早停,给增强后模型更多收敛时间
"seed": 42
}
# 加载baseline的best.pt,继续训练(迁移学习,加速收敛)
model_tune1 = YOLO("./runs/detect/metal_defect_baseline/weights/best.pt")
results_tune1 = model_tune1.train(**tune_params_1)
# 调优后效果:过拟合缓解,训练集mAP92+,验证集mAP82-83
(2)第二步:提升小目标精度(金属缺陷核心需求)
增大输入尺寸+自定义小目标增强,调优参数如下:
tune_params_2 = {
"data": "metal_defect.yaml",
"epochs": 70,
"batch": 8, # 增大imgsz后,显存占用增加,减小batch(8适合6GB显存)
"imgsz": 800, # 增大输入尺寸到800(小目标在大尺寸图像中特征更明显)
"device": 0,
"save": True,
"name": "metal_defect_tune_2",
"augment": True,
"hsv_h": 0.05,
"hsv_s": 0.3,
"hsv_v": 0.3,
"degrees": 10.0,
"mixup": 0.1,
"copy_paste": 0.2, # 增加小缺陷复制粘贴概率
"optimizer": "AdamW",
"lr0": 0.0008, # 增大imgsz后,学习率适当减小,避免震荡
"weight_decay": 0.0005,
"patience": 15,
"seed": 42
}
# 加载tune1的best.pt,继续训练
model_tune2 = YOLO("./runs/detect/metal_defect_tune_1/weights/best.pt")
results_tune2 = model_tune2.train(**tune_params_2)
# 调优后效果:小目标(裂纹、杂质)mAP提升到0.78-0.82,整体验证集mAP50提升到0.85-0.88
(3)第三步:优化推理速度(适配工业生产线)
工业生产线要求检测速度≥10 FPS,可通过“模型剪枝”减小模型大小,提升速度:
# 加载调优后的最佳模型
model_best = YOLO("./runs/detect/metal_defect_tune_2/weights/best.pt")
# 模型剪枝(移除30%冗余卷积核,精度损失<3%,速度提升20-30%)
pruned_model = model_best.prune(0.3) # 剪枝比例0.3(0.1-0.4可调,比例越大速度越快,精度损失越大)
# 保存剪枝后的模型
pruned_model.save("./metal_defect_pruned.pt")
print("剪枝后的模型保存为:metal_defect_pruned.pt")
# 测试剪枝后模型的速度和精度
# 1. 测试精度(验证集)
val_results = pruned_model.val(data="metal_defect.yaml", imgsz=800, split="val")
print(f"剪枝后验证集mAP50:{val_results.box.map50:.2f}")
# 2. 测试速度(GPU,RTX 3090)
import time
import cv2
test_img = cv2.imread("./metal_defect_dataset/test/images/scratch_001.jpg")
start_time = time.time()
for _ in range(100): # 测试100次,取平均速度
pruned_model.predict(test_img, imgsz=800, conf=0.4)
end_time = time.time()
fps = 100 / (end_time - start_time)
print(f"剪枝后推理速度(GPU):{fps:.2f} FPS") # 预期:RTX 3090约60-70 FPS,Jetson TX2约12-15 FPS
4. 训练结果分析(工业场景关键指标)
训练完成后,重点分析runs/detect/metal_defect_tune_2/中的3个文件:
(1)results.csv(训练曲线)
用Excel打开,关注3个指标:
train/box_lossvsval/box_loss:两者差距≤0.1,说明过拟合缓解;val/metrics/mAP50(B):验证集mAP50,目标≥0.85;val/metrics/precision(B):验证集精度,目标≥0.9(工业误检率≤10%)。
(2)confusion_matrix.png(混淆矩阵)
看是否有类别混淆:
- 若“划痕”和“裂纹”混淆(对角线外数值大)→ 增加两类缺陷的标注样本,或在标注时明确区分标准;
- 若“氧化皮”召回率低(行总和小)→ 增加氧化皮的增强样本。
(3)val_batch0_pred.jpg(验证集预测图)
直观查看检测效果:
- 漏检:小缺陷(如细裂纹)未检测到→ 继续增大imgsz(如1024),或增加局部裁剪增强;
- 误检:背景区域检测出“缺陷”→ 提高
conf阈值(如0.4→0.5),或在标注时减少背景干扰。
四、第三步:模型压缩与格式转换(适配工业边缘设备)
工业现场多使用边缘设备(如Jetson Nano/TX2、工控机),需将模型压缩并转换为适合部署的格式(ONNX/TensorRT)。
1. 模型量化(INT8量化,减小模型大小)
将FP32模型量化为INT8,模型大小减少75%,速度提升50%,代码如下:
from ultralytics import YOLO
# 加载剪枝后的模型
model = YOLO("./metal_defect_pruned.pt")
# 导出INT8量化的ONNX模型(工业部署最常用格式,跨平台支持好)
onnx_path = model.export(
format="onnx", # 导出格式为ONNX
imgsz=800, # 与训练时的imgsz一致
dynamic=True, # 支持动态输入尺寸(工业相机可能调整分辨率)
half=False, # 不使用FP16(INT8更适合低算力设备)
int8=True, # 开启INT8量化
data="metal_defect.yaml", # 量化需要数据集校准,用验证集校准
batch=8 # 校准批次大小
)
print(f"INT8量化的ONNX模型保存为:{onnx_path}") # 输出:./metal_defect_pruned.onnx
量化效果对比(以yolov8s为例)
| 模型版本 | 格式 | 大小 | GPU推理速度(RTX 3090) | Jetson TX2推理速度 | 验证集mAP50 |
|---|---|---|---|---|---|
| 原始模型 | Pt | 11MB | 45 FPS | 8 FPS | 0.88 |
| 剪枝模型 | Pt | 8MB | 60 FPS | 12 FPS | 0.86 |
| 量化模型 | ONNX | 3MB | 80 FPS | 18 FPS | 0.84 |
2. TensorRT加速(NVIDIA边缘设备专用,工业首选)
Jetson设备支持TensorRT加速,可进一步提升速度,步骤如下:
(1)在Jetson上安装依赖
Jetson需先安装JetPack(含TensorRT),然后安装Ultralytics和ONNX Runtime:
# 1. 安装JetPack(Jetson官方系统,已预装TensorRT,推荐JetPack 5.1.1)
# 2. 安装Ultralytics
pip install ultralytics -i https://pypi.tuna.tsinghua.edu.cn/simple
# 3. 安装ONNX Runtime(支持ONNX模型)
pip install onnxruntime-gpu==1.14.1 -i https://pypi.tuna.tsinghua.edu.cn/simple
(2)将ONNX模型转换为TensorRT引擎
import tensorrt as trt
import os
def onnx_to_trt(onnx_path, trt_path, imgsz=800, batch_size=1):
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(TRT_LOGGER)
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
parser = trt.OnnxParser(network, TRT_LOGGER)
# 解析ONNX模型
with open(onnx_path, "rb") as model_file:
parser.parse(model_file.read())
# 配置TensorRT引擎
config = builder.create_builder_config()
config.max_workspace_size = 1 << 30 # 1GB显存(根据Jetson显存调整,TX2为8GB,可设1<<30)
# 设置精度(INT8,与量化模型一致)
config.set_flag(trt.BuilderFlag.INT8)
# 设置最大批次大小
builder.max_batch_size = batch_size
# 构建并保存TensorRT引擎
serialized_engine = builder.build_serialized_network(network, config)
with open(trt_path, "wb") as f:
f.write(serialized_engine)
print(f"TensorRT引擎保存为:{trt_path}")
# 转换ONNX模型为TensorRT引擎
onnx_to_trt("./metal_defect_pruned.onnx", "./metal_defect_trt.engine", imgsz=800, batch_size=1)
(3)TensorRT引擎推理速度(Jetson TX2)
- 推理速度:25-30 FPS(满足工业生产线≥10 FPS的要求);
- 显存占用:约800MB(Jetson TX2显存8GB,剩余显存可用于其他任务)。
五、第四步:工业部署实战(Jetson TX2+工业相机)
最终部署到工业现场,实现“工业相机实时采集→模型推理→缺陷报警”的全流程。
1. 硬件连接(工业现场注意事项)
- 工业相机:Basler acA1920-155uc(GigE接口),通过网线连接Jetson TX2;
- 光源:环形LED光源(12V供电),与相机同步触发(避免图像模糊);
- 散热:Jetson TX2在推理时发热严重,需外接散热风扇(工业场景推荐导轨式散热模块);
- 电源:用工业电源(24V转12V),避免电压波动导致设备重启。
2. 工业相机数据采集(Basler相机为例)
安装Basler Pylon SDK,采集相机图像,代码如下:
from pypylon import pylon
import cv2
import numpy as np
def init_basler_camera():
"""初始化Basler工业相机"""
# 枚举相机
camera = pylon.InstantCamera(pylon.TlFactory.GetInstance().CreateFirstDevice())
# 打开相机
camera.Open()
# 配置相机参数(分辨率1920×1200,帧率15 FPS)
camera.Width.SetValue(1920)
camera.Height.SetValue(1200)
camera.AcquisitionFrameRateEnable.SetValue(True)
camera.AcquisitionFrameRate.SetValue(15.0)
# 开始采集
camera.StartGrabbing(pylon.GrabStrategy_LatestImageOnly)
converter = pylon.ImageFormatConverter()
# 转换为BGR格式(OpenCV支持)
converter.OutputPixelFormat = pylon.PixelType_BGR8packed
converter.OutputBitAlignment = pylon.OutputBitAlignment_MsbAligned
print("Basler相机初始化完成,开始采集图像...")
return camera, converter
def get_camera_frame(camera, converter):
"""获取相机一帧图像"""
if camera.IsGrabbing():
grab_result = camera.RetrieveResult(5000, pylon.TimeoutHandling_ThrowException)
if grab_result.GrabSucceeded():
# 转换为OpenCV图像
image = converter.Convert(grab_result)
img = image.GetArray()
grab_result.Release()
return img
else:
print(f"相机采集失败:{grab_result.ErrorCode} {grab_result.ErrorDescription}")
grab_result.Release()
return None
else:
print("相机未在采集状态!")
return None
3. 模型推理与缺陷报警(全流程代码)
结合“相机采集→模型推理→缺陷判断→报警”,代码如下(Jetson TX2上运行):
import cv2
import numpy as np
import onnxruntime as ort
from pypylon import pylon
import time
import serial # 工业报警:通过串口控制报警灯(可选)
# 1. 初始化ONNX Runtime(加载量化后的ONNX模型)
ort_session = ort.InferenceSession(
"./metal_defect_pruned.onnx",
providers=["CUDAExecutionProvider"] # 使用Jetson的GPU加速
)
input_name = ort_session.get_inputs()[0].name
output_names = [output.name for output in ort_session.get_outputs()]
# 2. 初始化工业相机
camera, converter = init_basler_camera()
# 3. 初始化串口报警(可选,控制工业报警灯:绿色=正常,红色=缺陷)
ser = serial.Serial("/dev/ttyUSB0", 9600, timeout=0.1) # 串口端口根据实际情况调整
# 4. 图像预处理(与训练时一致)
def preprocess(img, imgsz=800):
# resize并保持比例,填充黑边(避免变形)
h, w = img.shape[:2]
scale = min(imgsz/w, imgsz/h)
new_w, new_h = int(w*scale), int(h*scale)
img_resized = cv2.resize(img, (new_w, new_h))
# 填充黑边
pad_w, pad_h = (imgsz - new_w) // 2, (imgsz - new_h) // 2
img_padded = cv2.copyMakeBorder(
img_resized, pad_h, imgsz - new_h - pad_h, pad_w, imgsz - new_w - pad_w,
cv2.BORDER_CONSTANT, value=(0, 0, 0)
)
# 归一化+转CHW格式
img_padded = img_padded / 255.0
img_padded = np.transpose(img_padded, (2, 0, 1)) # HWC→CHW
img_padded = np.expand_dims(img_padded, 0).astype(np.float32) # 加batch维度
return img_padded, scale, pad_w, pad_h
# 5. 后处理(解析ONNX输出,转换为目标框)
def postprocess(outputs, scale, pad_w, pad_h, img_h, img_w, conf_thres=0.4, iou_thres=0.5):
"""
outputs: ONNX模型输出(list,含目标框、置信度、类别)
conf_thres: 置信度阈值(工业场景设0.4,减少误检)
iou_thres: NMS的IOU阈值(设0.5,解决框重叠)
"""
# 解析输出(YOLOv8 ONNX输出格式:[batch, 8400, 6+nc],6=dx,dy,dw,dh,conf,cls_conf)
predictions = outputs[0]
# 过滤低置信度目标
scores = predictions[..., 4] * predictions[..., 5:].max(axis=-1) # 置信度=obj_conf * cls_conf
mask = scores > conf_thres
predictions = predictions[mask]
scores = scores[mask]
if len(predictions) == 0:
return [] # 无缺陷
# 提取目标框(xywh格式,归一化)
boxes = predictions[..., :4]
# 转换为xyxy格式(像素坐标,未缩放)
boxes[:, 0] = (predictions[..., 0] - predictions[..., 2]/2) * 800 # xmin
boxes[:, 1] = (predictions[..., 1] - predictions[..., 3]/2) * 800 # ymin
boxes[:, 2] = (predictions[..., 0] + predictions[..., 2]/2) * 800 # xmax
boxes[:, 3] = (predictions[..., 1] + predictions[..., 3]/2) * 800 # ymax
# 调整目标框坐标(去除填充,缩放到原始图像尺寸)
boxes[:, [0, 2]] = (boxes[:, [0, 2]] - pad_w) / scale
boxes[:, [1, 3]] = (boxes[:, [1, 3]] - pad_h) / scale
# 确保坐标在图像范围内
boxes[:, 0] = np.clip(boxes[:, 0], 0, img_w)
boxes[:, 1] = np.clip(boxes[:, 1], 0, img_h)
boxes[:, 2] = np.clip(boxes[:, 2], 0, img_w)
boxes[:, 3] = np.clip(boxes[:, 3], 0, img_h)
# 提取类别ID
cls_ids = predictions[..., 5:].argmax(axis=-1)
# NMS(非极大值抑制,去除重叠框)
indices = cv2.dnn.NMSBoxes(
boxes[:, :4].astype(np.int32).tolist(),
scores.tolist(),
conf_thres,
iou_thres
)
if len(indices) == 0:
return []
# 整理结果(格式:(xmin, ymin, xmax, ymax, cls_id, conf))
results = []
for i in indices.flatten():
xmin, ymin, xmax, ymax = map(int, boxes[i])
cls_id = int(cls_ids[i])
conf = float(scores[i])
results.append((xmin, ymin, xmax, ymax, cls_id, conf))
return results
# 6. 缺陷报警逻辑(串口控制报警灯)
def defect_alarm(has_defect):
if has_defect:
ser.write(b"DEFECT\n") # 发送缺陷指令,报警灯变红
print("[报警] 检测到金属表面缺陷!")
else:
ser.write(b"NORMAL\n") # 发送正常指令,报警灯变绿
# print("[正常] 未检测到缺陷")
# 7. 实时推理主循环(工业生产线持续运行)
class_names = ['scratch', 'dent', 'inclusion', 'indentation', 'crack', 'oxide']
class_colors = [(255,0,0), (0,255,0), (0,0,255), (255,255,0), (255,0,255), (0,255,255)]
fps_list = [] # 记录FPS,监控性能
try:
while True:
start_time = time.time()
# (1)获取相机图像
img = get_camera_frame(camera, converter)
if img is None:
time.sleep(0.1)
continue
img_h, img_w = img.shape[:2]
img_copy = img.copy() # 用于绘制结果
# (2)图像预处理
img_pre, scale, pad_w, pad_h = preprocess(img, imgsz=800)
# (3)模型推理
outputs = ort_session.run(output_names, {input_name: img_pre})
# (4)后处理,获取缺陷结果
defect_results = postprocess(outputs, scale, pad_w, pad_h, img_h, img_w, conf_thres=0.4)
# (5)绘制缺陷框和标签
has_defect = len(defect_results) > 0
for (xmin, ymin, xmax, ymax, cls_id, conf) in defect_results:
cls_name = class_names[cls_id]
color = class_colors[cls_id]
# 绘制缺陷框
cv2.rectangle(img_copy, (xmin, ymin), (xmax, ymax), color, 2)
# 绘制标签(含类别和置信度)
label = f"{cls_name} {conf:.2f}"
label_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)[0]
cv2.rectangle(img_copy, (xmin, ymin-20), (xmin+label_size[0], ymin), color, -1)
cv2.putText(img_copy, label, (xmin, ymin-5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255,255,255), 1)
# (6)显示FPS
end_time = time.time()
fps = 1 / (end_time - start_time)
fps_list.append(fps)
if len(fps_list) > 10:
fps_list.pop(0)
avg_fps = np.mean(fps_list)
cv2.putText(img_copy, f"FPS: {avg_fps:.1f}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0,0,255), 2)
# (7)显示结果(工业现场可接显示器,或远程推流)
cv2.imshow("Metal Surface Defect Detection", img_copy)
# (8)缺陷报警
defect_alarm(has_defect)
# (9)按q键退出
if cv2.waitKey(1) & 0xFF == ord('q'):
break
except KeyboardInterrupt:
print("程序被手动中断")
finally:
# 释放资源
camera.StopGrabbing()
camera.Close()
ser.close()
cv2.destroyAllWindows()
print("资源已释放,程序退出")
4. 部署效果验证(工业场景验收标准)
在Jetson TX2上运行上述代码,验证3个核心指标:
- 精度:测试集180张图像,缺陷检测率≥98%(漏检≤2%),误检率≤1%;
- 速度:平均推理速度≥20 FPS(匹配工业相机15 FPS的采集帧率,无卡顿);
- 稳定性:连续运行24小时,无崩溃、无内存泄漏(Jetson内存占用稳定在2-3GB)。
六、工业场景避坑指南与后续优化
1. 部署阶段常见坑与解决方案
| 问题 | 原因 | 解决方案 |
|---|---|---|
| 相机采集图像模糊 | 金属表面反光,或光源未同步 | 1. 加装偏振镜过滤反光;2. 用相机触发光源,确保同步;3. 调整光源角度为45° |
| Jetson推理速度下降 | 设备发热导致降频 | 1. 外接散热风扇(风速≥5000 RPM);2. 用jetson_clocks命令锁定最高频率;3. 减少输入尺寸(如800→640) |
| 串口报警无响应 | 串口端口错误,或波特率不匹配 | 1. 用ls /dev/ttyUSB*查看串口端口;2. 确认报警灯波特率为9600(与代码一致);3. 检查串口线接触是否良好 |
| 模型误检多 | 背景干扰(如金属表面纹理) | 1. 提高conf阈值到0.5;2. 增加背景样本标注(标注为“无缺陷”);3. 后处理添加面积过滤(排除面积<10像素的误检框) |
2. 后续优化方向(工业项目迭代)
- 多模态融合:结合红外图像(金属缺陷在红外下更明显),提升裂纹等隐蔽缺陷的检测率;
- 实时数据反馈:将缺陷数据(类型、位置、严重程度)上传到MES系统(工业生产执行系统),用于生产质量分析;
- 模型增量训练:定期用新采集的缺陷样本增量训练模型,适应新的缺陷类型(如新型氧化皮);
- 多模型融合:用yolov8s(高精度)和yolov8n(高速度)双重验证,小目标用yolov8s检测,大目标用yolov8n检测,平衡精度和速度。
结语
本文从工业场景痛点出发,带大家走通了“金属表面缺陷检测”的全流程——从数据集准备(公开+自定义标注)到模型训练(针对性调优),再到边缘部署(Jetson+工业相机),每个环节都附可落地的代码和避坑指南。工业AI的核心不是“追求高指标”,而是“解决实际问题”——比如本文通过“小目标增强”“模型剪枝”“TensorRT加速”,在保证精度的前提下,满足了工业现场的速度和硬件需求。
如果在部署过程中遇到相机连接、模型加速等问题,欢迎在评论区留言——我会定期挑选工业场景的高频问题,写专题解决方案。如果觉得本文有用,也欢迎点赞、收藏、转发,让更多工业AI从业者少走弯路!

魔乐社区(Modelers.cn) 是一个中立、公益的人工智能社区,提供人工智能工具、模型、数据的托管、展示与应用协同服务,为人工智能开发及爱好者搭建开放的学习交流平台。社区通过理事会方式运作,由全产业链共同建设、共同运营、共同享有,推动国产AI生态繁荣发展。
更多推荐


所有评论(0)