📚 课程目录第2章 PCB板缺陷检测与分类

2. PCB板缺陷检测与分类

任务说明

本案例使用YOLOv8训练一个PCB板缺陷检测模型，识别三类常见缺陷：短路（short）、断路（open）、缺件（missing）。输入为PCB板图像，输出为每个缺陷的边界框位置与类别标签。

环境依赖

pip install ultralytics opencv-python

数据集准备

使用公开PCB缺陷数据集（如北京大学PCB_DATASET），按YOLO格式组织：

datasets/pcb/
├── images/train/   # 训练图像（.jpg）
├── images/val/     # 验证图像
├── labels/train/   # 对应YOLO格式标签（.txt）
└── labels/val/

每张图像的标签文件内容格式：class_id x_center y_center width height（归一化坐标）。类别映射：0: short, 1: open, 2: missing。

完整实现代码

from ultralytics import YOLO
import yaml
import os

# 1. 准备数据配置文件
data_config = {
    'path': 'datasets/pcb',
    'train': 'images/train',
    'val': 'images/val',
    'names': {0: 'short', 1: 'open', 2: 'missing'}
}
with open('pcb_data.yaml', 'w') as f:
    yaml.dump(data_config, f, default_flow_style=False)

# 2. 加载预训练模型并训练
model = YOLO('yolov8n.pt')  # 使用轻量级预训练权重

# 3. 训练模型
results = model.train(
    data='pcb_data.yaml',
    epochs=100,
    imgsz=640,
    batch=16,
    device='cuda',  # 无GPU时改为'cpu'
    patience=20,
    project='pcb_defect_run',
    name='exp1',
    exist_ok=True
)

# 4. 验证模型性能
metrics = model.val(data='pcb_data.yaml')
print(f"mAP50: {metrics.box.map50:.4f}")
print(f"mAP50-95: {metrics.box.map:.4f}")

# 5. 推理测试（单张图像）
def predict_defects(image_path, conf_thres=0.25):
    """检测PCB图像中的缺陷，返回类别名称和坐标"""
    results = model.predict(
        source=image_path,
        conf=conf_thres,
        save=False,
        verbose=False
    )
    detections = []
    for r in results:
        boxes = r.boxes
        for box in boxes:
            cls_id = int(box.cls[0])
            conf = float(box.conf[0])
            x1, y1, x2, y2 = box.xyxy[0].tolist()
            detections.append({
                'class': model.names[cls_id],
                'confidence': round(conf, 4),
                'bbox': [round(v, 1) for v in [x1, y1, x2, y2]]
            })
    return detections

# 6. 保存最终模型
best_model_path = 'pcb_defect_run/exp1/weights/best.pt'
final_model = YOLO(best_model_path)
final_model.export(format='onnx')  # 可选：导出ONNX便于部署

if __name__ == '__main__':
    # 示例推理
    test_image = 'datasets/pcb/images/val/sample.jpg'
    if os.path.exists(test_image):
        dets = predict_defects(test_image)
        for d in dets:
            print(f"{d['class']} conf={d['confidence']} bbox={d['bbox']}")
    else:
        print("请将测试图像放在指定路径后运行推理")

运行方式

# 训练（约需GPU，CPU训练时间较长）
python train_pcb.py

# 推理单张图像
python train_pcb.py  # 脚本末尾自动执行示例推理

验证方法与通过标准

训练收敛验证：训练完成后检查 pcb_defect_run/exp1/results.csv，观察验证集mAP50是否随epoch上升并趋于稳定（通常应达到0.85以上，具体取决于数据集质量）。
定量评估：运行验证代码后，终端输出mAP50和mAP50-95指标。若mAP50 ≥ 0.85，mAP50-95 ≥ 0.55，判定模型性能合格。
定性检测：将测试图像放入指定路径，运行推理脚本，观察输出结果：

# 预期输出示例（实际数值取决于训练结果）
short conf=0.9213 bbox=[112.5, 84.2, 156.3, 102.8]
open conf=0.8876 bbox=[320.1, 210.7, 355.9, 228.4]
missing conf=0.7932 bbox=[501.0, 402.3, 530.6, 418.9]

检测结果应满足：每个缺陷的类别标签正确，边界框基本贴合缺陷区域，置信度高于0.5。

常见问题

类别不均衡：若某类缺陷样本极少，可增加该类别的图像增强权重，或使用class_weights参数。
小目标漏检：PCB缺陷通常较小，可将imgsz提升至800或使用yolov8m以上模型。
误检背景纹理：若将电路走线误检为缺陷，可提高conf_thres至0.35，或增加背景负样本。