基于Segformer实现PCB缺陷检测(步骤 + 代码)

news2025/4/17 9:04:50

导读

本文主要介绍基于Segformer实现PCB缺陷检测，并给出步骤和代码。

背景介绍

PCB缺陷检测是电子制造的一个重要方面。利用Segformer等先进模型不仅可以提高准确性，还可以大大减少检测时间。传统方法涉及手动检查，无法扩展且容易出错。利用机器学习，特别是 Segformer模型，提供自动化且精确的解决方案。

实现步骤

下面是具体步骤：

【1】安装所需环境。首先，我们安装所需的库。其中，albumentations用于数据增强，transformers允许访问 Segformer等预训练模型，并xmltodict帮助解析数据集的XML注释。

pip install evaluate albumentations transformers accelerate xmltodict

【2】数据集。这个项目中使用的数据集由Roboflow提供。可以从下面链接获取：

https://universe.roboflow.com/diplom-qz7q6/defects-2q87r/dataset/16

该数据集分为测试文件夹和训练文件夹，由XML格式的图像及其相应注释组成。

# Create train and test setstrain_folder = "drive/..../train/images/"test_folder = "drive/.../validation/images/"
train_img_paths = sorted([train_folder + f for f in os.listdir(train_folder) if f.endswith("jpg")])test_img_paths = sorted([test_folder + f for f in os.listdir(test_folder) if f.endswith("jpg")])train_xml_paths = [f[:-3] + "xml" for f in train_img_paths]test_xml_paths = [f[:-3] + "xml" for f in test_img_paths]
train_ds = {"image_paths": train_img_paths, "xml_paths": train_xml_paths}test_ds = {"image_paths": test_img_paths, "xml_paths": test_xml_paths}

XML文件包含多边形标注信息，指示PCB图像上缺陷的位置。该函数process_mask读取XML标注信息并将其转换为掩码（类似图像的数组）。该掩模对应于PCB图像上的缺陷区域，基本上将缺陷与电路板的其余部分分开。

该函数首先使用OpenCV读取原始图像。在我们初始化与输入图像大小相同的蒙版之后。根据可视化标志，蒙版可以是3通道RGB蒙版（如彩色图像）或1通道灰度蒙版。最初，该掩码中的所有值都设置为零，这意味着没有缺陷。

【3】缺陷标注解析。对于每个缺陷标注信息：识别缺陷类型，提取缺陷的多边形形状，该多边形被绘制到初始化的蒙版上。最后，该函数弥合了XML标注信息和适合训练的格式之间的差距。给定 PCB图像及其相应的XML 注释，它会生成一个分割掩模，突出显示有缺陷的区域。掩模可以是适合训练模型的数字格式，也可以是用于人工检查的视觉格式。

def process_mask(img_path, xml_path, visualize=False):    img = cv2.imread(img_path)    num_dim = 3 if visualize else 1    mask = np.zeros((img.shape[0], img.shape[1], num_dim))
    # Read xml content from the file    with open(xml_path, "r") as file:      xml_content = file.read()
    data = xmltodict.parse(xml_content)
    # If has defect mask    if "object" in data["annotation"]:        objects = data["annotation"]["object"]
        # Single defects are annotated as a single dict, not a list        if not isinstance(objects, List):          objects = [objects]
        for obj in objects:          defect_type = obj["name"]          polygon = obj["polygon"]          poly_keys = list(polygon.keys())
          # Get number of (x, y) pairs - polygon coords          poly_keys = [int(k[1:]) for k in poly_keys]          num_poly_points = max(poly_keys)
          # Parse ordered polygon coordinates          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)
          # Draw defect segment on mask          fill_color = color_map[defect_type] if visualize else id_cat_map[defect_type]          mask = cv2.fillPoly(mask, pts=[poly_coords], color=fill_color)
    #Optional    if visualize:        cv2.imwrite("output.jpg", mask)        mask = Image.open("output.jpg")
    return mask

【4】探索性数据分析。在训练模型之前，最好先了解数据。在这里，我们检查缺陷类型的分布并在样本图像上可视化缺陷。

缺陷热力图显示了常见的缺陷位置，箱线图显示了缺陷尺寸的分布。

该函数旨在通过读取边界框详细信息来计算 XML 注释中存在的每个缺陷的大小。

def get_defect_sizes(xml_paths):    sizes = []    for xml_path in xml_paths:        with open(xml_path) as f:            data = xmltodict.parse(f.read())            objects = []            if 'object' in data['annotation']:                objects = data['annotation']['object']            if not isinstance(objects, list):                objects = [objects]            for obj in objects:                bndbox = obj['bndbox']                width = int(bndbox['xmax']) - int(bndbox['xmin'])                height = int(bndbox['ymax']) - int(bndbox['ymin'])                sizes.append(width * height)    return sizes

最后，群图重点关注缺陷尺寸在整个数据集中的分布和扩散。

【5】数据增强。该albumentations库用于人为扩展训练数据集，有助于提高模型的泛化能力。唯一指定的增强是水平翻转，它将以 50% 的概率水平翻转图像。

transform = A.Compose([    A.HorizontalFlip(p=0.5)])

【6】图像预处理。将图像及其掩模预处理为适合Segformer模型的格式。

preprocessor = SegformerImageProcessor()

我们将定义一个继承自PyTorchDataset类的自定义数据集类。这个自定义数据集允许我加载和预处理 PCB 图像及其相应的缺陷掩模。

使用 OpenCV 加载图像。使用前面讨论的函数生成缺陷掩模process_mask。使用之前初始化的图像预处理图像及其掩模SegformerImageProcessor。此步骤将图像转换为张量格式，并确保它们具有适合 Segformer 模型的大小和标准化。返回预处理的图像和掩模张量。

class DefectSegmentationDataset(Dataset):  def __init__(self, dataset, mode):    self.image_paths = dataset["image_paths"]    self.xml_paths = dataset["xml_paths"]
  def __len__(self):    return len(self.image_paths)
  def __getitem__(self, idx):    # Read image    image = cv2.imread(self.image_paths[idx])
    # Get mask    mask = process_mask(self.image_paths[idx], self.xml_paths[idx])    mask = mask.squeeze()    mask = Image.fromarray(mask.astype("uint8"), "L")
    # Return preprocessed inputs    inputs = preprocessor(image, mask, ignore_index=None, return_tensors="pt")    inputs["pixel_values"] = inputs["pixel_values"].squeeze(0)    inputs["labels"] = inputs["labels"].squeeze(0)    return inputs