今回調査した背景
皆さん御存知の通り、画像認識の技術は凄まじいスピードで発展しており、応用事例も日々増えています。それもあってGPUの需要はうなぎのぼりで、先日のNVidiaの決算で4半期純利益 3兆円らしくうらやましい限りです。ちなみに3兆円はリトアニアの歳入と大体同じになるらしいです。弊社の中心事業領域でもある画像の異常検知の技術発展も目覚ましく、一年に何回も検知性能のSOTAの報告がなされています。
ただ、PatchCoreなど学会や論文でよく見るAIモデルはあるものの、実際の工場の外観検査でデファクトスタンダードとして使われているものは今のところ見当たりません。その一つの理由としてはAIを使った異常検知にはGPUを使った重たい計算が必要となるのが一般的で、制限のある環境下で高速なラインを流れる工業製品の検査に用いるのは難しい部分が多いからと考えています。また、部品の大きさに対して小さな欠陥を取りたいというお困りごとをよくお聞きするのですが、高画素の画像に対する検査というのはGPUメモリの制約もありほとんど報告されていません。512x512サイズの画像に対する検査方法の提案がCVPR2024で採択されるなどまだ高画素画像に対応したAIは研究段階と言えます。
我が社ではこの状況下であえて古い特徴量のHLACを技術の中心に添えています。最新の技術と比べると検出性能が劣る場面もあるのですが、比較的高い検出性能を持ちながら10Mピクセルほどの高解像度の画像でも高速な検査が可能という優れた特性を持ち合わせています。学習枚数は50枚程度と比較的少なくて学習時間も短いため、現場でモデルの修正が必要となったときも素早く対応ができるという強みも持ち合わせています。
最近よく聞くお困りごととして、1品種は少量しか作られないが多くの品種があるラインの検査の自動化のご相談を受けます。1品種100個など超少量のケースの場合、弊社の技術でもチャレンジングな対象となっており、現在技術的に強化したい状況の一つです。そこで今回はfew-shotの異常検知モデルの技術調査を行い、詳細を紹介したいと思います。
今回の調査の方針としては弊社技術に近い特性を持ち合わせているものを対象とするため、
- few-shotで学習可能
- 学習時間が短い
- プロンプトエンジニアリングが必要ない
という選定基準で探しました。3番については検査の現場での扱いやすさという観点から、画像のみで性能調整ができるモデルを対象としました。今回はこの方針で調べたRegADというモデルについて紹介したいと思います。
RegADの紹介
RegADはECCV2022の論文で発表されたモデルです。
paper with codeでは2024年12月現在で95位につけており、few-shotに限ると3番目に出て来るモデルです。トップオブトップの性能ではないものの、数枚の検査対象に素早くアダプテーションできるという特性を持ち合わせています。
画像の位置合わせができる、つまり同カテゴリー画像からなら同じ特徴量が得られるネットワークなら、ファインチューニングなしでもカテゴリーに依存しないものを作れるというのがこの論文の基本アイデアです。この考えの下、まずは同じカテゴリーの正常画像から近い特徴量が得られるFeature Registration Networkを作ります。その後、実際の検査対象の数枚の正常画像からFeature Registration Networkで得られる特徴量の分布を推定します。検査の中では正常画像の分布から離れた特徴量が出る場合に異常と判定します。
画像の位置合わせができる、つまり同カテゴリー画像からなら同じ特徴量が得られるネットワークなら、ファインチューニングなしでもカテゴリーに依存しないものを作れるというのがこの論文の基本アイデアです。この考えの下、まずは同じカテゴリーの正常画像から近い特徴量が得られるFeature Registration Networkを作ります。その後、実際の検査対象の数枚の正常画像からFeature Registration Networkで得られる特徴量の分布を推定します。検査の中では正常画像の分布から離れた特徴量が出る場合に異常と判定します。
Feature Registration Network
今回のネットワークはSimSiam をベースにしたcontrastive learning(学習時に違うカテゴリーであるnegative画像使わないからnon-contrastive learningと呼ばれることもあるらしいです。ややこしい。。)をするものとなります。
こちらの学習で同じカテゴリーの画像は同じ特徴量が得られるネットワークを作ります。論文の中ではfew-shot learningさせたいカテゴリ以外の画像を使ってネットワークの学習を行っていました。ベースとしてはImageNetで事前学習したresnetを使っています。
本モデルではSpatial Transformer Network(STN)をネットワークの途中に入れるという工夫をしています。STNはタスクに適切なaffine変換を特徴に入れられるネットワークです。この工夫により回転、平行移動したときに同じような絵になる（つまり同じカテゴリーの画像)ときに特徴量がより揃えられるようです。
もとのSimSiamでもオーグメンテーションした画像で学習をすれば同様の性質は得られそうな気もしたのですが、論文のablation studyによるとSTNにより性能が上がるようです。MVTecのデータセットでの性能向上は小さいのですが、MPDDでは1割ほどと結構性能が向上しています。
few shotの正常画像特徴量分布の推定
正常の分布推定はPaDiMと同様の方法で行われます。
まずは正常の画像にrotation, translation, flipping, grayingのオーグメンテーションを加えて数を増やし、Feature Registration Networkに入れて特徴量を得ます。そして高さと幅の次元方向に特徴量テンソルを分割した個々の領域をパッチとして、パッチごとに特徴量のセットを作ります。この特徴量セットに対して、ガウス分布をフィットしてマハラノビス距離が測れるようにします。このような単純な分布推定を使うため、アダプテーションがとても高速です。
正常異常の判定
正常異常判定のために検査画像に対して各パッチのマハラノビス距離を測り、パッチの位置に合わせて配置して距離のマトリクスを作ります。そのマトリクスにSTNで行ったaffine変換の逆変換をかけて元画像に特徴量の位置を合わせます(公式のgithubのコードだとそうなっていない気がしますが、ひとまず論文の記述通りの説明にします)。その後、距離マトリクスを線形補間で画像のサイズまで引き伸ばしたあと、ガウスブラーによるスムージングをしてピクセルごとの異常値を求めます。ピクセル異常値の最大値をその画像の異常値と定義して、しきい値以上の異常値が出たら異常、それ以外を正常と判定します。
検査性能
MVTecデータセットの中から位置合わせがほぼいらなそうなwood と位置合わせが重要そうなscrewについて性能を見ていきたいと思います。イメージの異常値についてのROC-AUCは下のようになります。数字はアダプテーションする画像の10個の組み合わせに対する結果の平均です。
| 2-shot | 8-shot | |
| wood | 99.2 | 99.4 |
| screw | 52.5 | 63.4 |
これを見ると、screwではあまり性能が出ないみたいですね。公式githubのコードで実行すると、ピクセルごとの異常値のROC-AUCは95ほど出ているので、画像中の最大値のみを使って判定しているのがあまり良くないのかもしれません。woodではイメージ、ピクセルともにAUCが高いので論文の画像評価方法でもワークはするケースも有るようですが、画像評価方法の良し悪しは気になるところです。woodはあまり位置合わせがいらないケースであることも効いているかもしれないので、STNの効果など位置合わせの影響についてはもう少し知りたいと思いました。
google colabのT100のGPUで公式のコードを試してみたところ、アダプテーションの時間は8shotで大体5秒位でできました。用途にもよると思うのですが、多品種を切り替えながら学習・検査するときにこれくらいならストレスはないと感じました。
検査時間についてはアダプテーションと同じ環境で、一枚あたり150 msもかからないくらいのスループットが出せました。こちらも用途よって適用可能性は変わると思いますが、これまでお聞きしたことのあるお困りごとから考えると実際に使えるユースケースケースは多いと思います。
気になりごと
RegADの検査の中ではパッチごとにマハラノビス距離を出していますが、一つだけしきい値を設定します。このやり方で検査性能を出すにはパッチごとの異常値分布が揃っている必要が出てきますが、実際どうなっているのか気になったので調べてみました。
下のグラフは横軸パッチの異常値で縦軸アダプテーションに使っていない画像中パッチの累積密度です。距離の平均値が最大値、中間値、最小値となる3パッチの結果を重ねてみると、結構ズレているのがわかります。。screwに至っては分布の形自体も揃っていないので、もしかするとこれもイメージのAUCが低くなる原因かもしれません。
これを見て、画像オーグメンテーションで推定している分布とテストデータの分布はどれくらい一致するのかも気になりました。この分布があっていないとすると、パッチごとの分布が一致しないのは仕方ない部分があります。下は一つ代表のパッチを決めて、アダプテーションに使った画像のマハラノビス距離の累積分布とそれ以外の正常画像の累積分布です。
これを見ると2shotだと結構ズレているのがわかりますね。。検定などのちゃんとした調査はしていませんが、8shotでもあっているとは言い難い状況です。この状態でも性能が出るのは興味深いですが、距離分布を合わせる方法はもう少し検討の余地があるように感じました。
あとがき
paper with codeを中心に見ていたため調査中には気づかなかったのですが、この論文の続編が今年(2024年)に出ていました(https://ieeexplore.ieee.org/document/10704738)。こちらの論文ではWasserstein距離を使った最適なオーグメンテーションの選び方が提案されており、上の推定分布とテスト画像の分布ずれに対する対策が提案されていました。
今後、新作の調査を進めて今回気になったポイントが解消されているのか確認したいと思います。また、今回の方法が実際の工場で流れている部品画像に対してどの程度の性能を示すのかを調べてみたいと思いました。次以降の投稿で実部品に対して弊社技術と比較したときの強み弱みを皆様にお伝えできたらと思います。
弊社では機械学習エンジニアを中心にこのような調査を行っています。現在下のポジションの募集を行っていますので、ご興味を持たれたらぜひご応募ください！
