U2Net : シングルショットで物体の切り抜きを行う機械学習モデル

ailia SDKで使用できる機械学習モデルである「U2Net」のご紹介です。エッジ向け推論フレームワークであるailia SDKとailia MODELSに公開されている機械学習モデルを使用することで、簡単にAIの機能をアプリケーションに実装することができます。

U2Netの概要

U2Netはシングルショットで物体の切り抜きを行うことができる機械学習モデルです。人や猫などの画像を入力として、全景と背景を分離するためのアルファ値を計算することができます。U2Netにはモデルサイズが176.3MBのU2Netと、モデルサイズが4.7MBのU2NetPがあります。

U$^2$-Net: Going Deeper with Nested U-Structure for Salient Object Detection

（出典：https://github.com/NathanUA/U-2-Net）

U2Netのアーキテクチャ

U2NetはUNetを階層化した上で、直列に2つ接続した構造となっています。従来のモデルは、BackboneにImageNetで事前学習したClassifierのモデルを使用することが多かったのですが、背景分離タスクに最適化したモデル構造を提案し、SOTAに匹敵するパフォーマンスを出しています。

（出典：https://arxiv.org/abs/2005.09007）

U2Netは320x320の解像度で学習を行なっています。

NathanUA/U-2-Net

U2NetとU2NetPはネットワークアーキテクチャは同じで入力と出力のFeatureMapの枚数が異なります。U2NetではEn.5のFeatureMapが512まで広がりますが、U2NetPでは64で抑えられています。Convolutionの重みのサイズはFeatureMapの枚数に大きく依存するため、FeatureMapの枚数を抑制することでモデルサイズを1/37.5にすることができています。

ただし、パフォーマンスの改善については、論文で示されているように、GeforceGTX 1080TiでU2Netの30FPSが、U2NetPで40FPSと、1.33倍に留まっています。

（出典：https://arxiv.org/abs/2005.09007）

学習はSOD（Salient Object Dataset）を使用しています。

Salient Objects Dataset (SOD) — Elder Laboratory

U2Netの使用方法

ailia SDKでは下記のサンプルでU2Netを使用可能です。

ailia-models/background_removal/u2net at master · axinc-ai/ailia-models

下記のコマンドで、U2Netを使用して、WEBカメラを入力とした人物切り抜きを行うことができます。

python3 u2net.py -v 0

また、任意のビデオファイルに人物切り抜きを行いファイルに保存することもできます。

python3 u2net.py -v input.mp4 -s output.mp4

実行例です。

-a smallを指定すると、モデルサイズの小さなU2NetPを使用することができます。

python3 u2net.py -v 0 -a small

ailia SDKはPythonだけでなく、CやC#から使用することもできます、下記は、XcodeからU2Netを使用するサンプルとなっています。

axinc-ai/ailia-xcode

ax株式会社はAIを実用化する会社として、クロスプラットフォームでGPUを使用した高速な推論を行うことができるailia SDKを開発しています。ax株式会社ではコンサルティングからモデル作成、SDKの提供、AIを利用したアプリ・システム開発、サポートまで、 AIに関するトータルソリューションを提供していますのでお気軽にお問い合わせください。