Detecting defects in PCBs with YOLOX using OpenMMLab
Object Detection using OpenMMLab with a new architecture in the ever famous YOLO series — YOLOX. We will be trying to detect different defects in Printed Circuit Boards.
PCB (Printed Circuit Board)
I know I know, you must be asking, what is a PCB? You aren't? Well, for someone who doesn't know what a PCB is, here's a definition from Wikipedia:
¹A printed circuit board ( PCB) is a laminated sandwich structure of conductive and insulating layers. PCBs have two complementary functions. The first is to affix electronic components in designated locations on the outer layers by means of soldering. The second is to provide reliable electrical connections (and also reliable open circuits) between the component's terminals in a controlled manner often referred to as PCB design.
I bet you must have seen a PCB at least once in your life but probably weren't bothered to know what it was. Did you ever look at the back of the DVD Reader? Here's an image of a PCB on a DVD Reader from Wikipedia:
PCBs are everywhere. Almost all electronic devices have a PCB somewhere hidden in them. There are a lot of instances where these PCBs can be defective either while designing them or after they are in use. Here's a list of a few of the commonly known defects in a PCB as listed on the internet³ and an example image from a freely available dataset ⁴.
- Opens
- Excessive solder
- Component shifting
- Cold joints
- Solder bridges
- Webbing and splashes
- Lifted pads
We aren't going to go deep into what they exactly mean as this is not what the blog is about. But, looking at it from a Computer Engineer's perspective who understands a bit of Computer Vision and Deep Learning, it seems like detecting a defect in a digital image of a PCB is a solvable problem.
We will be using mmdetection⁶ from OpenMMLab ⁸ to detect defects in the PCB images. OpenMMLab⁸ is a deep learning library that has pre-trained models of most of the state-of-the-art implementations in the Computer Vision domain. It has implementations for almost all the well-known vision problems like classification, object detection & segmentation, pose estimation, image generation, object tracking, and so much more.
YOLOX: Exceeding YOLO series in 2021⁷
In this blog, we will be using YOLOX⁷ that we will be fine-tuning with mmdetection⁶. YOLOX ⁷ is a state-of-the-art model released in 2021 which is an improvement in the YOLO series. The authors have made some significant improvements which are listed down below.
- Introduction of SimOTA for label assignment
- Removed anchor boxes
- Special attention to data augmentation like MixUp with jittering
- Separate heads for detection and classification
The previous YOLO series from v3 to v5 have all had a single prediction head which included bounding box predictions, classification score predictions as well as objectness score predictions as seen in the top half of the image above.
This was changed in the YOLOX⁷ series where the authors have chosen to use a decoupled head with separate heads for all predictions. As seen in the bottom half of Fig2 above, the detection head and classification head are in separate heads. This helps in improving the convergence time during training (as shown in Fig3) and slightly improves the model accuracy.
The speed of the model does take a hit due to this as the number of parameters is significantly increased with the separation into two heads. As we can see in Fig4, the YOLOX-L is a tad bit slower than YOLOv5-L. It also has Nano and Tiny versions specifically built for edge devices that have much lower parameters.
Compared with the previous state-of-the-art object detection models, they do have an improvement in terms of Average Precision, with a slight degradation in the FPS.
Finally, as the great Linus Torvalds once said,
Talk is cheap. Show me the code.
Let's jump right in!
Fine-tuning YOLOX with mmdetection
We have an open-source PCB Defect Dataset known as DeepPCB⁵. This dataset consists of 1500 image pairs with each one having a defect-free template image and an image with defects that has bounding box annotations of 6 common types of defects viz. open, mouse-bite, short, spur, spurious copper, and pin-hole. You can get a bit more in-depth information about it from the dataset link ⁵ if you need it.
The images are of dimensions 640×640 which are perfect in our case as YOLOX⁷ is trained on the same dimensions.
OpenMMLab⁸
OpenMMLab makes it very easy to fine-tune state-of-the-art models with very little code change. It has a comprehensive API for specific use-cases. We will be using mmdetection⁶ for fine-tuning YOLOX⁷ on the DeepPCB⁵ dataset.
Dataset Format
Note: The PCB Defect Dataset is an open-source dataset taken from the DeepPCB Github repo with MIT License which can be referred to here.
We need to modify the dataset into COCO format or the Pascal VOC format to retrain the model. This is required for mmdetection⁶ to load our custom dataset for training. We will be going with the COCO format for our training purposes. You don't need to go through the hassle to convert the dataset to COCO format as it is already done for you. You can directly download the converted dataset from here. The whole dataset is the same as in DeepPCB⁵ with only the addition of train and test JSON files with annotations in COCO format for training.
I won't be going through the conversion to COCO format as you can find many documentations going through the process like it has been mentioned in the mmdetection documentation as well. Feel free to go through the script and don't mind leaving a comment if you need any help with it. The script to convert this data set to COCO format is shared here.
Dataset Configuration
The next step is to modify the dataset configuration to use our custom dataset. We will need to add/modify specific things like the number of classes, annotation paths, dataset path, number of epochs, base config path, and some dataloader arguments.
We will be copying a prewritten YOLOX-s configuration and modifying it for our dataset. The rest of the configuration, like, augmentation, optimizers, and other hyperparameters will be the same.
We aren't going to change a lot as the main motive of this blog is to get familiar with the problem at hand, try out the state-of-the-art YOLOX architecture, and experiment with the mmdetection library. We will be naming this file as yolox_s_config.py and using it for training.
We will be adding the class names and changing the number of classes for our prediction head. The base path needs to be changed as the config will be loaded from the root directory instead of the configs directory.
_base_ = ['configs/_base_/schedules/schedule_1x.py', 'configs/_base_/default_runtime.py']classes = ('open', 'short', 'mousebite', 'spur', 'copper', 'pin-hole')bbox_head = dict(type='YOLOXHead', num_classes=6, in_channels=128, feat_channels=128)
We need to modify the train dataset loader slightly to use our classes and annotations path.
train_dataset = dict(
type='MultiImageMixDataset',
dataset=dict(
type=dataset_type,
classes=classes,
ann_file='train.json',
img_prefix='',
pipeline=[
dict(type='LoadImageFromFile'),
dict(type='LoadAnnotations', with_bbox=True)
],
filter_empty_gt=False,
),
pipeline=train_pipeline)We need to do the same for validation and test set here. We aren't going to use a separate test set here, instead, we will use the same set for validation and test.
data = dict(
samples_per_gpu=8,
workers_per_gpu=4,
persistent_workers=True,
train=train_dataset,
val=dict(
type=dataset_type,
classes=classes,
ann_file='test.json',
img_prefix='',
pipeline=test_pipeline),
test=dict(
type=dataset_type,
classes=classes,
ann_file='test.json',
img_prefix='',
pipeline=test_pipeline))We will only be training the model for 20 epochs and getting the validation results every 5 epochs. We don't need to train it any longer as we get decent results in just 20 epochs. Though it is not the best model, it is decent for the sake of the blog.
max_epochs = 20
interval = 5You can also copy the whole edited configuration from the gist here and use it to train the model. You can thank me later! 😉
Training
We are good to go with the dataset part. The next thing we need to do is train the model. The best part about the mmdetection⁶ library is, that everything concerning training is already done for you. All you need to do is run the training script from the tools directory and pass the path to the dataset configuration we created above.
python3 tools/train.py yolox_s_config.pyEt Voilà! 🔥
You have successfully trained the YOLOX⁷ model to detect PCB defects using the DeepPCB⁵ dataset with the help of mmdetection⁶.
Inference
Let’s see how our model performs on some examples. You must be wondering, looking at how easy it was to train the model, there must be a command to run inference on an image? There is! But, don’t let the easy procedure of training the model spoil you. Let’s write some code for inference, but to your joy, it is less than 10 lines of code.
from mmdet.apis import init_detector, inference_detector, show_result_pyplotconfig_file = 'yolox_s_config.py'checkpoint_file = 'best_bbox_mAP_epoch_20.pth'
device = 'cuda:0'# init a detector
model = init_detector(config_file, checkpoint_file, device=device)# inference the demo image
image_path = 'demo.jpg'op = inference_detector(model, image_path)
show_result_pyplot(model, image_path, op, score_thr=0.6)
This will display an image with bounding boxes drawn on it with predicted class names. Here’s an example image from the dataset with the model predictions.
And we did it! 🎊
You can find the code for training and inference on Google Colab here. You can also try our pre-trained model and run an inference with it. The pre-trained model can be downloaded here. Feel free to play around with it and do let me know in the comments if you run into any issues. Happy to help! 😉
Conclusion
Today, we learned about a new problem prevalent in the real world and tried solving the problem with a state-of-the-art model like YOLOX⁷.
We also used mmdetection⁶ which is one of the leading open-source libraries in the deep learning community to train an object detection model. It would be unfair if I don’t mention how mmdetection⁶ has spoiled us in getting around a problem so quickly and easily with barely any custom scripting.
I hope you learned something new today! Feel free to run through the code.
At Deep Learning Analytics, we are extremely passionate about using Machine Learning to solve real-world problems. We have helped many businesses deploy innovative AI-based solutions. Contact us through our website here if you see an opportunity to collaborate.
References
- PCB Definition: https://en.wikipedia.org/wiki/Printed_circuit_board
- PCB Image: https://en.wikipedia.org/wiki/File:SEG_DVD_430_-_Printed_circuit_board-4276.jpg
- Common PCB defects: https://blog.matric.com/7-types-of-pcb-soldering-defects-to-watch-out-for
- PCB Defect Dataset Paper: https://arxiv.org/abs/1902.06197
- PCB Defect Dataset: https://github.com/tangsanli5201/DeepPCB
- MMDetection: https://github.com/open-mmlab/mmdetection
- YOLOX Paper: https://arxiv.org/abs/2107.08430
- OpenMMLab: https://openmmlab.com/
