PaDiM : 再学習不要で不良品検知を行う機械学習モデル
ailia SDKで使用できる機械学習モデルである「PaDiM」のご紹介です。エッジ向け推論フレームワークであるailia SDKとailia MODELSに公開されている機械学習モデルを使用することで、簡単にAIの機能をアプリケーションに実装することができます。
PaDiMの概要
PaDiMは2020年11月に発表された、不良品検知を行う機械学習モデルです。正常品のみの画像から再学習不要で不良品検知が実装可能です。不良品検知のデータセットであるMVTec ADでSOTAを達成しています。
不良品検知のアーキテクチャ
機械学習で不良品検知を行う方法として、下記の3つの方法があります。
- 正常品と不良品から学習（正常品、不良品共に10000枚以上必要）
- 正常品のみから学習（正常品が10000枚以上必要）
- 正常品のみから画像の特徴量の分布を使って判定（正常品が240枚以上必要）
1.は通常のClassifierのモデルで、正常品と不良品の画像からResNetなどで分類モデルを学習し、不良品検知を行うものです。不良個所の可視化にはGradCamを使用します。また、近年では距離学習も使用されています。しかし、工場のラインでは、正常品の画像は多く取得できますが、不良品の画像はあまり取得できないという課題があります。
2.は、そのような背景で開発されたもので、正常品のみからAutoEncoderなどで学習を行います。AutoEncoderでは、入力された画像を低次元のベクトルに写像し、低次元のベクトルから元の画像を復元するモデルアーキテクチャです。正常品のみから学習したAutoEncoderは、正常品の特徴しか学習していないため、不良品の画像を入れても正常品の画像が出力されます。その特徴を利用して、入力画像とAutoEncoderの復元画像の差分を取ると異常個所を検出できます。近年では、GANも併用されています。しかし、製品ごとに機械学習の再学習が必要で負荷が高いという問題があります。
3.は、近年注目されている方法で、機械学習モデルを特徴抽出のみに使用します。ImageNetで学習された機械学習モデルの中間層のデータは、画像の特徴をよく示しています。そこで、入力された画像をResNetなどに通し、中間層のデータを抽出して特徴量とした上で、正常品の特徴量の分布と、評価画像の特徴量の分布の差分を取ることで不良品検知を行います。
PaDiMは3番目のアプローチに属しており、K-NNによるクラスタリングが必要だったSPADEのアーキテクチャを改良し、Gaussian Distributionだけのシンプルなアーキテクチャにも関わらず、SOTAを達成しました。AUROCで計算した精度は97.9です。
PaDiMのアーキテクチャ
PaDiMの学習では、正常品の画像列に対して、WideResNet50を適用し、Feature Mapを計算、特徴ベクトルを計算します。特徴ベクトルの計算においては、layer1〜3の浅い特徴量を使用します。特徴ベクトルは、Feature Mapの画像平面の各画素ごとに計算します。
計算した特徴ベクトルから、正常品の画像群の平均値を計算し、共分散行列を計算します。共分散行列は、スカラー値における分散の一般化となります。
平均と共分散行列から、正常品のN次元正規分布を計算することができます。
下記が特徴ベクトルの計算の図解です。N枚の正常品の画像から、(W,H)=(224,224)の解像度を持つCNNに入力し、Feature Mapの画像平面の各画素ごとに共分散行列Σと平均μを計算します。
学習結果は、共分散行列と平均となります。
PaDiMの推論では、テスト用の画像に対して、学習時と同様に特徴ベクトルを計算し、事前に計算しておいた共分散行列と平均からマハラノビス距離を計算します。マハラノビス距離は、普通の距離を一般化したもので、2つの確率変数ベクトルの非類似性を計算することができます。
各画素ごとにマハラノビス距離を計算し、マハラノビス距離を画像全体で正規化したものがヒートマップ、マハラノビス距離の最大値が異常度になります。これより、異常度が一定異常の場合、異常品と判定します。
PaDiMの演算量の削減
PaDiMの演算量を削減する方法として、論文では、特徴ベクトルの次元数を削減する方法や、異なるアーキテクチャのCNNを用いることが提案されています。次元数の削減においては、PCA（主成分分析）とランダム抽出を比較し、ランダム抽出の方が性能が高いことが示されています。
PaDiMの再現実装においては、resnet18では448次元の特徴ベクトルを100次元に、wide_resnet50_2では1792次元の特徴ベクトルを550次元に、ランダム抽出によって削減しています。
不良品検知のデータセット
不良品検知の評価用のデータセットにはMVTec ADが使用されています。
MVTec ADのデータセットは、trainフォルダのgoodに正常画像、testフォルダのgoodに正常画像、それ以外のフォルダ名に異常画像が格納されています。また、ground_truthに異常個所のマスクが格納されています。異常個所のマスクはpixel ROCAUCを計算するのに使用します。
アーキテクチャ上、浅い特徴量を使用するため、大きな位置ずれには弱くなっています。そのため、検出対象画像のセンタリングは必要です。
MVTec ADのデータセットは高精度にセンタリングされているため、論文では、実アプリケーションへの応用に向けて、学習時に256x256から224x224へのRandomCropと、+-10度のRandomRotationを入れる、Modified versionが提案されています。
PaDiMの使用方法
Pytorch
PaDiMのリポジトリをCloneします。
mvtec_anomaly_detectionにデータセットを配置します。
datasets/mvtec.pyのCLASS_NAMESを書き換えます。
CLASS_NAMES = ['grid']python3 main.pyを実行します。Pytorchの場合、7GB程度、メモリを使用します。MVTec ADのカーペットのデータセット（279枚）を使用して学習を行なった場合、MacBookPro 13で2分程度で学習が完了します。
$ python3 main.py
| feature extraction | train | grid |: 100%|██████| 9/9 [01:30<00:00, 10.08s/it]
| feature extraction | test | grid |: 100%|███████| 3/3 [00:28<00:00, 9.60s/it]
image ROCAUC: 0.957
pixel ROCAUC: 0.965
Average ROCAUC: 0.957
Average pixel ROCUAC: 0.965実行が完了すると、mvtec_resultに検出結果のヒートマップが格納されます。検出結果の例は下記となります。カーペットの不良を検知することができており、不良個所を可視化できています。
ailia SDK
ailia MODELSのPaDiMフォルダに移動します。
trainフォルダに正常品の画像を配置します。この例では、MVTec ADのbottleのgoodフォルダにある画像を使用しています。
学習作業を行います。
python3 padim.py --train_dir train学習結果がtrain.pklに格納されます。Resnet18の場合は127.9MBになります。
学習済みのtrain.pklを使用して、異常検知を行います。
python3 pading.py -i input.png -f train.pkl -th 0.5異常度がコンソールに表示されます。この値が大きいほど異常度が高くなります。
Anomaly score: 25.189980 #異常品
Anomaly score: 8.388268 #学習に使用した正常品
Anomaly score: 11.866211 #学習に使用していない正常品異常箇所がoutput.pngに出力されます。
異常箇所に関する最適な検出しきい値は、GroundTruthのマスク画像から算出することができます。gt_masksフォルダに、入力画像の名前に対応するマスク画像を入れておくことで、推奨しきい値が出力されます。
INFO padim.py (392) : Optimal threshold: 0.592428GroundTruthのマスク画像が存在しない場合は、output.pngのGroundTruthは空になります。
また、異常箇所については、画素ごとの異常度の最大値で正規化されてるため、正常品の場合は物体全体が異常になるような出力となります。画像全体の異常度が低い場合は異常箇所の画像は無視してください。
なお、入力画像の代わりにフォルダを与えることで、まとめて複数の画像の検証を行うことも可能です。
python3 pading.py -i input_folder -s output_folder -f train.pkl -th 0.5実行例です。
PaDiMのその他の応用
PaDiMは不良品検知の他に監視カメラの異常検知にも応用することが可能です。PaDiMはShanghaiTech Campus datasetでCAVGA（Convolutional Adversarial Variational autoencoder with Guided Attention ）の85%を超える91.2%のAUROCスコアを達成しています。
関連研究
MahalanobisADではEfficientNet B4の1〜7階層のFeatureMapにGlobalAveragePoolingを適用して取得したFeatureVectorに対してマハラノビス距離を計算します。MahalanobisADでは画像全体で一つのマハラノビス距離を計算します。PaDiMのように画素ごとにマハラノビス距離を計算しているわけではないため、異常箇所の可視化は行えません。異常箇所が必要な場合はGradCAMによってヒートマップを生成します。
ax株式会社はAIを実用化する会社として、クロスプラットフォームでGPUを使用した高速な推論を行うことができるailia SDKを開発しています。ax株式会社ではコンサルティングからモデル作成、SDKの提供、AIを利用したアプリ・システム開発、サポートまで、 AIに関するトータルソリューションを提供していますのでお気軽にお問い合わせください。
