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

• xuebinqin/U-2-Net

Author

高橋かずひと(https://twitter.com/KzhtTkhs)

License

U-2-Net-ONNX-Sample is under Apache 2.0 License.

License(Image)

猫の画像と女性の画像はフリー素材ぱくたそ様の写真を利用しています。