U-2-Net-ONNX-Sample
U-2-NetのPythonでのONNX推論サンプルです。
ONNXに変換したモデルも同梱しています。
変換自体を試したい方はColaboratoryなどでU-2-Net-Convert2ONNX.ipynbを使用ください。
Requirement(ONNX変換)
- Pytorch 1.9.0 or later
- onnx 1.10.2 or later
- onnx-simplifier 0.3.6 or later
Requirement(ONNX推論)
- gdown 4.2.0 or later
- OpenCV 4.5.3.56 or later
- onnxruntime-gpu 1.9.0 or later
※onnxruntimeでも動作しますが、推論時間がかかるのでGPUをお勧めします
Demo(u2net、u2netp、u2net_human_seg)
デモの実行方法は以下です。
デモ実行前に、Googleドライブからu2net.onnx、u2netp.onnx、u2net_human_seg.onnxの何れかをダウンロードし、modelディレクトリに格納してください。
実行対象のモデルが格納されていない場合、gdownを用いたダウンロードが実行されます(gdownインストール必要)
python sample_u2net_onnx.py- --device
カメラデバイス番号の指定
デフォルト：0 - --movie
動画ファイルの指定 ※指定時はカメラデバイスより優先
デフォルト：指定なし - --width
カメラキャプチャ時の横幅
デフォルト：640 - --height
カメラキャプチャ時の縦幅
デフォルト：360 - --model
ロードするモデルの格納パス
デフォルト：model/u2net.onnx 選択肢：model/u2net.onnx、model/u2netp.onnx、model/u2net_human_seg.onnx - --input_shape
モデルの入力サイズ
デフォルト：320
Demo(u2net_portrait)
デモの実行方法は以下です。
デモ実行前に、Googleドライブからu2net_portrait.onnxの何れかをダウンロードし、modelディレクトリに格納してください。
u2net_portrait.onnxが格納されていない場合、gdownを用いたダウンロードが実行されます(gdownインストール必要)
python sample_u2net_onnx.py- --device
カメラデバイス番号の指定
デフォルト：0 - --movie
動画ファイルの指定 ※指定時はカメラデバイスより優先
デフォルト：指定なし - --width
カメラキャプチャ時の横幅
デフォルト：640 - --height
カメラキャプチャ時の縦幅
デフォルト：360 - --model
ロードするモデルの格納パス
デフォルト：model/u2net_portrait.onnx - --input_shape
モデルの入力サイズ
デフォルト：512
Reference
Author
高橋かずひと(https://twitter.com/KzhtTkhs)
License
U-2-Net-ONNX-Sample is under Apache 2.0 License.
License(Image)
猫の画像と女性の画像はフリー素材ぱくたそ様の写真を利用しています。