4. 焊缝质量检测

任务目标

构建一个基于YOLOv8的焊缝缺陷检测程序，能够从焊缝X射线图像中识别并定位气孔（porosity）、裂纹（crack）和未熔合（lack_of_fusion）三类缺陷，输出带边界框标注的结果图像。

环境依赖

• Python 3.9+
• ultralytics（YOLOv8）
• opencv-python
• matplotlib（可选，用于显示结果）

安装命令：pip install ultralytics opencv-python

数据集准备

使用GDXray焊缝X射线公开数据集子集（或自备标注数据），目录结构如下：

dataset/
├── images/train/   # 训练图像（jpg/png）
├── images/val/     # 验证图像
└── labels/train/   # YOLO格式txt标注（class x_center y_center width height）
    └── labels/val/

类别映射：0: porosity, 1: crack, 2: lack_of_fusion

完整实现代码

import cv2
from ultralytics import YOLO
import os

# ========== 1. 训练配置 ==========
def train_model():
    """在焊缝缺陷数据集上训练YOLOv8模型"""
    model = YOLO("yolov8n.pt")  # 使用预训练权重初始化
    
    model.train(
        data="weld.yaml",        # 数据集配置文件（见下方说明）
        epochs=100,
        imgsz=640,
        batch=16,
        patience=20,
        project="weld_defect",
        name="yolov8n_weld",
        device="cpu"             # 有GPU可改为"0"
    )
    print("训练完成，最佳权重保存在 weld_defect/yolov8n_weld/weights/best.pt")

# ========== 2. 推理检测 ==========
def detect_weld_defects(image_path, weights_path="weld_defect/yolov8n_weld/weights/best.pt"):
    """
    对单张焊缝图像执行缺陷检测
    返回：标注后的图像（BGR格式）和检测结果列表
    """
    # 加载训练好的模型
    model = YOLO(weights_path)
    
    # 读取原始图像
    img = cv2.imread(image_path)
    if img is None:
        raise FileNotFoundError(f"无法读取图像: {image_path}")
    
    # 执行推理
    results = model.predict(source=img, conf=0.25, iou=0.45, verbose=False)
    
    # 获取检测框信息
    boxes = results[0].boxes
    class_names = model.names  # {0: 'porosity', 1: 'crack', 2: 'lack_of_fusion'}
    
    # 在图像上绘制检测结果
    annotated_img = img.copy()
    defect_list = []
    
    for box in boxes:
        # 提取坐标和类别
        x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())
        cls_id = int(box.cls[0])
        conf = float(box.conf[0])
        label = f"{class_names[cls_id]} {conf:.2f}"
        
        # 根据类别选择颜色
        color_map = {
            0: (0, 255, 0),    # 气孔 - 绿色
            1: (0, 0, 255),    # 裂纹 - 红色
            2: (255, 0, 0)     # 未熔合 - 蓝色
        }
        color = color_map.get(cls_id, (255, 255, 0))
        
        # 绘制边界框和标签
        cv2.rectangle(annotated_img, (x1, y1), (x2, y2), color, 2)
        cv2.putText(annotated_img, label, (x1, y1 - 10),
                   cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)
        
        defect_list.append({
            "class": class_names[cls_id],
            "confidence": conf,
            "bbox": [x1, y1, x2, y2]
        })
    
    return annotated_img, defect_list

# ========== 3. 批量处理入口 ==========
if __name__ == "__main__":
    # 训练模式（首次运行取消注释）
    # train_model()
    
    # 推理模式
    test_image = "test_weld.jpg"  # 替换为实际图像路径
    result_img, defects = detect_weld_defects(test_image)
    
    # 保存结果
    output_path = "result_weld.jpg"
    cv2.imwrite(output_path, result_img)
    print(f"检测结果已保存至: {output_path}")
    
    # 打印检测到的缺陷
    if defects:
        print(f"共检测到 {len(defects)} 处缺陷:")
        for i, d in enumerate(defects, 1):
            print(f"  {i}. {d['class']} (置信度: {d['confidence']:.2f}) "
                  f"位置: {d['bbox']}")
    else:
        print("未检测到缺陷")
    
    # 显示结果（可选）
    # cv2.imshow("Weld Defect Detection", result_img)
    # cv2.waitKey(0)
    # cv2.destroyAllWindows()

数据集配置文件（weld.yaml）

path: dataset  # 数据集根目录
train: images/train
val: images/val

nc: 3  # 类别数量
names: ['porosity', 'crack', 'lack_of_fusion']

运行与验证

1. 训练模型：在代码中取消注释train_model()，运行python weld_detection.py，等待训练完成（CPU约需数小时，GPU约30分钟）。
2. 测试检测：将test_image路径改为实际焊缝图像，再次运行脚本。
3. 验证标准：
   • 输出图像中缺陷位置有彩色边界框标注，类别名称和置信度清晰可见；
   • 控制台打印的缺陷列表与图像标注一致；
   • 对已知含气孔的图像，绿色框应准确框住气孔区域；对含裂纹图像，红色框应沿裂纹走向覆盖缺陷。

常见问题

• 训练不收敛：检查标注文件是否与图像一一对应，类别ID是否在0-2范围内。
• 检测漏检严重：增加训练轮数（epochs=150）或使用更大模型（yolov8s.pt）。
• 误检过多：提高置信度阈值（conf=0.35~0.5）。