flyinghu123 / CPR

This is an official implementation of the paper : "Target before Shooting: Accurate Anomaly Detection and Localization under One Millisecond via Cascade Patch Retrieval" (Accepted by IEEE TIP)

CPR

Official PyTorch implementation of CPR

Datasets

We use the MVTec AD dataset for experiments. And use DTD data set to simulate anomalous image.

The data directory is as follows:

data
├── dtd
│   ├── images
│   ├── imdb
│   └── labels
└── mvtec
    ├── bottle
    │   ├── ground_truth
    │   ├── license.txt
    │   ├── readme.txt
    │   ├── test
    │   └── train
    ...
    └── zipper
        ├── ground_truth
        ├── license.txt
        ├── readme.txt
        ├── test
        └── train

Installation

pytorch

conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia

pip install -r requirements.txt

generate foreground and global retrieval result

python tools/generate_foreground.py
python tools/generate_retrieval.py

Training

generate synthetic data

python tools/generate_synthetic_data.py -fd log/foreground_mvtec_DenseNet_features.denseblock1_320

bash train.sh

Testing

python test.py -fd log/foreground_mvtec_DenseNet_features.denseblock1_320/ --checkpoints weights/{category}.pth

Pretrained Checkpoints

Download pretrained checkpoints here and put the checkpoints under <project_dir>/weights/.

Baidu Netdisk: https://pan.baidu.com/s/1FTE4b2G8nVZt4lUyaP-kIQ?pwd=ky7j

Acknowledgement

We borrow some codes from PatchCore, MemSeg and SuperPoint

Citation

@misc{li2023target,
      title={Target before Shooting: Accurate Anomaly Detection and Localization under One Millisecond via Cascade Patch Retrieval}, 
      author={Hanxi Li and Jianfei Hu and Bo Li and Hao Chen and Yongbin Zheng and Chunhua Shen},
      year={2023},
      eprint={2308.06748},
      archivePrefix={arXiv},
      primaryClass={cs.CV}
}