📚 课程目录第4章基于深度学习的焊缝缺陷检测

4. 基于深度学习的焊缝缺陷检测

任务目标

构建并训练一个CNN模型，对焊缝X射线图像中的气孔、裂纹两类缺陷进行分类与定位。分类指判断图像中存在哪类缺陷；定位指以边界框形式标出缺陷在图像中的位置。本案例使用公开的GDXray焊缝数据集子集，实现从图像预处理到模型训练、推理与可视化验证的完整闭环。

环境与依赖

Python ≥ 3.8
torch ≥ 1.12
torchvision ≥ 0.13
opencv-python ≥ 4.5
numpy ≥ 1.21
matplotlib ≥ 3.5
tqdm ≥ 4.60

数据集准备

从GDXray焊缝子集下载包含气孔（porosity）与裂纹（crack）标注的图片。将图片放入data/images/，标注文件（VOC XML格式）放入data/annotations/。本案例代码自动解析XML中的边界框坐标与类别名。

完整实现代码

以下代码实现：数据加载与增强、基于Faster R-CNN的缺陷检测模型训练、推理与边界框可视化。模型使用在COCO上预训练的ResNet-50骨干，迁移学习适配两类焊缝缺陷。

import os, cv2, torch, numpy as np
import xml.etree.ElementTree as ET
from torch.utils.data import Dataset, DataLoader
from torchvision.models.detection import fasterrcnn_resnet50_fpn
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision import transforms as T
import matplotlib.pyplot as plt
from tqdm import tqdm

CLASSES = ['__background__', 'porosity', 'crack']

class WeldDataset(Dataset):
    def __init__(self, img_dir, ann_dir, transforms=None):
        self.img_dir, self.ann_dir = img_dir, ann_dir
        self.transforms = transforms
        self.imgs = [f for f in os.listdir(img_dir) if f.endswith('.jpg')]

    def __len__(self):
        return len(self.imgs)

    def __getitem__(self, idx):
        img_name = self.imgs[idx]
        img_path = os.path.join(self.img_dir, img_name)
        img = cv2.imread(img_path)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        h, w = img.shape[:2]

        xml_path = os.path.join(self.ann_dir, img_name.replace('.jpg', '.xml'))
        tree = ET.parse(xml_path)
        boxes, labels = [], []
        for obj in tree.findall('object'):
            name = obj.find('name').text
            if name not in CLASSES:
                continue
            bnd = obj.find('bndbox')
            xmin = float(bnd.find('xmin').text)
            ymin = float(bnd.find('ymin').text)
            xmax = float(bnd.find('xmax').text)
            ymax = float(bnd.find('ymax').text)
            boxes.append([xmin, ymin, xmax, ymax])
            labels.append(CLASSES.index(name))

        boxes = torch.as_tensor(boxes, dtype=torch.float32)
        labels = torch.as_tensor(labels, dtype=torch.int64)
        image = T.ToTensor()(img)

        if self.transforms:
            image = self.transforms(image)

        return image, {'boxes': boxes, 'labels': labels}

def collate_fn(batch):
    return tuple(zip(*batch))

def get_model(num_classes):
    model = fasterrcnn_resnet50_fpn(pretrained=True)
    in_features = model.roi_heads.box_predictor.cls_score.in_features
    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
    return model

def train_one_epoch(model, loader, optimizer, device):
    model.train()
    total_loss = 0
    for images, targets in tqdm(loader, desc='Training'):
        images = [img.to(device) for img in images]
        targets = [{k: v.to(device) for k, v in t.items()} for t in targets]
        loss_dict = model(images, targets)
        loss = sum(loss for loss in loss_dict.values())
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    return total_loss / len(loader)

def visualize_prediction(model, img_path, device, conf_thresh=0.5):
    model.eval()
    img = cv2.imread(img_path)
    img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img_tensor = T.ToTensor()(img_rgb).unsqueeze(0).to(device)

    with torch.no_grad():
        pred = model(img_tensor)[0]

    for box, label, score in zip(pred['boxes'], pred['labels'], pred['scores']):
        if score < conf_thresh:
            continue
        x1, y1, x2, y2 = box.cpu().numpy().astype(int)
        cls_name = CLASSES[label]
        color = (255, 0, 0) if cls_name == 'porosity' else (0, 255, 0)
        cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)
        cv2.putText(img, f'{cls_name} {score:.2f}', (x1, y1-5),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1)

    plt.figure(figsize=(10, 8))
    plt.imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
    plt.axis('off')
    plt.title('Defect Detection Result')
    plt.show()

if __name__ == '__main__':
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    dataset = WeldDataset('data/images', 'data/annotations')
    loader = DataLoader(dataset, batch_size=2, shuffle=True, collate_fn=collate_fn)

    model = get_model(len(CLASSES)).to(device)
    optimizer = torch.optim.SGD(model.parameters(), lr=0.005, momentum=0.9, weight_decay=0.0005)

    num_epochs = 10
    for epoch in range(num_epochs):
        avg_loss = train_one_epoch(model, loader, optimizer, device)
        print(f'Epoch {epoch+1}/{num_epochs}, Avg Loss: {avg_loss:.4f}')

    torch.save(model.state_dict(), 'weld_defect_model.pth')
    print('Model saved to weld_defect_model.pth')

    # 推理验证
    test_img = 'data/images/sample_test.jpg'
    if os.path.exists(test_img):
        visualize_prediction(model, test_img, device)

运行方式

# 安装依赖
pip install torch torchvision opencv-python numpy matplotlib tqdm

# 准备数据后直接运行
python weld_defect_detection.py

验证方法与判定标准

训练过程中观察loss值持续下降至接近收敛（通常低于0.1）。训练完成后，程序自动对sample_test.jpg执行推理并显示可视化结果。判定通过的客观标准：

模型能在测试图像上以≥0.5置信度检出缺陷，且边界框与真实缺陷位置重叠度（IoU）≥0.5；
类别标签正确区分气孔（蓝色框）与裂纹（绿色框）；
若数据集中有标注的验证集，可计算mAP@0.5，达到≥0.7视为通过。

常见问题

CUDA内存不足：降低batch_size至1，或减小图像输入尺寸。
XML解析错误：确认标注文件与图片文件名一一对应，且类别名严格为porosity或crack。
检测效果差：增加训练轮次至20-30，或使用数据增强（随机翻转、亮度调整）。