A095-segformer模型实现pcb缺陷检测

1. 项目简介

3. 数据集与预处理

4. 模型架构

1. 模型结构的逻辑 本项目采用了Segformer架构，这是一种基于Transformer的语义分割模型，结合了轻量化和高效性的设计。Segformer模型通过高效的混合Transformer编码器和解码器实现像素级语义分割。编码器利用Transformer模块捕获全局上下文信息，同时保留图像的细节特征，并通过多尺度特征聚合提高分割精度。具体来说，本项目加载了NVIDIA发布的“mit-b0”预训练模型，以充分利用大规模预训练权重的优势，并进行了模型微调以适应工业缺陷分割任务。模型在输出层映射到多类标签，用于预测不同的缺陷类别（如短路、虚焊和错误安装等）。
2. 模型的整体训练流程 模型训练分为数据加载、模型构建与配置、训练与评估几个主要步骤。首先，创建自定义的数据集类DefectSegmentationDataset，从图像路径和XML标注文件中加载图像和对应的缺陷mask，并通过SegformerImageProcessor进行预处理。然后，加载预训练的Segformer模型并设置训练参数，如学习率（0.00001）、训练周期（50轮）、批量大小（1）等。通过Hugging Face的TrainingArguments进行训练配置，使用步骤（steps）作为评估和保存策略，以在每隔指定的步数后评估模型性能和保存检查点。

5. 核心代码详细讲解

1. 数据预处理与特征工程

def process_mask(img_path, xml_path, visualize=False):# 读取图像，并初始化一个空的掩码（mask），大小与图像一致
    img = cv2.imread(img_path)
    # 如果visualize=True，则创建一个3通道的RGB掩码，否则创建单通道掩码
    num_dim = 3 if visualize else 1
    mask = np.zeros((img.shape[0], img.shape[1], num_dim))
    # 读取XML文件内容并解析为字典格式with open(xml_path, "r") as file:
        xml_content = file.read()
    data = xmltodict.parse(xml_content)
# 检查XML标注中是否包含缺陷信息if "object" in data["annotation"]:
        objects = data["annotation"]["object"]
# 如果只有一个缺陷对象，则将其转换为列表格式，以便统一处理if not isinstance(objects, List):
            objects = [objects]
# 遍历每个缺陷对象，获取缺陷类型和多边形坐标for obj in objects:
            defect_type = obj["name"]
            polygon = obj["polygon"]
            poly_keys = list(polygon.keys())
# 解析多边形坐标点的数量
            poly_keys = [int(k[1:]) for k in poly_keys]
            num_poly_points = max(poly_keys)
# 将多边形的坐标点依次提取并存入数组
            poly_coords = []for i in range(1, num_poly_points + 1):
                poly_coords.append([int(float(polygon[f"x{i}"])),int(float(polygon[f"y{i}"]))
                ])
            poly_coords = np.array(poly_coords)
# 使用OpenCV在掩码中填充多边形区域
            fill_color = color_map[defect_type] if visualize else id_cat_map[defect_type]
            mask = cv2.fillPoly(mask, pts=[poly_coords], color=fill_color)
    # 如果是可视化模式，保存掩码图像if visualize:
        cv2.imwrite("output.jpg", mask)
        mask = Image.open("output.jpg")
    return mask

• img = cv2.imread(img_path): 使用OpenCV读取图像文件，并存入img变量中。
• mask = np.zeros((img.shape[0], img.shape[1], num_dim)): 初始化一个空的掩码数组，与图像的大小相同。num_dim为3时创建RGB掩码，为1时创建单通道掩码。
• xmltodict.parse(xml_content): 解析XML格式的标注文件，将其转换为字典格式，方便后续的数据提取。
• if "object" in data["annotation"]: 检查标注文件中是否存在缺陷对象。
• if not isinstance(objects, List): 处理单个缺陷对象，将其包装为列表，便于统一遍历处理。
• for obj in objects: 遍历所有缺陷对象，提取缺陷类型和多边形坐标。
• poly_coords.append([...]): 将多边形的每个坐标点提取并存入poly_coords数组。
• mask = cv2.fillPoly(...): 使用OpenCV的fillPoly函数，在掩码上绘制缺陷多边形区域。
• cv2.imwrite("output.jpg", mask): 如果启用了可视化模式，将掩码图像保存并加载为PIL图像进行展示。

2. 模型架构构建

id2label = {0: "background", 1: "Short_circuit", 2: "Dry_joint", 3: "Incorrect_installation"}
label2id = {label: id for id, label in id2label.items()}
num_labels = len(id2label)

# 加载预训练的Segformer模型
pretrained_model_name = "nvidia/mit-b0"
model = SegformerForSemanticSegmentation.from_pretrained(
    pretrained_model_name,
    id2label=id2label,
    label2id=label2id
)

• id2label 和 label2id: 建立类别标签和类别ID之间的映射，确保模型可以正确地将预测输出转换为实际类别标签。
• num_labels = len(id2label): 计算总的类别数，用于定义模型的输出维度。
• pretrained_model_name = "nvidia/mit-b0": 指定使用的预训练模型名称，这里选择了Segformer的MIT-B0版本，轻量且高效。
• SegformerForSemanticSegmentation.from_pretrained(...): 从Hugging Face Transformers库中加载预训练的Segformer模型，并配置类别标签映射，以适应语义分割任务。

3. 模型训练与评估

# 设置训练参数
epochs = 50
lr = 0.00001
batch_size = 1# 创建训练参数对象
training_args = TrainingArguments("segformer-b0-finetuned-pcb-outputs",
    learning_rate=lr,
    num_train_epochs=epochs,
    per_device_train_batch_size=batch_size,
    per_device_eval_batch_size=batch_size,
    save_total_limit=3,
    evaluation_strategy="steps",
    save_strategy="steps",
    save_steps=500,
    eval_steps=250,
    logging_steps=1,
    eval_accumulation_steps=5,
    load_best_model_at_end=True,
)

• epochs = 50: 设置训练轮数为50次，以确保模型有足够的时间进行优化。
• lr = 0.00001: 设置学习率为0.00001，控制模型的参数更新步伐，防止训练过程中出现梯度爆炸或梯度消失。
• batch_size = 1: 设置每个设备上的批量大小为1，适应高分辨率图像分割任务的内存需求。
• TrainingArguments(...): 定义训练过程的参数，包括学习率、训练轮数、保存策略、评估策略、日志记录频率等。
• save_total_limit=3: 限制保存的模型检查点数量，以节省存储空间。

metric = evaluate.load("mean_iou")

def compute_metrics(eval_pred):with torch.no_grad():
        logits, labels = eval_pred
        logits_tensor = torch.from_numpy(logits)
# 调整预测输出的大小与标签尺寸一致
        logits_tensor = nn.functional.interpolate(
            logits_tensor,
            size=labels.shape[-2:],
            mode="bilinear",
            align_corners=False,
        ).argmax(dim=1)
        pred_labels = logits_tensor.detach().cpu().numpy()
# 使用Mean IoU作为评估指标
        metrics = metric._compute(
            predictions=pred_labels,
            references=labels,
            num_labels=len(id2label),
            ignore_index=None,
            reduce_labels=preprocessor.do_reduce_labels,
        )
# 提取每个类别的精度和IoU，并更新到metrics字典中
        per_category_accuracy = metrics.pop("per_category_accuracy").tolist()
        per_category_iou = metrics.pop("per_category_iou").tolist()
        metrics.update({f"accuracy_{id2label[i]}": v for i, v in enumerate(per_category_accuracy)})
        metrics.update({f"iou_{id2label[i]}": v for i, v in enumerate(per_category_iou)})
return metrics

• metric = evaluate.load("mean_iou"): 加载Mean IoU评估指标，用于语义分割任务。
• compute_metrics(eval_pred): 定义一个函数来计算评估指标，接收模型的预测结果和标签作为输入。
• logits_tensor = nn.functional.interpolate(...): 使用双线性插值调整预测输出的大小，使其与真实标签的尺寸一致。
• logits_tensor.argmax(dim=1): 获取每个像素的类别预测索引。
• metrics = metric._compute(...): 计算Mean IoU及每个类别的精度和IoU，并返回结果。

6. 模型优缺点评价