STFPM-Paddle

April 13, 2022 · View on GitHub

1 简介

This is the unofficial code based on PaddlePaddle of BMVC 2021 paper:
Student-Teacher Feature Pyramid Matching for Anomaly Detection

attention

STFPM算法由预先培训的教师网络和结构相同的学生网络组成。学生网络通过将特征与教师网络中的对应特征相匹配来学习无异常图像的分布。多尺度特征匹配用于增强鲁棒性。这种分层特征匹配使学生网络能够从特征金字塔接收多层次知识的混合,从而允许各种大小的异常检测。 在推理过程中,比较了教师网络和学生网络的特征金字塔。差异越大,异常发生的概率越高。 本项目基于PaddlePaddle框架复现了STFPM,并在MvTec数据集上进行了实验。 论文:

2 复现精度

在MvTec数据集的测试效果如下表。

Image-Level AUC

AvgCarpetGridLeatherTileWoodBottleCableCapsuleHazelnutMetal NutPillScrewToothbrushTransistorZipper
anomalib0.8930.9540.9820.9890.9490.9610.9790.8380.7590.9990.9560.7050.8350.9970.8530.645
Paddle0.9370.9720.9880.9990.9920.9931.0000.9330.8171.0000.9840.9330.8560.7720.9460.866

Pixel-Level AUC

AvgCarpetGridLeatherTileWoodBottleCableCapsuleHazelnutMetal NutPillScrewToothbrushTransistorZipper
anomalib0.9510.9860.9880.9910.9460.9490.9710.8980.9620.9810.9420.8780.9830.9830.8380.972
Paddle0.9670.9920.9880.9940.9590.9540.9880.9540.9790.9870.9740.9700.9820.9890.8170.982

image-level auc的Mean为0.937。 pixel-level auc的Mean为0.967。

达到验收指标。

训练日志:logs AIStudio预训练权重和日志:AIStudio预训练权重

3 数据集

数据集网站:MvTec数据集

AiStudio上的数据集:MVTec-AD

这里采用AiStudio上的数据集上的数据集,下载后用如下命令解压到自己指定的PATH/MVTec路径下:

tar xvf mvtec_anomaly_detection.tar.xz

4 环境依赖

  • 硬件: Tesla V100 16G >= 1

  • 框架:

    • PaddlePaddle >= 2.2.0

快速开始

第一步:克隆本项目

# clone this repo
git clone git@github.com:CuberrChen/STFPM-Paddle.git
cd STFPM-Paddle

第二步:训练模型

MVTec共有15个类别的子数据集,因此每个类别都需要单独训练一个模型; 在训练时,通过category参数来指定类别数据进行训练。depth表示resnet深度。mvtec_ad指定上述数据集路径PATH/MVTec 。 train_val表示是否在训练时开启指标计算。model_save_path指定模型保存路径,会在这个路径下生成category的目录。 seed表示随机数种子,这里根据给定参考项目取42.

开始训练: 以carpet为例:

python train.py train --depth=18 --mvtec_ad=PATH/MVTec --category carpet --epochs 100 --train_val=True --model_save_path=./output --seed 42

第三步:验证

需要指定训练好的模型参数路径checkpoint=output/carpet/best.pdparams

python val.py --depth=18 --mvtec_ad=PATH/MVTec --category carpet --checkpoint=output/carpet/best.pdparams

第四步:预测

这里需要指定单张图片路径image_path以及保存预测结果路径save_path,会在生成预测结果predict.png

python predict.py --checkpoint=output/carpet/best.pdparams --image_path=/home/aistudio/data/carpet/test/color/000.png --save_path=predict.png

原图:

原图

检测结果:

检测

第五步:TIPC

详细日志在test_tipc/output

TIPC: TIPC: test_tipc/README.md

首先安装auto_log,需要进行安装,安装方式如下: auto_log的详细介绍参考https://github.com/LDOUBLEV/AutoLog。

git clone https://github.com/LDOUBLEV/AutoLog
cd AutoLog/
pip3 install -r requirements.txt
python3 setup.py bdist_wheel
pip3 install ./dist/auto_log-1.2.0-py3-none-any.whl

进行TIPC:

bash test_tipc/prepare.sh test_tipc/configs/STFPM/train_infer_python.txt 'lite_train_lite_infer'

bash test_tipc/test_train_inference_python.sh test_tipc/configs/STFPM/train_infer_python.txt 'lite_train_lite_infer'

TIPC结果:

tipc测试结果截图

tipc测试结果

  • 预训练模型的静态图导出与推理测试:
python export_model.py --depth 18 --img_size=256 --model_path=./output/carpet/best.pdparams --save_dir=./output
python infer.py --use_gpu=True --model_file=output/model.pdmodel --input_file=/home/aistudio/data/carpet/test/color/000.png --params_file=output/model.pdiparams

可正常导出与推理。

5 代码结构与说明

代码结构

├── data
├── datasets
├── models
├── logs
├── output
├── test_tpic
├── README.md
├── train.py
├── val.py   
├── predict.py
├── export_model.py
├── infer.py
└── requirements.txt              

6 模型信息

相关信息:

信息描述
作者xbchen
日期2022年4月
框架版本PaddlePaddle==2.2.1
应用场景异常检测
硬件支持GPU、CPU
在线体验notebook

7 说明

  • 感谢百度提供的算力支持。