deeplab v3+でオリジナルデータを学習させる

deeplab v3+ で自分のデータセットを使ってセグメンテーション出来るよう学習させてみました。
deeplab v3+のモデルと詳しい説明はこちら
github.com

データセットの準備

まず学習させるためのデータセットを作成します。
今回は画像中の木をセグメンテーションすることにしました。なので、背景領域と木領域の2クラスに分類することを学習させます。

用意するもの

①元の画像
②ラベル画像　を１セットとしたデータセットを用意します。


元の画像はjpg、ラベル画像はpng(モノクロ画像)で用意しました。
ラベル画像は、背景をrgb(0,0,0)、木領域をrgb(1,1,1)としてペイントソフトで手動で作成しました。もし２クラス以上にに分類したいときはr(n,n,n)でさらに塗り分ければ良いです。
ただし初期設定ではrgb(255,255,255)は無視されます。これはsegmentation_dataset.py で設定されているので変更可能です。

元画像とラベル画像は必ず同じ名前(例：1.jpgと1.png)をつけて保存します。

フォルダ分け

用意したデータセットをフォルダ分けします。
/deeplab/datasets/にoriginalという名前で新規ファイルを用意します。
そして以下のようなフォルダ群を新規作成します。

original
-img
  -train
  -trainval
  -val
-lbl
  -train
  -trainval
  -val
-lst
  -train
  -trainval
  -val

imgのtrainには訓練用画像、trainvalには訓練時の評価用画像、valには評価画像を入れます。
lblにもimgで入れた画像とセットとなるラベル画像をそれぞれ入れます。
lstには画像の名前を並べたテキストファイルを用意します。

例えばimg/train、val/trainに以下のような画像セットを入れていた場合は

このように拡張子を除いた画像の名前を羅列したテキストファイルを作成してlst/trainに保存します。

TFRecordに変換

データセットが用意できたら、次はTensorFlow用の形式である「TFRecord」に変換する必要があります。
これは/deeplab/datasets/build_voc2012_data.pyを利用して変更します。

/deeplab/datasets/download_and_convert_voc2012.shを見てみると、以下のような記述があります。

echo "Converting PASCAL VOC 2012 dataset..."
python ./build_voc2012_data.py \
  --image_folder="${IMAGE_FOLDER}" \
  --semantic_segmentation_folder="${SEMANTIC_SEG_FOLDER}" \
  --list_folder="${LIST_FOLDER}" \
  --image_format="jpg" \
  --output_dir="${OUTPUT_DIR}"

これを参考にしてオリジナルデータセットをTFRecordに変換するconvert_dataset.shを作成します。
以下のように記述すると訓練データのデータセットを変換出来ます。

python ./build_voc2012_data.py \
  --image_folder="./original/img/train" \
  --semantic_segmentation_folder="./original/lbl/train" \
  --list_folder="./original/lst/train" \
  --image_format="jpg" \
  --output_dir="./original/tfrecord" \

最後のoutput_dirでどこに変換後のTFRecordファイルを保存するか指定しています。
指定したフォルダに変換後のファイルが作成されていればOKです。

設定の追加

segmentation_dataset.pyにデータセットの設定が記載されているので、ここに自分で作成したデータセットの設定を記述します。
以下のように新しいデータセットの内容を記述します。

_ORIGINAL_INFORMATION = DatasetDescriptor(
    splits_to_sizes={
        'train': 50,  # num of samples in images/training
        'val': 10,  # num of samples in images/validation
    },
    num_classes=3,
    ignore_label=255,
)

また、下記のように、'original'という設定を_DATASETS_INFORMATIONに追加します。
これで学習実行時にデータセットにoriginalを指定すると、自作のデータセット設定を適応します。

_DATASETS_INFORMATION = {
    'cityscapes': _CITYSCAPES_INFORMATION,
    'pascal_voc_seg': _PASCAL_VOC_SEG_INFORMATION,
    'ade20k': _ADE20K_INFORMATION,
    'original':_ORIGINAL_INFORMATION
}

学習

準備が完了したので学習を行ってみます。
学習はdeeplab/test_local.shのファイルの一部を変更して行います。
このtest_local.shは、実行するとPASCALデータセットを用いた学習を行います。
これを自作のデータセットを参照するようにフォルダパス等を変更します。

フォルダパスを変更した後、--dataset=　の部分を_DATASETS_INFORMATIONに新規追加したoriginを参照するように変更します。

--dataset = origin

これでtest_local.shを実行することによって自作データセットで学習が行えます。

結果

訓練画像60枚で500エポック程度学習を回した結果