大家好,我是南木——专注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支持“检测+分割”(后续若需缺陷分割,无需重新标注),步骤如下:

  1. 安装LabelStudio:pip install label-studio
  2. 启动LabelStudio:终端输入label-studio,自动打开浏览器界面;
  3. 新建项目:点击“Create Project”,输入项目名(如“金属表面缺陷检测”);
  4. 导入数据:点击“Data Import”,上传采集的金属图像(支持批量上传);
  5. 配置标注类型:点击“Settings→Labeling Interface→Object Detection”,添加6个标签(Scratch、Dent、Inclusion、Indentation、Crack、Oxide);
  6. 开始标注:点击“Label”进入标注界面,用“Rectangle”工具框选缺陷,选择对应标签,标注完成后点击“Submit”。
(3)标注格式导出与转换

LabelStudio标注完成后,导出为“YOLO”格式(直接适配YOLOv8):

  1. 点击“Export”→选择“YOLO”→点击“Export”下载压缩包;
  2. 解压后,将“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_loss vs val/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个核心指标:

  1. 精度:测试集180张图像,缺陷检测率≥98%(漏检≤2%),误检率≤1%;
  2. 速度:平均推理速度≥20 FPS(匹配工业相机15 FPS的采集帧率,无卡顿);
  3. 稳定性:连续运行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. 后续优化方向(工业项目迭代)

  1. 多模态融合:结合红外图像(金属缺陷在红外下更明显),提升裂纹等隐蔽缺陷的检测率;
  2. 实时数据反馈:将缺陷数据(类型、位置、严重程度)上传到MES系统(工业生产执行系统),用于生产质量分析;
  3. 模型增量训练:定期用新采集的缺陷样本增量训练模型,适应新的缺陷类型(如新型氧化皮);
  4. 多模型融合:用yolov8s(高精度)和yolov8n(高速度)双重验证,小目标用yolov8s检测,大目标用yolov8n检测,平衡精度和速度。

结语

本文从工业场景痛点出发,带大家走通了“金属表面缺陷检测”的全流程——从数据集准备(公开+自定义标注)到模型训练(针对性调优),再到边缘部署(Jetson+工业相机),每个环节都附可落地的代码和避坑指南。工业AI的核心不是“追求高指标”,而是“解决实际问题”——比如本文通过“小目标增强”“模型剪枝”“TensorRT加速”,在保证精度的前提下,满足了工业现场的速度和硬件需求。

如果在部署过程中遇到相机连接、模型加速等问题,欢迎在评论区留言——我会定期挑选工业场景的高频问题,写专题解决方案。如果觉得本文有用,也欢迎点赞、收藏、转发,让更多工业AI从业者少走弯路!

在这里插入图片描述

Logo

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

更多推荐