Image Classification of PCBs and its Web Application (Flask) | by Utkarsh Ankit | Towards Data Science

Hello, This blog is about creating an Image Classification model for PCBs(Printed Circuit boards) to detect defective PCBs and classify them as Good or Bad. So for that, we will create a Deep Learning Model and try to get the best possible results along with proper visualisation of each step. After creating the working model, we will create a Web Application for the same using Flask. You can find the code for this project on my GitHub along with the setup instruction file below. So, let’s go ahead and flex.
Requirements
Jupyter NotebookPython 3.5 or Later. (3.7 preferable)Libraries — Tensorflow, Keras, Pandas, Numpy, Matplotlib, OpenCV, Pillow, Flask. [all latest version]
Jupyter Notebook
Python 3.5 or Later. (3.7 preferable)
Libraries — Tensorflow, Keras, Pandas, Numpy, Matplotlib, OpenCV, Pillow, Flask. [all latest version]
Understanding the Dataset
PCBs (Printed Circuit Boards) is a widely used electronic component, used almost in any electronic device. Sometimes during its mass production in factories leads to some anomalies which lead to a defective product. Electronic companies have to bear that loss. So this deep learning Image Classification model helps to identify Good and Bad Circuits at production time only.
So, the dataset we are using is a binary image of PCB, in which the Black part denotes the circuit and the white parts signify the empty circuit areas. The dots or noise in the image shows the anomalies.
As we can see the image has no anomalies and defects in it. Now, we will see another image of a PCB with defects.
As here the defects are marked, it is an example of a defective PCB. You can download the dataset from here.
Data Preparation/Preprocessing
Now after we download the dataset, we can see the dimension of each image is 640x640. It will be better if we reduce the dimension of them so that it will be easy for training the data, but to make sure the quality of the image didn’t get affected, we will compress the image using the Anti-Aliasing technique.
Here is my link to my GitHub repo which helps to do the same. Use the Image Resizer tool there and you will get the desired result.
So, after resizing the image to 224x224, we save them into a different folder. I saved 1,095 images of somewhat equally mixed good and bad PCBs. So, that folder is going to be your Training Data set. Pick some other images (different from the training set) and put them into a separate folder to work as your Test Data set. Make sure the test folder is a good mixture of good and bad images with a proper ratio maintained.
Now I will check my training dataset that it has a proper ratio of both types of images or in other words, to check the possible “class imbalance”. Well for my reference, I named the good PCB image as Good().jpeg and bad as Bad().jpeg. Then I labelled the good images as ‘0’ and bad images as ‘1’.
And then we plot their graph:-
So, we can observe from the graph that our classes are pretty balance. Now for Cross-Validation, we will split out the training dataset into training and validation sets. As I mentioned earlier I took 1,095 images. I split them as 80% training data and 20% validation set.
train_df, validate_df = train_test_split(df, test_size=0.20, random_state=42)train_df = train_df.reset_index(drop=True)validate_df = validate_df.reset_index(drop=True)
And then we plot both training and testing image ratios:-
⇙
This is the plotted graph of the final “training dataset”, i.e. approximately 400+ images in both classes.
⇖
⇙
This is the plotted graph of the “validation dataset”, i.e. approximately 100+ images in both classes.
⇖
___
Now I am going to apply one of the greatest tools according to me i.e. Keras ImageDataGenerator. It is a very useful tool as it helps to do Data Augmentation in real-time while training the model. As we know our training data is already less, so applying augmentation not only helps us to make our model robust but also saves memory.
Now, we add the path and load the images in the generator.
So, now we can say that our Data Preparation/Preprocessing part is complete.
Model Implementation
Now it’s time to design the brain. Here I used the “MobileNet” model which is pre-trained on the ImageNet dataset. Well, MobileNet is a Deep Learning based Convolution Neural Network Architecture. It is a lightweight architecture and uses depthwise separable convolutions. You can learn more about it from its paper here.
As I mentioned we are using the MobileNet model pre-trained of ImageNet Dataset, so here we are using the concept of Transfer Learning. Therefore instead of training the whole model, we will freeze the base layer and add some other layers according to our requirements.
As we are using pre-trained weights, we will make sure that all the weights are non-trainable except the last few dense layers. Well, I set the first 20 layers to be non-trainable.
for layer in model.layers[:20]:    layer.trainable=Falsefor layer in model.layers[20:]:    layer.trainable=True
But if you want to set all layers to be non-trainable, you can use this instead.
for layer in model.layers:        layer.trainable=False
So, we are done with the Model Implementation part now.
Training
So, our data preparation is done along with model implementation. Now, it’s time to train our model. Before I proceed further, I want to mention that I used two Callbacks named “Early Stopping” and “Learning Rate Reduction”.
We need to consider the number of Epochs carefully as too many of them lead to overfitting the data, whereas too few leads to underfitting. Early Stopping allows you to specify an arbitrarily large number of training epochs and stop training when the model stops improving on a hold out validation dataset. Here is a link if you want to learn more about it.
ReduceLROnPlateau reduces the Learning Rate when a metric is stopped improving. Models often benefit from reducing the learning rate by a factor of 2–10 once learning stagnates. This callback monitors a quantity and if no improvement is seen for a ‘patience’ number of epochs, the learning rate is reduced. You can check this from here.
Now, we will start training our model. Well, I trained my model on Google Colab as it gives certain advantages, I will advise everyone to try to use Google Colab, you just need to upload or link the dataset on google drive and then mount it on colab using this simple call by giving right folder location:-
from google.colab import drivedrive.mount('/content/drive')
Now we will set all the variables in the model.fit.generator() and then we will start to train our model.
Well, I set my epoch as 50 but it stopped it at 42 as it stops improving. I plotted two graphs from the data we extracted while training the model.
.
⇙
“Loss vs Epochs” graph.
⇖
.
⇙
“Accuracy vs Epochs” graph.
⇖
.
Note:- The Blue line denotes the training and Red line denotes the validation variable respectively in both the graphs.
So, after training the model, I got a Validation Accuracy of 95.71%.
Now, we save the trained weights in the .h5 or .hdf5 file (both are the same).
model.save_weights("model.h5")
If you trained it via Google Colab, you can save the weights in the drive directly like this:-
model.save('/content/drive/My Drive/folder (1)/model.h5')
Therefore our training is complete, now it’s time to check the results.
Prediction and Results
So first, we will load our model.
model = load_model(r'C:/Users/nEW u/Flask/model.h5')
Then we will use ImagedataGenerator to get to the test folder and call the image files collectively. And will call the predict function and set it with the image generator function.
Now, we will plot our results along with the images. I took the help of the Matplotlib library.
And here comes our result:-
Note:- As I mentioned earlier, I named the images files as good and bad for my reference. So, the initial part is the name of the image file and the result is in the Bracket. Therefore you can compare the Result with the actual Image quality as per the image name.
So, here our Image Classification model is complete successfully. Now, it's time to create a Web Application. I am going to use Flask for it.
Flask is a Python’s API used to create web applications. So, for start, I am gonna give you a basic code snippet that we usually follow while using Flask.
Hope you got a bit idea, Now let’s start implementing.
Functions
First, we must write functions to call the model, for processing the image before we pass it to our model and prediction function which let our image pass the model and return the result.
For calling the model:-
For processing the image:-
Prediction function (it uses the“load_image” function):-
Design
As we are done with our functions, we will proceed forward to design our web pages. So, here we will need two pages, On the first page, it will ask to select and upload the image. The second page will show the image along with the result.
Important Note:- While using Flask we must be very attentive as we need to create a separate folder. In that folder, we will create our python file in which we use flask, and we also need to create two subfolders named “Templates” and “Static”. Templates are where we save our HTML files, and in the Static folder, we save images and other stuff. You must follow this convention if you want to save yourself from the error cycle.
So, as we need two pages, I designed two HTML pages accordingly. For design purposes, I used “Bootstrap”. Well, Bootstrap is an open-source Front-End framework used to create websites and web apps. It contains various HTML, CSS and JavaScript-based design templates for UI interface elements such as Buttons and Forms etc.
So, here is my first web page HTML file named home.html:-
Notice on the third line, bootstrap link for the design is used (CSS based template).
And here is the second page HTML file, saved as predict.html:-
Here both CSS and JavaScript templates are used.
Compilation and Deployment
Our model is ready, functions are ready, templates are also ready. Now it’s time to compile all those things and enjoy the magic.
First, we run the Flask:-
app = Flask(__name__)
Then we load the model by calling the function:-
get_model()
And now we call both the pages along with functions and templates:-
Well, now its time for the final step i.e. to run everything and to generate URL:-
Now a URL is generated (notice in blue). Click on it and we will get our first page like this:-
Then we choose the image file which we want to check and then click upload. After pressing the upload button, it redirects to the next page with url+”/predict”. It will show the image with the result:-
Here we are with the final result. So, the implementation of a Deep Learning model-based Web Application for defect detection in PCB is hereby successful.
You can clone the repository by:-
git clone https://github.com/utk-ink/Defect-Detection-of-PCB.git
Or
Here is my GitHub’s link for the whole code:-
github.com
Hope you liked it!!!
Well, this blog is not monetized, I just wanted to share my knowledge in the best possible way. Please do give Claps, Follow and Share as it motivates me to write more. Thank You!!