This project aims to develop a deep learning-based anomaly detection system for target images. This system can be integrated into a factory assembly line as an inspection tool for manufactured products. Our approach uses the combination of AutoEncoder and triplet loss to detect defective(anomalies) images.
The first step in our pipeline is data loading and preprocessing. The images are labelled as OK and NG, where 'NG' represent the anomaly(defected) images. Image info, which is image path and image label is stored in 'image_path.csv'. The images in this CSV file represent the original images. After storing image info, the next part is to extract a rectangle around the circular cross-section (area of concern). Below are the examples for the same.
Drawn rectangular box is extraced from the image and then used in the training process.
For drawing rectangular boxes around circular cross section, we are using Hough Circle AlgorithmWe store the image path of cropped images in 'crop_image_paths.csv'. On top of of that each cropped image is subjected to a set of augmentations. These are:
- Rotation - (0, 90, 180, 270)
- Flip - (None, Horizontal, Vertical)
Using combinations of these augmentations we get 11 augmented images from one original image. This heavy augmentaion is required because we have a very limited dataset for our intial experiment.
We are using an Autoencoder in association with triplet loss. For each forward pass, model calculates two different loss.
- Reconstrunction loss
- Triplet loss
-
Standard Triplet loss = max(0, d(a,p) - d(a,n) + margin) . Our objective in triplet loss is to make d(a,n) >= d(a,p) + margin
-
LossLess Triplet loss =
-
N here is the dimension of the last embedding layer.
-
Let us say that dp = -ln(-(fa - fp)2 / N +1 + ) and dn = -ln(-(N - (fa - fp)2) / N +1 + )
-
The dp term in the equation that is underlined by red is representative of the anchor positive distances seen in the linear equation.
-
As the anchor-positive distance that is (fa - fp)2 increases the cost value dp also increases non linearly as shown in the red line graph.
-
The term dn, underlined by yellow, represents the difference between margin(N) and distance between anchor and negative. (N - (fa - fn)2)
-
As the distance between (fa - fn)2 decreases the loss value of dn increases non linearly as shown in the yellow line graph.
-
Combining the effect of both these equations, the model will try to decrease (fa - fp)2 and increase (fa - fn)2 such that both dn and dp approaches 0.
To create image_path.csv file, run the following command from the utils directory.
python3 make_image_df.py
Make sure that we are following this structure to store the images
images/
|
| - 20201028/
| |
| | - image_1
| |
| | - image_2
|
| - 20201207/
| |
| | - image_1
| |
| | - image_2
|
| - augmented_images/
|
| - cropped_images/
To generate cropped and augmented images run the following command from the utils folder.
python3 augment_crop_image.py
This script will add the cropped images in the images/cropped_images/ folder and generate crop_image_paths.csv file which is used in the traning process.
To start the training, run the following command from the src directory.
python3 train.py --num_epochs=100 --learning_rate=0.0001 --image_df='../crop_image_paths.csv' --batch_size=64 --save_interval=50
- num_epochs - Number of epochs for which the model will run
- learning_rate - Set learning rate for the model
- image_df - Describe the image df path which store image patha and corresponding label
- batch_size - Set the batch size
- save_interval - Number of epcohs interval after which we save model weigths and evaluation results.
Model weights, training and validation loss information and evaluation results are all stored in the saved_models directory
After the traning process gets completed, training ans validation dataframes are created inside the reports folder. A html file is also saved which display the data and metric information in table form.
This project uses Pytorch (Machine Learning Library) for writing deep learning architecture. This code can be used with or without GPU, user don't need to change anything to run the repo with GPU acceleration. It is already taken care in the code itself.
For macOS and Linux
python3 -m pip install --user --upgrade pip
For windows
py -m pip --version
VirtualEnv is used to manage Python packages for different projects. Using virtualenv allows you to avoid installing Python packages globally which could break system tools or other projects. You can install virtualenv using pip. For macOS and Linux
python3 -m pip install --user virtualenv
python3 -m venv env
For Windows:
py -m pip install --user virtualenv
py -m venv env
The second argument is the location to create the virtual environment. Generally, you can just create this in your project and call it env
For macOs and Linux
source path/to/env/bin/activate
For Windows
.\env\Scripts\activate
pip3 install -r requirement.txt