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

Kazuki Kyakuno
axinc
Published in
5 min readJul 28, 2020

--

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

U2Netの概要

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

(出典:https://github.com/NathanUA/U-2-Net

U2Netのアーキテクチャ

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

(出典:https://arxiv.org/abs/2005.09007

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

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)を使用しています。

U2Netの使用方法

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

下記のコマンドで、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を使用するサンプルとなっています。

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

--

--

Kazuki Kyakuno
axinc

CTO — ax Inc. , Managing Director — AXELL