目录

项目介绍

一、数据分析

1）数据集的定义

2）训练的脚本配置

3）数据准备

1、数据标注

2、EDA数据分析

二、YOLO模型训练与调优

1）数据增强优化

2）模型选择和图片选择

3）离线数据增强 (扩充训练集)

4）推理时增强TTA

还是以上一节内容yolo模型介绍（二）中钢铁缺陷检测的数据来做AI质检内容。我们的目标仍然是缺陷检测，但是流程上会更完善。

项目介绍

在制造业产线中，产品表面缺陷检测是品质把控的关键一环。传统人工目视质检长期面临效率、标准一致性、人力成本的多重压力，AI 视觉质检正在成为工业落地的主流方案。

目前工业视觉质检主要存在两套技术路线，二者定位不同，可以互为补充。

第一套是YOLO 系列目标检测方案，也是工业落地的主流选择，支持自定义标注数据完成训练，推理速度快，能够满足产线实时检测的要求；

第二套是VLM 多模态大模型，例如 Qwen‑VL，依靠提示词就可以实现零样本缺陷检测，不需要大量标注训练数据，适合快速做原型验证，但检测精度有限，很难直接用于正式产线。

本项目以钢铁缺陷检测数据集为例，打造AI质检系统，项目整体流程如下：

除此之外，还用OpenClaw + Skills 来搭建一个AI检测助理。如下是这个项目需要的工具栈：

本篇内容介绍前两个模块，即数据分析处理和YOLO模型训练与调优。

一、数据分析

如下，是我们整个项目给到的数据：

1）数据集的定义

我们的数据集定义在dataset.yaml这个文件里。在yolo模型介绍（二）里面我们介绍过dataset.yaml文件，它里面包含了数据集所在的位置，数据集的划分（训练集，测试集和验证集）以及数据标注中6种缺陷。

2）训练的脚本配置

如下是yolov12.yaml文件里的部分内容。nc表示我们要识别的缺陷的种类。

Scales下面这些是yolo的尺寸，从n到x层数越来越多，参数也越来越多。还记得上篇内容中第一个实例种有这样一句话吗

model = YOLO('./yolov12n.pt')

如果想用V12最强的模型，我们将参数改为./yolov12x.pt即可。

yolov12.yaml文件后面还有一些其他有关框架的内容，一般都不需要改。了解即可。

3）数据准备

1、数据标注

在实际项目中，最大的工作量是数据标注，要准备很多标注好的数据进行训练。一般使用labelImg或者label-studio工具做数据标注，yolo模型介绍（二）中介绍了labelImg的使用方法了。在本项目中，数据已经标好了。

2、EDA数据分析

EDA（explore data analysis）属于是在标注数据后，模型训练前的数据分析。我们可以读取标注文件的内容来进行分析。我们读取的是.txt格式的标注文件，格式如下：

第一个参数是dataset.yaml对应的缺陷类型，后面四个参数代表的是经过归一化处理的缺陷的位置(x,y,w,h)。

step 1：读取所有的标注文件，按照图片的名字，存入annotations这个结构中。

def load_all_labels(labels_dir): """读取所有标注文件, 返回 {filename: [(class_id, x, y, w, h), ...]}""" annotations = {} for label_file in glob.glob(os.path.join(labels_dir, "*.txt")): basename = os.path.splitext(os.path.basename(label_file))[0] boxes = [] with open(label_file, "r") as f: for line in f: parts = line.strip().split() if len(parts) == 5: cls_id = int(parts[0]) x, y, w, h = float(parts[1]), float(parts[2]), float(parts[3]), float(parts[4]) boxes.append((cls_id, x, y, w, h)) annotations[basename] = boxes return annotations

step 2：根据缺陷Id，统计每种缺陷的数量以及占所有缺陷的比例。并统计了每张图片缺陷的数量。

def analyze_class_distribution(annotations): """分析类别分布""" class_counts = Counter() boxes_per_image = [] for basename, boxes in annotations.items(): boxes_per_image.append(len(boxes)) for cls_id, *_ in boxes: class_counts[cls_id] += 1 print("=" * 60) print("类别分布统计") print("=" * 60) total = sum(class_counts.values()) for cls_id in sorted(class_counts.keys()): count = class_counts[cls_id] ratio = count / total * 100 bar = "#" * int(ratio) print(f" {CLASS_NAMES[cls_id]:30s} | {count:5d} ({ratio:5.1f}%) {bar}") print(f" {'总计':30s} | {total:5d}") print(f"\n图片数量: {len(annotations)}") print(f"每张图的标注框数: 最少{min(boxes_per_image)}, 最多{max(boxes_per_image)}, " f"平均{np.mean(boxes_per_image):.1f}") print() return class_counts, boxes_per_image

step 3：分析所有标注框的宽度，高度和面积。

def analyze_box_sizes(annotations): """分析标注框尺寸分布""" widths = [] heights = [] areas = [] aspect_ratios = [] class_areas = defaultdict(list) for basename, boxes in annotations.items(): for cls_id, x, y, w, h in boxes: widths.append(w) heights.append(h) area = w * h areas.append(area) aspect_ratios.append(w / max(h, 1e-6)) #避免除0 class_areas[cls_id].append(area) print("=" * 60) print("标注框尺寸分析 (归一化坐标)") print("=" * 60) print(f" 宽度: 最小{min(widths):.3f}, 最大{max(widths):.3f}, 平均{np.mean(widths):.3f}") print(f" 高度: 最小{min(heights):.3f}, 最大{max(heights):.3f}, 平均{np.mean(heights):.3f}") print(f" 面积: 最小{min(areas):.4f}, 最大{max(areas):.4f}, 平均{np.mean(areas):.4f}") print() return widths, heights, areas, aspect_ratios, class_areas

step 4：统计出最多数量的缺陷和最少数量的缺陷，判断不均衡的比例。

def check_class_imbalance(class_counts): """评估类别不均衡程度""" counts = list(class_counts.values()) max_count = max(counts) min_count = min(counts) ratio = max_count / max(min_count, 1) print("=" * 60) print("类别不均衡评估") print("=" * 60) print(f" 最多类别数量: {max_count}") print(f" 最少类别数量: {min_count}") print(f" 不均衡比例: {ratio:.1f}:1") if ratio > 5: print(" [!] 类别严重不均衡, 建议采取以下措施:") print(" 1. 对少数类别进行过采样/数据增强") print(" 2. 使用类别权重(class weights)") print(" 3. 使用Focal Loss") elif ratio > 2: print(" [!] 存在一定的类别不均衡, 建议关注少数类别的召回率") else: print(" [OK] 类别分布较均衡") print()

结果输出：

除此之外，根据上述函数返回的值，可以将这些数据绘制成可视化的统计图：

如果经过上述的分析发现缺陷数量不均衡，就需要补充相对应的缺陷。

在目标检测模型需要预测图上哪里有目标、目标多大。 => 如果让模型从零开始预测绝对坐标，难度很高。Anchor的思路是：先在图上铺一组预设的参考框，模型只需预测目标相对于参考框的偏移量。上面的analyze_box_sizes函数分析了各个缺陷的长度，宽度和长宽比等信息，这些信息可以影响anchor预测参考框的尺寸。

假设图上有个30x50 的缺陷：

方式A (无anchor)：模型直接预测x=120, y=80, w=30, h=50 --> 数值跨度大，学习困难

方式B (有anchor)：预设一个32x48 的参考框，模型只需预测dx=-2, dy=+1, dw=+2 --> 小偏移量， 容易学

二、YOLO模型训练与调优

1）数据增强优化

数据增强是在不新增真实采集图片的前提下，通过图像变换算法对原始图片做各类修改，模拟现实中不同拍摄条件、物体形态，生成大量 “虚拟样本”。如果不用数据增强的话，我们的训练数据就只有我们给定的1400张。如果使用增强方法，那么图片可能被处理，使我们参与训练的图片增加很多。

以下是训练的核心代码，其实这个代码在yolo模型介绍（二）中已经见过。数据增强分为基础增强和高级增强。我把各个参数标出来，大家可以看下：

model = YOLO(YOLOV12_YAML) # 训练参数 train_kwargs = dict( data=dataset_yaml, epochs=epochs, batch=batch, imgsz=args.imgsz, patience=args.patience if args.patience else (50 if not is_cpu else epochs), device=device, workers=workers, pretrained=False, save=True, verbose=True, plots=False, project=os.path.join(BASE_DIR, "runs"), name="steel_train", exist_ok=True, # 高级增强: YOLO核心技术 mosaic=1.0, # Mosaic: 4张图拼成1张, 增加小目标检测能力 mixup=0.1, # MixUp: 两张图混合, 增加模型泛化能力 copy_paste=0.1, # Copy-Paste: 将目标复制粘贴到其他图片, 适合目标较少的场景 # erasing=0.3, # 擦除: 随机遮挡部分区域, 增加鲁棒性 # 几何变换: 缺陷方向不固定 scale=0.5, # 缩放: 模拟不同距离拍摄 fliplr=0.5, # 水平翻转: 缺陷可能出现在任何位置 # flipud=0.5, # 垂直翻转: 钢铁表面缺陷方向不固定 # degrees=15.0, # 旋转: 缺陷角度多变 # translate=0.15, # 平移: 模拟相机位置变化 # 颜色变换: 模拟光照和材料差异，HSV色调增强 # hsv_h=0.02, # hsv_s=0.7, # hsv_v=0.5, # perspective=0.001, # 透视变换: 轻微的视角变化 # shear=5.0, # 剪切变换 # cos_lr=True, # 余弦退火学习率, 训练更稳定 ) results = model.train(**train_kwargs)

调优思路：

1. 缺陷可能出现在任何方向 -> 增加旋转和翻转

2. 缺陷大小差异大 -> 增加多尺度训练

3. 光照不均匀 -> 增加HSV增强

4. 数据量有限 -> 加强Mosaic和MixUp

2）模型选择和图片选择

调优思路：

1.选用更大的模型，如yolov12x.pt

2.imgsz选用更大尺寸，cos_lr=True

权衡: 这种调优方法精度更高, 但速度更慢、显存需求更大（需要减小batch的大小）

3）离线数据增强 (扩充训练集)

Albumentations是专为计算机视觉设计的高性能图像增强库，支持边界框、分割掩膜同步变换，比PIL 快10-100倍。提供300+种变换（模糊、畸变、天气模拟等），是工业检测离线增强的首选工具。这种方法和上面的两种方法相比：

- YOLO内置增强: 每个epoch随机变换, 不改变数据集大小

- 离线增强: 物理扩充数据集, 增加数据量

这二者可以结合使用

def offline_augmentation(images_dir, labels_dir, output_dir, multiply=3): try: import cv2 import albumentations as A except ImportError: print("[!] 需要安装: pip install albumentations opencv-python") return out_images = os.path.join(output_dir, "images") out_labels = os.path.join(output_dir, "labels") os.makedirs(out_images, exist_ok=True) os.makedirs(out_labels, exist_ok=True) # 针对钢铁表面缺陷设计的增强管道 transform = A.Compose([ A.HorizontalFlip(p=0.5), A.VerticalFlip(p=0.5), A.RandomRotate90(p=0.5), # CLAHE: 对比度受限自适应直方图均衡化, 增强缺陷可见性 A.CLAHE(clip_limit=4.0, tile_grid_size=(8, 8), p=0.3), A.RandomBrightnessContrast(brightness_limit=0.2, contrast_limit=0.2, p=0.5), A.GaussNoise(var_limit=(10.0, 50.0), p=0.2), A.Blur(blur_limit=3, p=0.1), A.ShiftScaleRotate(shift_limit=0.1, scale_limit=0.15, rotate_limit=30, border_mode=cv2.BORDER_REFLECT, p=0.5), ], bbox_params=A.BboxParams( format='yolo', label_fields=['class_labels'], min_visibility=0.3, )) image_files = [f for f in os.listdir(images_dir) if f.endswith(('.jpg', '.png'))] total_generated = 0 for img_file in image_files: img_path = os.path.join(images_dir, img_file) label_path = os.path.join(labels_dir, os.path.splitext(img_file)[0] + ".txt") if not os.path.exists(label_path): continue image = cv2.imread(img_path) if image is None: continue image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) bboxes = [] class_labels = [] with open(label_path, "r") as f: for line in f: parts = line.strip().split() if len(parts) == 5: class_labels.append(int(parts[0])) bboxes.append([float(x) for x in parts[1:5]]) # 复制原图 basename = os.path.splitext(img_file)[0] shutil.copy2(img_path, os.path.join(out_images, img_file)) shutil.copy2(label_path, os.path.join(out_labels, basename + ".txt")) # 生成增强副本 for i in range(multiply): try: augmented = transform(image=image, bboxes=bboxes, class_labels=class_labels) aug_image = augmented['image'] aug_bboxes = augmented['bboxes'] aug_labels = augmented['class_labels'] if len(aug_bboxes) == 0: continue aug_name = f"{basename}_aug{i}" aug_img = cv2.cvtColor(aug_image, cv2.COLOR_RGB2BGR) cv2.imwrite(os.path.join(out_images, f"{aug_name}.jpg"), aug_img) with open(os.path.join(out_labels, f"{aug_name}.txt"), "w") as f: for cls_id, bbox in zip(aug_labels, aug_bboxes): f.write(f"{cls_id} {bbox[0]:.6f} {bbox[1]:.6f} {bbox[2]:.6f} {bbox[3]:.6f}\n") total_generated += 1 except Exception as e: continue print(f"原始图片: {len(image_files)} 张") print(f"新增图片: {total_generated} 张") print(f"输出目录: {output_dir}")

4）推理时增强TTA

测试时增强 (Test Time Augmentation)对同一张图做多次不同变换, 取所有预测的并集。优点是不需要重新训练, 直接提升推理精度；缺点是推理速度变慢(约2-3倍)

results_tta = model.predict( source=test_images_dir, conf=0.25, iou=0.45, augment=True, # 开启TTA save=True, project=os.path.join(BASE_DIR, "runs_enhanced"), name="predict_tta", )