Industrial KNN-based Anomaly Detection

⭐ Now has streamlit support! ⭐ Run $ streamlit run streamlit_app.py

This repo aims to reproduce the results of the following KNN-based anomaly detection methods:

1. SPADE (Cohen et al. 2021) - knn in z-space and distance to feature maps
2. PaDiM* (Defard et al. 2020) - distance to multivariate Gaussian of feature maps
3. PatchCore (Roth et al. 2021) - knn distance to avgpooled feature maps

* actually does not have any knn mechanism, but shares many things implementation-wise.

Install

$ pipenv install -r requirements.txt

Note: I used torch cu11 wheels.

Usage

CLI:

$ python indad/run.py METHOD [--dataset DATASET]

Results can be found under ./results/.

Code example:

from indad.model import SPADE

model = SPADE(k=5, backbone_name="resnet18")

# feed healthy dataset
model.fit(...)

# get predictions
img_lvl_anom_score, pxl_lvl_anom_score = model.predict(...)

Custom datasets

📂datasets
 ┗ 📂your_custom_dataset
  ┣ 📂 ground_truth/defective
  ┃ ┣ 📂 defect_type_1
  ┃ ┗ 📂 defect_type_2
  ┣ 📂 test
  ┃ ┣ 📂 defect_type_1
  ┃ ┣ 📂 defect_type_2
  ┃ ┗ 📂 good
  ┗ 📂 train/good

$ python indad/run.py METHOD --dataset your_custom_dataset

Results

📝 = paper, 👇 = this repo

Image-level

Pixel-level

PatchCore-10 was used.

Hyperparams

The following parameters were used to calculate the results. They more or less correspond to the parameters used in the papers.

spade:
  backbone: wide_resnet50_2
  k: 50
padim:
  backbone: wide_resnet50_2
  d_reduced: 250
  epsilon: 0.04
patchcore:
  backbone: wide_resnet50_2
  f_coreset: 0.1
  n_reweight: 3

Progress

• Datasets
• Code skeleton
• Config files
• CLI
• Logging
• SPADE
• PADIM
• PatchCore
• Add custom dataset option
• Add dataset progress bar
• Add schematics
• Unit tests

Design considerations

• Data is processed in single images to avoid batch statistics interference.
• I decided to implement greedy kcenter from scratch and there is room for improvement.
• torch.nn.AdaptiveAvgPool2d for feature map resizing, torch.nn.functional.interpolate for score map resizing.
• GPU is used for backbones and coreset selection. GPU coreset selection currently runs at:
  • 400-500 it/s @ float32 (RTX3080)
  • 1000+ it/s @ float16 (RTX3080)

Acknowledgements

• hcw-00 for tipping sklearn.random_projection.SparseRandomProjection

References

SPADE:

@misc{cohen2021subimage,
      title={Sub-Image Anomaly Detection with Deep Pyramid Correspondences}, 
      author={Niv Cohen and Yedid Hoshen},
      year={2021},
      eprint={2005.02357},
      archivePrefix={arXiv},
      primaryClass={cs.CV}
}

PaDiM:

@misc{defard2020padim,
      title={PaDiM: a Patch Distribution Modeling Framework for Anomaly Detection and Localization}, 
      author={Thomas Defard and Aleksandr Setkov and Angelique Loesch and Romaric Audigier},
      year={2020},
      eprint={2011.08785},
      archivePrefix={arXiv},
      primaryClass={cs.CV}
}

PatchCore:

@misc{roth2021total,
      title={Towards Total Recall in Industrial Anomaly Detection}, 
      author={Karsten Roth and Latha Pemula and Joaquin Zepeda and Bernhard Schölkopf and Thomas Brox and Peter Gehler},
      year={2021},
      eprint={2106.08265},
      archivePrefix={arXiv},
      primaryClass={cs.CV}
}