This is the official implementation for the paper 'Frequency-aware Image Restoration for Industrial Visual Anomaly Detection'. https://arxiv.org/abs/2309.07068 If you have any questions, you could contact [email protected]
Different data_loader.py files correspond to different high-frequency extractors. Code of FAIRm (using morphological gradient) can be found in https://github.com/liutongkun/EdgRec
The method is evaluated on:
the MVTec AD dataset: https://www.mvtec.com/company/research/datasets/mvtec-ad
the VisA dataset: https://github.com/amazon-science/spot-diff
#20240828: The dataloader.py with “newaug” introduces NSA as an aditional synthesized anomaly, achieving a 99.1% image-level AUROC on the MVTec AD dataset.
The original code uses the DTD dataset to create synthesized anomalies, so you first need to download it
:
the DTD dataset (optional): https://www.robots.ox.ac.uk/~vgg/data/dtd/
Then
python train.py --gpu_id 0 --obj_id -1 --lr 0.0001 --bs 8 --epochs 800 --data_path /home/b211-3090ti/Anomaly-Dataset/mvtec_ad/ --anomaly_source_path /home/b211-3090ti/Anomaly-Dataset/dtd/images --log_path /home/b211-3090ti/FAIR/checkpoints_mvtecad/ --checkpoint_path /home/b211-3090ti/FAIR/checkpoints_mvtecad/ --visualize
Change all the involved paths to your own paths
activate line 16 and line 76 in data_loaderbhpfnoDTD.py
self.images = sorted(glob.glob(root_dir+"/*/*.JPG"))
mask_file_name = file_name.split(".")[0]+".png"
It's also feasible to train it without extra data, just activate line 3 in train.py:
from data_loaderbhpfnoDTD import MVTecTrainDataset
Then
python train.py --gpu_id 0 --obj_id -1 --lr 0.0001 --bs 8 --epochs 800 --data_path /home/b211-3090ti/Anomaly-Dataset/mvtec_ad/ --log_path /home/b211-3090ti/FAIR/checkpoints_mvtecad/ --checkpoint_path /home/b211-3090ti/FAIR/checkpoints_mvtecad/ --visualize
python test.py --gpu_id 0 --base_model_name FAIR_0.0001_800_bs8 --data_path /home/b211-3090ti/Anomaly-Dataset/mvtec_ad/ --checkpoint_path /home/b211-3090ti/FAIR/checkpoints_mvtecad/
Change all the involved paths to your own paths
If you want to visualize the results, add
--saveimages
MVTec AD
https://drive.google.com/file/d/1hbl_k_hKgxo_IejNNu8daEYJkfLN3Wxk/view?usp=sharing #20240828 using new synthesized anomalies
VisA
https://drive.google.com/file/d/1G2YosWUgJzGWq5S4OsPFHE0jShXmHHhd/view?usp=sharing #20240828 using new synthesized anomalies
We use the codes from https://github.com/VitjanZ/DRAEM, https://github.com/taikiinoue45/RIAD, and https://www.mvtec.com/company/research/datasets/mvtec-3d-ad, https://github.com/hmsch/natural-synthetic-anomalies
A big thanks to their great work