注目物体検出のU^2-Netやーる（Windows10、Python3.6）

はじめに

注目物体検出のU^2-Net (U square net)をやってみました。CPUで動くよ。

システム環境

• Windows10（RTX2080 Max-Q、i7-8750H、RAM16GB）
• Anaconda 2020.02
• Python 3.6

導入

U^2-Netをクローンします。

U^2-Net用の環境を作成します。

conda create -n u2net python=3.6
conda activate u2net
cd U-2-Net-master
pip install numpy==1.15.2
pip install scikit-image==0.14.0
pip install Pillow==5.2.0
pip install scipy
pip install torch==1.0.0 torchvision==0.2.1 -f https://download.pytorch.org/whl/torch_stable.html
pip install matplotlib

u2net.pthをsaved_models/u2net/に、u2netp.pthをsaved_models/u2netp/に置きます。

86行目でCPUを指定します。

net.load_state_dict(torch.load(model_dir, map_location={'cuda:0': 'cpu'}))

test_data\test_imagesフォルダに入力画像を置きます。
test_data\u2net_results\にtest_imagesフォルダを作成します。出力画像がここに保存されます。

下記を実行します。

python u2net_test.py

before

after

お疲れ様でした。

おまけ

Background-Mattingとの比較

U^2-Net	Background-Matting

グリーンバックはもういらない！？Background-Mattingでどこでも合成（Windows10、Python 3.6）
https://qiita.com/SatoshiGachiFujimoto/items/f5583a89f751f88fbac4