PCB defect detection method based on improved YOLOv7

Abstract

The present invention discloses a PCB defect detection method based on an improved YOLOv7. The method comprises the following steps: first, randomly dividing a PCB defect dataset into a training set, a validation set, and a test set; then constructing an improved YOLOv7 detection model; secondly, preprocessing the data in the training set; then inputting the preprocessed training set data and validation set data into the improved YOLOv7 detection model, and starting model training until the loss curve converges; finally, inputting the test set images into the trained improved YOLOv7 detection model to obtain recognition results for PCB defect images. The present invention relates to a PCB defect detection method based on an improved YOLOv7. The method utilizes an attention mechanism to enable the receptive field to focus more on defect features and ignore irrelevant features, and takes effective measures to reduce the size of the model. The method achieves higher accuracy, faster speed, and smaller model size for PCB defect detection, making it suitable for industrial inspection environments.

Description

TECHNICAL FIELD

[0001] The application relates to the technical field of image processing and target detection, and particularly relates to a PCB defect detection method based on an improved YOLOv7. BACKGROUND

[0002] A PCB (Printed Circuit Board) is an important basic component in modern electronic products, and the performance thereof is related to the function realization of the entire product, so the quality of the PCB needs to be strictly monitored in the production process to screen out PCBs with defects. Defects existing in the PCB mainly include a missing hole, a mouse bite, an open circuit, a short circuit, strayness and miscellaneous copper, etc. Currently, common detection means are manual visual detection, electrical detection and traditional visual detection, etc. However, with the progress of technology, the manufacturing process of the PCB gradually develops towards high performance, high integration and high complexity, and the traditional detection method is difficult to adapt to these changing trends, and is prone to false detection and missed detection problems.

[0003] PCB defect detection belongs to target detection, and target detection is an important task in the field of computer vision. With the continuous development and breakthrough of deep learning technology, the precision and speed of target detection have been greatly improved. Nowadays, the target detection technology based on deep learning has been practically applied in automatic driving, industrial defect detection, intelligent medical treatment and many other fields. The detection method based on deep learning can automatically extract the features of data, and the detection precision is very good under the condition that the data volume is sufficient, thereby providing a feasible scheme for solving the existing problems of PCB defect detection.

[0004] YOLO is one of the mainstream algorithms based on the deep learning concept, and is a high-performance single-stage detector. The application of YOLO to the PCB defect detection scene can take into account the precision and speed required in the industrial detection environment, and can greatly save the cost.

[0005] At present, there are few public literatures on the research on the PCB defect detection by using the YOLO algorithm. Meanwhile, with the continuous progress of deep learning technology, the YOLO algorithm is also constantly updated and replaced. The existing few public literatures on the PCB defect detection by using the YOLO algorithm all have the problems of low precision and slow speed. SUMMARY

[0006] The application aims to provide a PCB defect detection method based on an improved YOLOv7 to solve the above problems.

[0007] In order to achieve the above-mentioned purpose, the application provides the following technical scheme.

[0008] A PCB defect detection method based on an improved YOLOv7 comprises the following steps.

[0009] S1, divide the PCB defect data set containing six types of defects into a training set, a validation set and a test set, and the division process adopts a random form;

[0010] S2, construct an improved YOLOv7 detection model;

[0011] S3, pre-process the data in the training set, including the process of Mosaic data enhancement and the process of clustering algorithm for clustering the labeled target bounding box;

[0012] S4, input the pre-processed training set data and validation set data into the improved YOLOv7 detection model, start model training, and until the loss curve converges;

[0013] S5, input the test set image into the trained improved YOLOv7 detection model, and obtain the recognition result of the PCB defect image.

[0014] Preferably, in step S1, the six types of defects specifically include: pinholes, mouse bites, open circuits, short circuits, stray, and copper impurities.

[0015] Preferably, in step S1, the process of dividing the PCB defect data set specifically includes: dividing the PCB defect data set into a training set, a validation set and a test set in a random form, the division ratio is 8:1:1, the labels are put into the corresponding folder, and the files are stored in yolo format.

[0016] Preferably, in step S2, the improved YOLOv7 detection model is constructed, and the process specifically includes: selecting a SiLU version of YOLOv7 tiny network as an improved basis, and improving in the following manner:

[0017] S21, replace the high-resolution detection head: small targets are more difficult to extract than regular targets, and are easy to lose in deep-level extraction, therefore, the features of the first feature map are introduced into the head part of the network, and a high-resolution small target detection head is set, and a detection head with the lowest resolution is deleted to reduce the parameter amount and simplify the model;

[0018] S22, replace the ordinary 3*3 convolution in the feature extraction structure in the backbone with DCNV3: PCB defect detection needs both shallow and deep features, deformable convolution is a kind of convolution variant, which adds bias to convolution, so that convolution has the ability to deform, thereby obtaining multi-scale information through different receptive fields, which can meet the needs of PCB defect detection;

[0019] S23, select NWD as a new evaluation index and introduce it into the loss function, and use NWD Loss to improve the problem that IoU is sensitive to small target bbox offset;

[0020] S24, adding GAM in the head part of YOLOv7: GAM improves recognition accuracy by selectively focusing on the parts needed in the channel and space to extract relevant information, and is added in two places before the detection head.

[0021] Preferably, in step S3, the Mosaic data enhancement and target bounding box clustering are performed, and the specific process includes: first, four pictures are randomly selected from the data set, then the data is augmented, including flipping, scaling, color gamut change operation, and then the picture and frame combination is performed; based on the labeled target bounding box of the training set, 9 small to large anchor frames are obtained in advance by the K-means clustering algorithm.

[0022] Preferably, in step S4, the model training is performed, and the specific process includes:

[0023] S41, inputting the data enhanced training set picture into the Backbone network for feature extraction to obtain a feature map;

[0024] S42, inputting the feature map into the head network to fuse the features of different feature maps to obtain a detection frame;

[0025] S43, screening the detection frame according to a pre-set threshold to obtain a prediction frame;

[0026] S44, inputting the verification set picture into the improved YOLOv7 model to obtain a detection result according to the prediction frame;

[0027] S45, calculating the loss function of the improved YOLOv7 model in the training according to the detection result, and constantly adjusting the parameters of the model so that the loss function of the model is constantly reduced until the training is finally completed.

[0028] Preferably, the step S5 has the specific process that: the recognition effect of the model is verified, the test set picture data is preprocessed, the weight file obtained by training is input in a unified format, and the target recognized by the model is framed on the picture, that is, the PCB defect image recognition result is obtained.

[0029] The application has the beneficial effects that:

[0030] The PCB defect detection method based on the improved YOLOv7 optimizes the YOLOv7 network, replaces the high-resolution detection head, so as to solve the problem that the PCB defect target is small and difficult to detect; the ordinary 3*3 convolution in the feature extraction structure of the backbone is replaced by DCNV3 to extract multi-scale feature information and introduce rich shallow features into the head; the NWDLoss is used to improve the problem that the IoU is sensitive to the small target bbox offset; finally, the GAM is added to make the network pay more attention to the effective features. The PCB defect detection method based on the improved YOLOv7 enhances the model according to the characteristics of the PCB defect, uses the attention mechanism to make the receptive field pay more attention to the defect features and ignore irrelevant features, and takes effective measures to reduce the volume of the model, compared with the original YOLOv7 and other existing detection methods, the improved YOLOv7 has higher detection precision, faster speed and smaller model volume, and is suitable for industrial detection environment. BRIEF DESCRIPTION OF DRAWINGS

[0031] Figure 1 : schematic diagram of the D-ELAN structure designed in the embodiment of the application;

[0032] Figure 2 : structure schematic diagram of the GAM used in the embodiment of the application;

[0033] Figure 3 : overall structure diagram of the improved YOLOv7 algorithm designed in the embodiment of the application;

[0034] Figure 4 : training result diagram of the improved algorithm in the embodiment of the application;

[0035] Figure 5 : PCB defect detection effect diagram in the embodiment of the application. DETAILED DESCRIPTION

[0036] The application will be further described below in conjunction with the embodiments. It should be noted that the embodiments are only examples and illustrations of the concept of the application, and those skilled in the art can make various modifications, supplements or substitutions to the described embodiments or use similar ways to replace, as long as the modifications, supplements or substitutions do not deviate from the concept of the application or exceed the scope defined by the claims, and should be regarded as falling within the protection scope of the application.

[0037] Embodiment 1:

[0038] As shown in Figures 1-5 , a PCB defect detection method based on an improved YOLOv7 includes the following steps:

[0039] S1. Divide the PCB defect dataset containing six types of defects into training set, validation set and test set. The division process adopts random form.

[0040] Download the PCB defect dataset from the data sharing platform of the Intelligent Robotics Open Laboratory of Peking University. This dataset includes 693 images and corresponding labels, covering six defect categories: leaks, rat bites, opens, shorts, strays, and scrap copper. Create a Python program called split.py to randomly partition the PCB defect dataset into training, validation, and test sets with an 8:1:1 split ratio. This results in a training set of 555 images, a validation set of 69 images, and a test set of 69 images. Place the labels in the corresponding folders and save the files in YOLO format.

[0041] S2. Build an improved YOLOv7 detection model.

[0042] In step S2, an improved YOLOv7 detection model is constructed. The process specifically includes: selecting the SiLU version of the YOLOv7 tiny network as the improvement basis, and improving the network in the following way:

[0043] S21, Replace the high-resolution detection head:

[0044] The YOLOv7 tiny network consists of input, backbone, and head. The image is resized to 640*640 after input. After feature extraction, feature maps of 160*160*64, 80*80*128, 40*40*256, and 20*20*512 are obtained.

[0045] Small targets are more difficult to extract than regular targets and are easily lost in deep extraction. Therefore, for small target defects on PCBs, the features of the first feature map are introduced into the network head, and a high-resolution small target detection head is set up. At the same time, the lowest-resolution detection head is deleted to reduce the number of parameters and streamline the model. Specifically, the 160*160*64 feature map is combined with the upsampled feature map through a CBS module, and a 160*160 high-resolution detection head is added. The less effective 20*20 detection head is then deleted to reduce the number of parameters and improve the running speed. To do this, open the network configuration folder cfg, modify the yaml file that records the YOLOv7 tiny network structure, and create a new high-resolution detection head according to the above steps.

[0046] S22, replace the common 3*3 convolution in the feature extraction structure of the backbone with DCNV3 (Deformable Convolutional Networks): PCB defect detection requires both shallow and deep features. Deformable convolution is a variant of convolution that adds a bias to the convolution, allowing the convolution to deform and obtain multi-scale information through different receptive fields, which can meet the needs of PCB defect detection.

[0047] Deformable convolution can change its receptive field to obtain multi-scale feature information through different receptive fields, and combining shallow and deep features is beneficial to the positioning and recognition of small targets. Specifically, the deformable convolution mechanism code DCNV3 is added in the common.py file, and a judgment condition is added in the yolo.py file. Finally, the yaml file recording the YOLOv7tiny network structure is modified, the 3*3 convolution conv in the feature extraction structure is replaced with DCNV3, and the modified structure is named D-ELAN. The structure is specifically as shown in Figure 1 Figure 1 It can be seen that in the D-ELAN structure, multiple features are fused. When the features enter the D-ELAN structure, there are two routes. One is to directly adjust the feature channel through CBS to participate in fusion. The CBS module is composed of convolution, BN layer and SiLU activation function. The other is to pass through a CBS module and two DCNV3 modules, and introduce the features generated by the three modules into the fusion. A CBS module is used to finally adjust the channel of the fused features.

[0048] S23, select NWD as a new evaluation index and introduce it into the loss function, and use NWD Loss to improve the problem that IoU is sensitive to small target bbox offset.

[0049] IoU is very sensitive to the bbox offset of small targets, but it is limited by the bbox offset of medium and large targets, so IoU is not suitable for small target detection. In order to detect a large number of small target defects in PCB defects, NWD is introduced as a new evaluation index and introduced into the loss function, and NWD (Normalized Wasserstein Distance) Loss is used to improve the problem that IoU is sensitive to small target bbox offset.

[0050] NWD is an evaluation index suitable for small targets. The similarity between the predicted target and the real target is calculated by calculating the Gaussian distribution corresponding to them, that is, the Wasserstein distance between them is calculated according to formula (1):

[0051]

[0052] its​ C is a constant closely related to the data set.

[0053] Specific operation: open the loss.py file, add the calculation method of the Wasserstein distance in the code, and add the NWD index in the function of calculating the loss.

[0054] S24, add GAM to the head part of YOLOv7.

[0055] GAM (Global Attention Mechanism) is a global attention mechanism, and the structure of GAM is shown in Figure 2 GAM introduces a channel attention module with multiple perceptrons and a convolutional spatial attention module. The channel attention uses three-dimensional arrangement to retain three-dimensional information, which can amplify the cross-dimensional channel and spatial dependency. The spatial attention module can fuse spatial information. GAM attention mechanism can reduce information loss, improve global feature interaction, and enhance model performance.

[0056] Adding GAM in YOLOv7 network helps the network focus on key target information and ignore irrelevant target information. GAM extracts relevant information by selectively focusing on the required parts in the channel and space, thereby improving recognition accuracy. GAM is added at two positions before the detection head.

[0057] Specific operation: add the code to implement GAM in common.py, add the judgment condition in yolo.py file, modify the yaml file recording the structure of YOLOv7 tiny network, and add GAM at two positions before the detection head. The specific positions are marked in the overall network structure. The overall structure of the improved YOLOv7 network is shown in Figure 3 The improved YOLOv7 network extracts features by Backbone. First, the feature scale is adjusted by two CBS modules, which are composed of convolution, BN layer and SiLU activation function. Then, the D-ELAN module and the MP module are alternately used to extract features and downsample, obtaining feature maps of different sizes. After Backbone, pyramid pooling is used to improve the multi-scale expression ability of the model, and then multiple upsampling and feature fusion are performed to fuse the multi-scale features extracted by the above Backbone. Subsequently, a bottom-up feature fusion process is performed, and GAM attention mechanism is added in this part to reduce information loss and focus on useful features. Finally, three target detection heads are used to predict the target.

[0058] S3, pre-process the data in the training set, including the process of Mosaic data augmentation and the process of clustering algorithm for the labeled target bounding box.

[0059] Mosaic data augmentation: first randomly select four pictures in the dataset, then perform data augmentation operations including flipping, scaling, color gamut changes (changing brightness, saturation, hue, etc.), and then combine the pictures and the frame.

[0060] Target bounding box clustering: YOLOv7 will use the labeled target bounding box of the training set to obtain 9 anchor boxes of different sizes through the K-means clustering algorithm in advance.

[0061] S4, input the pre-processed training set data and validation set data into the improved YOLOv7 detection model, start model training, and stop until the loss curve converges.

[0062] Build an experimental environment: The experimental platform uses Ubuntu 20.04 as the operating system, PyTorch 1.11.0 as the framework, Nvidia RTX3090 (24G memory) as the GPU, Python version 3.8, and Cuda version 11.3.

[0063] Set the training parameters: batch size = 16, initial learning rate lr = 0.001, and training round number epoch = 200.

[0064] The model training has a specific process, which includes:

[0065] S41, input the data enhanced training set picture into the Backbone network for feature extraction to obtain a feature map;

[0066] S42, input the feature map into the head network to fuse the features of different feature maps to obtain a detection frame;

[0067] S43, filter the detection frame according to the pre-set threshold to obtain a prediction frame;

[0068] S44, input the validation set picture into the improved YOLOv7 model to obtain a detection result according to the prediction frame;

[0069] S45, calculate the loss function of the improved YOLOv7 model in training according to the detection result, and constantly adjust the parameters of the model to make the loss function of the model constantly decrease until the training is completed.

[0070] After starting the training, the loss curve gradually tends to be stable, and after completing 200 epochs, the precision, recall, and average precision mean mAP are used as indicators to evaluate the performance of the model. The calculation methods of P and R are shown in formulas (2) and (3):

[0071]

[0072]

[0073] Among them, TP means correct classification; FP means predicting a negative sample as a positive example; FN means predicting a positive sample as a negative example; TN means predicting a negative sample as a negative example.

[0074] With P as the ordinate and R as the abscissa, the area enclosed by the PR curve and the coordinate axis is the AP value, and mAP represents the mean AP of all categories in the data set.

[0075] During the training process, the P, R and mAP curves are as follows Figure 4 As shown in the figure, the P curve and R curve gradually rise during the training process, with slight fluctuations possibly due to the imbalance of training data, and then tend to be flat. The loss of the model also converges in this process. The P value is finally at 0.97, and the R value is finally at 0.9. The mAP of the model finally reaches 0.95, and the volume is only 6.7MB.

[0076] S5. Input the test set images into the trained improved YOLOv7 detection model to obtain the recognition results of the PCB defect images. The specific process is as follows:

[0077] To verify the recognition effect of the model, pre-process the test set image data, input the weight file obtained by training in a unified format, and frame the target recognized by the model on the image to obtain the PCB defect image recognition result. The partial PCB defect image recognition obtained in this embodiment is as follows Figure 5 As shown by Figure 5 It can be seen that the two types of defects, leaks and open circuits, which occupy a very small proportion of the image, are accurately marked, indicating that the improved model also has strong detection capabilities for small targets.

[0078] The present invention provides a PCB defect detection method based on improved YOLOv7. The YOLOv7 network is optimized and a high-resolution detection head is replaced to address the problem that PCB defect targets are small and difficult to detect. The ordinary 3*3 convolution in the feature extraction structure in the backbone is replaced with DCNV3 to extract multi-scale feature information and introduce rich shallow features into the head. NWDLoss is used to improve the problem that the IoU is sensitive to the bbox offset of small targets. Finally, GAM is added to make the network focus more on effective features.

[0079] The application discloses a PCB defect detection method based on improved YOLOv7, which enhances the model according to the characteristics of PCB defects, uses an attention mechanism to make the receptive field pay more attention to defect features and ignore irrelevant features, and takes effective measures to reduce the volume of the model.

[0080] The above is an exemplary description of the application, and obviously, the specific implementation of the application is not limited by the above method. As long as the method concept and technical solution of the application are used for non-essential improvement or the concept and technical solution of the application are directly applied to other occasions without improvement, they are within the protection scope of the application.

Claims

1. A PCB defect detection method based on improved YOLOv7, characterized in that: The following steps are involved: S1. Divide the PCB defect dataset containing six types of defects into training set, validation set and test set. The division process adopts random form. S2. Build an improved YOLOv7 detection model; S3. Preprocess the data in the training set, including the process of mosaic data enhancement and the process of clustering the labeled target bounding boxes using the clustering algorithm; S4. Input the preprocessed training set data and validation set data into the improved YOLOv7 detection model and start model training until the loss curve converges; S5. Input the test set images into the trained improved YOLOv7 detection model to obtain the recognition results of the PCB defect images; In step S2, an improved YOLOv7 detection model is constructed, and the process specifically includes: The SiLU version of the YOLOv7 tiny network is selected as the basis for improvement and improved in the following way: S21. Replace the high-resolution detection head: Small objects are more difficult to extract than regular objects and are easily lost in deep extraction. Therefore, the features of the first feature map are introduced into the head part of the network, and a high-resolution small object detection head is set. At the same time, the lowest resolution detection head is deleted to reduce the number of parameters and streamline the model. S22. Replace the ordinary 3*3 convolution in the feature extraction structure of the backbone with DCNV3: PCB defect detection requires both shallow and deep features. Deformable convolution is a variant of convolution that adds a bias to the convolution to make it deformable. This allows multi-scale information to be obtained through different receptive fields, which meets the needs of PCB defect detection. S23. Select NWD as a new evaluation metric and introduce it into the loss function. Use NWD Loss to improve the problem that IoU is sensitive to the offset of small target bbox. S24. Add GAM to the head of YOLOv7: GAM extracts relevant information by selectively focusing on the required parts of the channel and space, thereby improving recognition accuracy. It is added in two positions before the detection head.

2. A PCB defect detection method based on improved YOLOv7 according to claim 1, characterized in that: In step S1, the six types of defects specifically include: leak holes, rat bites, open circuits, short circuits, strays, and miscellaneous copper.

3. A PCB defect detection method based on improved YOLOv7 according to claim 1, characterized in that: In step S1, the process of dividing the PCB defect data set specifically includes: randomly dividing the PCB defect data set into a training set, a validation set, and a test set with a division ratio of 8:1:1, putting the labels into corresponding folders, and storing the files in yolo format.

4. A PCB defect detection method based on improved YOLOv7 according to claim 1, characterized in that: In step S3, the Mosaic data is enhanced and the target bounding boxes are clustered. The specific process includes: first, four pictures are randomly selected from the data set, and then the data is augmented, including flipping, scaling, and color gamut change operations, and then the pictures and boxes are combined; based on the labeled target bounding boxes in the training set, 9 anchor boxes from small to large are obtained a priori through the K-means clustering algorithm.

5. A PCB defect detection method based on improved YOLOv7 according to claim 4, characterized in that: In step S4, the model training process includes: S41, input the data-enhanced training set images into the Backbone network for feature extraction to obtain a feature map; S42, input the feature map into the head network, fuse the features of different feature maps, and obtain the detection box; S43, screening the detection frame according to a preset threshold to obtain a prediction frame; S44. Input the validation set images into the improved YOLOv7 model and obtain the detection results based on the predicted boxes. S45. Calculate the loss function of the improved YOLOv7 model during training based on the detection results, and continuously adjust the parameters of the model so that the loss function of the model continues to decrease until the training is finally completed.

6. A PCB defect detection method based on improved YOLOv7 according to claim 5, characterized in that: The specific process of step S5 is as follows: verifying the recognition effect of the model, preprocessing the test set image data, inputting the trained weight file in a unified format, and framing the target recognized by the model on the image to obtain the PCB defect image recognition result.

Images