PCB defect detection based on yolov5

1. Brief description

Automatic printed circuit board (PCB) defect detection based on image recognition technology has the advantages of high efficiency, accuracy, reliability, high degree of automation and strong scalability, which greatly improves the detection efficiency and saves labor costs. Traditional detection methods usually require a lot of manpower and time, and are easily affected by factors such as the environment, worker skills, and fatigue. However, the defect detection method based on image recognition can realize automatic defect detection without manual intervention and improve production efficiency.

In addition, image recognition technology can continuously improve its recognition ability through continuous learning and training, adapt to more types of defects, and has strong scalability.

2. Dataset introduction

Using the PCB dataset released by Peking University, it contains 1386 images and has 6 types of defects: leaks, rat bites, open circuits, short circuits, stray copper, and miscellaneous copper. Download address: Peking University Intelligent Robot Open Laboratory

3. Dataset Construction

First divide the data set into train, val, and the processing code  split_train_val.py is as follows:

# -*- coding: utf-8 -*-
import os
import random
import argparse

parser = argparse.ArgumentParser()

#xml文件的地址，根据自己的数据进行修改 xml一般存放在Annotations下
parser.add_argument('--xml_path', default='Annotations', type=str, help='input xml label path')

#数据集的划分，地址选择自己数据下的ImageSets/Main
parser.add_argument('--txt_path', default='ImageSets/Main', type=str, help='output txt label path')

opt = parser.parse_args()

trainval_percent = 0.9
train_percent = 0.8
xmlfilepath = opt.xml_path
txtsavepath = opt.txt_path
total_xml = os.listdir(xmlfilepath)
if not os.path.exists(txtsavepath):
    os.makedirs(txtsavepath)

num = len(total_xml)
list_index = range(num)
tv = int(num * trainval_percent)
tr = int(tv * train_percent)
trainval = random.sample(list_index, tv)
train = random.sample(trainval, tr)

file_trainval = open(txtsavepath + '/trainval.txt', 'w')
file_test = open(txtsavepath + '/test.txt', 'w')
file_train = open(txtsavepath + '/train.txt', 'w')
file_val = open(txtsavepath + '/val.txt', 'w')

for i in list_index:
    name = total_xml[i][:-4] + '\n'
    if i in trainval:
        file_trainval.write(name)
        if i in train:
            file_train.write(name)
        else:
            file_val.write(name)
    else:
        file_test.write(name)

file_trainval.close()
file_train.close()
file_val.close()
file_test.close()

Then convert the VOC label format to yolov5 format, the code is as follows:

# -*- coding: utf-8 -*-
import xml.etree.ElementTree as ET
import os
from os import getcwd

sets = ['train', 'val']
classes = ["missing_hole","mouse_bite","open_circuit","short","spur","spurious_copper"]  # 改成自己的类别
abs_path = os.getcwd()
print(abs_path)


def convert(size, box):
    dw = 1. / (size[0])
    dh = 1. / (size[1])
    x = (box[0] + box[1]) / 2.0 - 1
    y = (box[2] + box[3]) / 2.0 - 1
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh
    return x, y, w, h


def conver_annotation(image_id):
    in_file = open('Annotations/%s.xml' % (image_id), encoding='UTF-8')
    out_file = open('labels/%s.txt' % (image_id), 'w')
    tree = ET.parse(in_file)
    root = tree.getroot()
    size = root.find('size')
    w = int(size.find('width').text)
    h = int(size.find('height').text)
    for obj in root.iter('object'):
        difficult = obj.find('difficult').text
        #difficult = obj.find('Difficult').text
        cls = obj.find('name').text
        if cls not in classes or int(difficult) == 1:
            continue
        cls_id = classes.index(cls)
        xmlbox = obj.find('bndbox')
        b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text), float(xmlbox.find('ymax').text))
        b1, b2, b3, b4 = b

        # 标注越界修正
        if b2 > w:
            b2 = w
        if b4 > h:
            b4 = h
        b = (b1, b2, b3, b4)
        bb = convert((w, h), b)
        out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')


wd = getcwd()
for image_set in sets:
    if not os.path.exists('labels/'):
        os.makedirs('labels/')
    image_ids = open('ImageSets/Main/%s.txt' % (image_set)).read().strip().split()
    list_file = open('%s.txt' % (image_set), 'w')
    for image_id in image_ids:
        list_file.write(abs_path + '/images/%s.jpg\n' % (image_id))
        conver_annotation(image_id)
    list_file.close()

4. Model training

(1) Create a dataset configuration file

Create your own data set configuration file pcb.yaml in the yolov5/data/ directory, the content is as follows, nc means that the category is 6 categories

# PCB data

# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
train: ./data/PCB_DATASET/train.txt
val: ./data/PCB_DATASET/val.txt

nc: 6  # number of classes

names: ['missing_hole', 'mouse_bite', 'open_circuit', 'short', 'spur', 'spurious_copper']

(2) Create a model file

In the yolov5/models/ directory, copy a version model as your own model. Here I use yolov5s, and then change the number of categories in the file to 6:

cp yolov5s.yaml yolov5s_pcb.yaml

# Parameters
nc: 6  # number of classes
depth_multiple: 0.33  # model depth multiple
width_multiple: 0.50  # layer channel multiple
anchors:
  - [10,13, 16,30, 33,23]  # P3/8
  - [30,61, 62,45, 59,119]  # P4/16
  - [116,90, 156,198, 373,326]  # P5/32

(3) Model training

Execute the following command to train:

python train.py --img 640 --batch 16 --epochs 100 --data ./data/pcb.yaml --weights yolov5s.pt

After 100 epochs training, the effect is as follows:

[email protected] reached 0.953, the effect is still good.

(4) Experimental verification effect

 

 

---

reference:

PCB defect detection based on yolov5, introducing CVPR 2023 BiFormer: Vision Transformer with Bi-Level Routing Attention to improve detection accuracy - Programmer Sought