(Image from MVTec AD datasets https://www.mvtec.com/company/research/datasets/mvtec-ad/)
Input shape : (n, 3, 224, 224)
Left to right: input, ground truth, predicted mask, segmentation result
Automatically downloads the onnx and prototxt files on the first run. It is necessary to be connected to the Internet while downloading.
In order to get the feature vector of the normal product, it is necessary to prepare the file of the normal product.
By default, normal files are got from the train
directory.
To calculate the threshold you also need to prepare some test images and corresponding ground truth images.
By default, test images are taken from the images
directory and ground truth images are taken from the gt_masks
directory.
For the sample image, first download MVTec AD datasets and place bottle/train/good/*.png
files to the train
directory.
For the sample image.
$ python3 mahalanobisad.py
If you want to specify the input image, put the image path after the --input
option.
You can use --savepath
option to change the name of the output file to save.
$ python3 mahalanobisad.py --input IMAGE_PATH --savepath SAVE_IMAGE_PATH
The feature vectors created from files in the train directory are saved to the pickle file.
From the second time, by specifying the pickle file by --feat
option,
it can omit the calculation of the feature vector of the normal product.
The name of the pickle file created is the same as the name of a normal product file directory.
$ python3 mahalanobisad.py --feat train.pkl
You can specify the directory of normal product files with the --train_dir
option.
$ python3 mahalanobisad.py --train_dir train
The threshold is calculated using the test images and the ground truth images, but you can also give a threshold.
$ python3 mahalanobisad.py --threshold 300
By adding the --aug
option, you can process with augmentation.
(default is processing without augmentation)
$ python3 mahalanobisad.py --aug
By adding the --aug_num
option, you can specify amplification factor of data by augmentation.
(default is 5)
$ python3 mahalanobisad.py --aug --aug_num 5
Pytorch
ONNX opset=11
efficientnet-b0.onnx.prototxt
efficientnet-b1.onnx.prototxt
efficientnet-b2.onnx.prototxt
efficientnet-b3.onnx.prototxt
efficientnet-b4.onnx.prototxt
efficientnet-b5.onnx.prototxt
efficientnet-b6.onnx.prototxt
efficientnet-b7.onnx.prototxt