Skip to content

3クラス(肌、服、髪)のセマンティックセグメンテーションを実施するモデル(A model that performs semantic segmentation of 3 classes(skin, clothes, hair))

License

Notifications You must be signed in to change notification settings

Kazuhito00/Skin-Clothes-Hair-Segmentation-using-SMP

Repository files navigation

[Japanese/English]

⚠Attention⚠

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

Skin-Clothes-Hair-Segmentation-using-SMP

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

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

  • DeepLabV3+
    エンコーダー:timm-mobilenetv3_small_100
    パラメータ数:約216万
  • PAN(Pyramid Attention Network)
    エンコーダー:timm-mobilenetv3_small_100
    パラメータ数:約102万
  • U-Net++
    エンコーダー:timm-mobilenetv3_small_100
    パラメータ数:約371万

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
  • --file
    動画ファイルの指定 ※動画指定時はカメラより優先
    デフォルト: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

Author

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

License

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

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

About

3クラス(肌、服、髪)のセマンティックセグメンテーションを実施するモデル(A model that performs semantic segmentation of 3 classes(skin, clothes, hair))

Topics

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published