Skip to content

liutongkun/EdgRec

Repository files navigation

2023/7/17:
We refine the EdgRec to FAIR, which achieve 98.6 image-level performance on MVTec AD and 96.9 on VisA with a simple encoder-decoder, code will be available at https://github.com/liutongkun/FAIR
This repository will no longer be updated


This is the code for the paper entitled "Reconstruction from edge image combined with color and gradient difference for industrial surface anomaly detection" (https://arxiv.org/abs/2210.14485)
Currently the paper is under review, so we only provide a brief introduction here, which will be refined later. If you have any questions, you could also contact ltk98633@stu.xjtu.edu.cn, both in English or Chinese is OK.

The code is implemented in pytorch, and most basic codes are from https://github.com/VitjanZ/DRAEM , a big thanks to their great work. We also use the codes from https://github.com/taikiinoue45/RIAD and https://www.mvtec.com/company/research/datasets/mvtec-3d-ad , thanks them a lot

Our method can achieve 97.8 AUROC, (average result, slightly different each re-implementation) detection performance on the MVTec AD dataset. 

For the original paper, to train the model, you need to download MVTec AD dataset and the DTD dataset(for creating pseduo anomalies). Here we introduce a new creation of pseudo anomalies. You could also use 'data_loader-noDTD.py' to replace the orignal 'data_loader.py' which does not require any external dataset while achieves almost the same performance.

The parameters for training and testing can be found in 'train&test_para.txt'. You need to change the '--data_path' and '--anomaly_source_path'(when using data_loader-noDTD.py, this is not required)  to your own paths which include the mvtec ad dataset and the dtd dataset respectively.
For testing, you could visualize the result images by add '--saveimages' (cost more time), otherwise it will only show the performance (faster)

Like:
python train.py --gpu_id 0 --obj_id -1 --lr 0.0001 --bs 8 --epochs 800 --data_path .../mvtec_ad/ --anomaly_source_path .../dtd/images --log_path . --checkpoint_path . --visualize

python test.py --gpu_id 0 --base_model_name xxx(your checkpoint name) --data_path .../mvtec_ad/ --checkpoint_path xxx(your checkpoint path) --saveimages

Also you can reduce the epochs or network channels (originally 128) for fast implementation, as we didn't carefully adjusted these parameters (we have tried to reduce the channel to 64, the performance is around 97.4)     

MVTec AD pretrained:
https://drive.google.com/file/d/1npAqAo0UVlnI3_SpRfVIMtma3R_Cs_PD/view?usp=share_link
or Baidu disk:
https://pan.baidu.com/s/1o4i-tlJvB9em-l6o8o_vKw 
txk8 

MVTec 3D AD pretrained:
https://drive.google.com/file/d/1M0T61ET1XZOfmX_kmXebHp6Jdr44fdp4/view?usp=share_link
or
https://pan.baidu.com/s/1bD-gvo8OmIWWBZv6pe9l8Q 
lkjo 

About

Reconstruction from edge image combined with color and gradient difference for industrial surface anomaly detection

Resources

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Languages