Detect Defects on Printed Circuit Boards Using YOLO v4 Network

This example shows how to detect, localize, and classify defects in printed circuit boards (PCBs) using a you only look once version 4 (YOLO v4) deep neural network.

PCBs contain individual electronic devices and their connections. Defects in PCBs can result in poor performance or product failures. By detecting defects in PCBs, production lines can remove faulty PCBs and ensure that electronic devices are of high quality.

Download Pretrained YOLOv4 Detector

By default, this example downloads a pretrained version of the YOLOv4 object detector using the helper function downloadTrainedNetwork. The helper function is attached to this example as a supporting file. You can use the pretrained network to run the entire example without waiting for training to complete.

trainedPCBDefectDetectorNet_url = "https://ssd.mathworks.com/supportfiles/"+ ...
    "vision/data/trainedPCBDefectDetectorYOLOv4.zip";
downloadTrainedNetwork(trainedPCBDefectDetectorNet_url,pwd);
load("trainedPCBDefectDetectorYOLOv4.mat");

List the classes that this network is trained to classify.

classNamesStr = ["missing hole","mouse bite","open circuit", ...
    "short","spur","spurious copper"];

Download PCB Defect Data Set

This example uses the PCB defect data set [1] [2]. The data set contains 1,386 images of PCB elements with synthesized defects. The data has six types of defect: missing hole, mouse bite, open circuit, short, spur, and spurious copper. Each image contains multiple defects of the same category in different locations. The data set contains bounding box and coordinate information for every defect in every image. The size of the data set is 1.87 GB.

Specify dataDir as the desired location of the data set. Download the data set using the downloadPCBDefectData helper function. This function is attached to the example as a supporting file.

dataDir = fullfile(tempdir,"PCBDefects");
downloadPCBDefectData(dataDir)

Perform Object Detection

Read a sample image from the data set.

sampleImage = imread(fullfile(dataDir,"PCB-DATASET-master","images", ...
    "Missing_hole","01_missing_hole_01.jpg"));

Predict the bounding boxes, labels, and class-specific confidence scores for each bounding box by using the detect (Computer Vision Toolbox) function.

[bboxes,scores,labels] = detect(detector,sampleImage);

Display the results.

imshow(sampleImage)
showShape("rectangle",bboxes,Label=labels);
title("Predicted Defects")

Prepare Data for Training

Create an image datastore that reads and manages the image data.

imageDir = fullfile(dataDir,"PCB-DATASET-master","images");
imds = imageDatastore(imageDir,FileExtensions=".jpg",IncludeSubfolders=true);

Create a file datastore that reads the annotation data from XML files. Specify a custom read function that parses the XML files and extracts the bounding box information. The custom read function, readPCBDefectAnnotations, is attached to the example as a supporting file.

annoDir = fullfile(dataDir,"PCB-DATASET-master","Annotations");
fds = fileDatastore(annoDir,ReadFcn=@readPCBDefectAnnotations, ...
    FileExtensions=".xml",IncludeSubfolders=true);

Save the labeled bounding box data as a box label datastore.

annotations = readall(fds);
tbl = struct2table(vertcat(annotations{:}));
blds = boxLabelDatastore(tbl);

Get the names of the object classes as a categorial vector.

classNames = categories(blds.LabelData{1,2})

classNames = 6×1 cell
    {'missing_hole'   }
    {'mouse_bite'     }
    {'open_circuit'   }
    {'short'          }
    {'spur'           }
    {'spurious_copper'}

Combine the image and box label datastores.

ds = combine(imds,blds);

Analyze Object Class Distribution

Measure the distribution of class labels in the data set by using the countEachLabel (Computer Vision Toolbox) function. The classes in this data set are balanced.

countEachLabel(blds)

ans=6×3 table
         Label         Count    ImageCount
    _______________    _____    __________

    missing_hole        497        115    
    mouse_bite          492        115    
    open_circuit        482        116    
    short               491        116    
    spur                488        115    
    spurious_copper     503        116    

Partition Data

Split the data set into training, validation, and test sets. Because the total number of images is relatively small, allocate a relatively large percentage (80%) of the data for training. Allocate 10% for validation and the rest for testing.

numImages = ds.numpartitions;
numTrain = floor(0.8*numImages);
numVal = floor(0.1*numImages);

shuffledIndices = randperm(numImages);
dsTrain = subset(ds,shuffledIndices(1:numTrain));
dsVal = subset(ds,shuffledIndices(numTrain+1:numTrain+numVal));
dsTest = subset(ds,shuffledIndices(numTrain+numVal+1:end));

Preprocess Training Data

Specify the network input size.

inputSize = [800 960 3];

Resize the training data by using the transform function with custom resizing operations specified by the preprocessPCBDefectData helper function. The helper function is attached to the example as a supporting file. The preprocessPCBDefectData helper function resizes the images and the bounding boxes to the size of the network input.

dsTrain = transform(dsTrain,@(data)preprocessPCBDefectData(data,inputSize));

Estimate Anchor Boxes

Estimate ten anchor boxes based on the size of objects in the preprocessed training data by using the estimateAnchorBoxes (Computer Vision Toolbox) function. This example uses a YOLO v4 deep learning network with two output feature maps, so split the anchor boxes into two sets with five anchor boxes each.

anchorBoxes = estimateAnchorBoxes(dsTrain,10);
anchorBoxes = {anchorBoxes(1:5,:); anchorBoxes(6:end,:)};

Augment Training Data

Augment the training data by using the transform function with custom preprocessing operations specified by the augmentDataForPCBDefectDetection helper function. The helper function is attached to the example as a supporting file. The augmentDataForPCBDefectDetection function applies these augmentations to the input data:

• Random horizontal flip

• Resizing by a random scale factor in the range [1, 1.1]

dsTrain = transform(dsTrain,@augmentDataForPCBDefectDetection);

Create YOLO v4 Object Detector Network

Create the YOLO v4 object detector by using the yolov4ObjectDetector (Computer Vision Toolbox) function. Specify the name of the pretrained YOLO v4 detection network trained on COCO data set [3]. Specify the class names, the estimated anchor boxes, and the network input size.

detectorToTrain = yolov4ObjectDetector("tiny-yolov4-coco",classNames,anchorBoxes, ...
    InputSize=inputSize);

Specify Training Options

Specify network training options using the trainingOptions function. Train the object detector using the Adam solver for 500 epochs. Specify the ValidationData name-value argument as the validation data.

options = trainingOptions("adam", ...
    MiniBatchSize=32, ...
    MaxEpochs=500, ...
    BatchNormalizationStatistics="moving", ...
    Shuffle="every-epoch", ...
    VerboseFrequency=250, ...
    ValidationFrequency=250, ...
    CheckpointPath=dataDir, ...
    ValidationData=dsVal, ...
    ResetInputNormalization=false, ...
    DispatchInBackground=true);

Train Detector

To train the detector, set the doTraining variable in the following code to true. Train the detector by using the trainYOLOv4ObjectDetector (Computer Vision Toolbox) function.

Train on one or more GPUs, if available. Using a GPU requires Parallel Computing Toolbox™ and a CUDA® enabled NVIDIA® GPU. For more information, see GPU Computing Requirements (Parallel Computing Toolbox). Training takes about 24 hours on an NVIDIA Titan RTX™ with 24 GB of memory.

doTraining = false;
if doTraining       
    [detector,info] = trainYOLOv4ObjectDetector(dsTrain,detectorToTrain,options);
    modelDateTime = string(datetime("now",Format="yyyy-MM-dd-HH-mm-ss"));
    save(tempdir+filesep+"trainedPCBDefectDetector-"+modelDateTime+".mat", ...
        "detector");
end

Evaluate Detector

Evaluate the trained object detector by measuring the average precision. Precision quantifies the ability of the detector to correctly classify objects.

Detect the bounding boxes for all test images.

detectionResults = detect(detector,dsTest);

Calculate the average precision score for each class by using the evaluateDetectionPrecision (Computer Vision Toolbox) function. Also calculate the recall and precision values from each detected defect. Recall quantifies the ability of the detector to detect all relevant objects for a class.

[averagePrecision,recall,precision] = evaluateDetectionPrecision(detectionResults,dsTest);

Display the average precision score for each class.

table(classNames,averagePrecision)

ans=6×2 table
        classNames         averagePrecision
    ___________________    ________________

    {'missing_hole'   }        0.93464     
    {'mouse_bite'     }        0.75237     
    {'open_circuit'   }        0.85178     
    {'short'          }        0.97002     
    {'spur'           }        0.96218     
    {'spurious_copper'}        0.83631     

A precision/recall (PR) curve highlights how the precision of a detector at varying levels of recall. The ideal precision is 1 at all recall levels. Plot the PR curve for the test data.

class = 1;
plot(recall{class},precision{class})
title("Average Precision for '"+classNamesStr(class)+ ...
    "' Defect: "+num2str(averagePrecision(class),"%0.2f"))
xlabel("Recall")
ylabel("Precision")
grid on

References

[1] Huang, Weibo, and Peng Wei. "A PCB Dataset for Defects Detection and Classification." arXiv, January 23, 2019. https://arxiv.org/abs/1901.08204.

[2] PCB-DATASET. Accessed December 20, 2022. https://github.com/Ironbrotherstyle/PCB-DATASET.

[3] Bochkovskiy, Alexey, Chien-Yao Wang, and Hong-Yuan Mark Liao. “YOLOv4: Optimal Speed and Accuracy of Object Detection.” arXiv, April 22, 2020. https://arxiv.org/abs/2004.10934.