[Japanese/English]

Warning

モデルを訓練するために使用したデータセットは自前で収集したものです。
背景や服の色、肌の色によって認識率が大きく下がる可能性があります。
モデルの使用に際し、損害や問題が発生しても、作成者(高橋)は一切責任を負いません。

Note

イラストにも対応するようにしましたが、データセット数が不足しているため、精度はイマイチです。

Note

このリポジトリは試験的なリポジトリです。
人物のセグメンテーションを実施したい場合、MediaPipe：SelfieMulticlass をおすすめします。

Skin-Clothes-Hair-Segmentation-using-SMP

3クラス(肌、服、髪)のセマンティックセグメンテーションを実施するモデルです。
Segmentation Models Pytorchを使用しています。
　
　　

本リポジトリでは、以下3種類のモデルのpthファイルとonnxファイルを用意しています。

• DeepLabV3+
  エンコーダー：timm-mobilenetv3_small_100
  パラメータ数：約216万
• DeepLabV3+
  エンコーダー：timm-mobilenetv3_large_100
  パラメータ数：約471万

Requirement

• torch 1.9.0 or later
• torchsummary 1.5.1 or later
• albumentations 1.0.3 or later
• matplotlib 3.2.2 or later
• onnx-simplifier 0.3.6 later
• opencv-python 3.4.2 or later
• onnxruntime 1.8.1 or later

Requirement(demo_onnx_xxxx.py)

demo_onnx_simple.py や demo_onnx_image_overlay.py のみを使用する場合は以下2つをインストールしてください。

• opencv-python 3.4.2 or later
• onnxruntime 1.8.1 or later

Dataset

自前で撮影した画像やインターネット上から収集した画像を合計452枚使用しています。
データセットは非公開です。

アノテーションはGrabCut-Annotation-Toolを用いて実施しており、クラスの割り当ては以下の通りです。

• クラスID 1：肌
• クラスID 2：服
• クラスID 3：髪

Training

Google Colaboratoryで [Colaboratory]Skin_Clothes_Hair_Segmentation_using_SMP.ipynb を開き、上から順に実行してください。
ノートブックが実行できるように、数枚のデータセットをコミットしてありますが、あくまで学習サンプルなので、
本リポジトリで公開しているモデルの精度には及びません。

Demo

デモの実行方法は以下です。

python demo_onnx_simple.py

python demo_onnx_image_overlay.py

• --device
  カメラデバイス番号の指定
  デフォルト：0
• --video
  動画ファイルの指定 ※動画指定時はカメラより優先
  デフォルト：None
• --image
  画像ファイルの指定 ※画像指定時はカメラより優先
  デフォルト：None
• --width
  カメラキャプチャ時の横幅
  デフォルト：960
• --height
  カメラキャプチャ時の縦幅
  デフォルト：540
• --mirror
  画像を反転表示するか否か
  デフォルト：指定なし
• --model
  使用するモデル
  デフォルト：'02.model/DeepLabV3Plus(timm-mobilenetv3_small_100)_452_2.16M_0.8385/best_model_simplifier.onnx'
• --score
  セマンティックセグメンテーションの判定閾値
  デフォルト：0.5

Reference

• Segmentation Models Pytorch
• GrabCut-Annotation-Tool

Author

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

License

Skin-Clothes-Hair-Segmentation-using-SMP is under MIT License.

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