DeepLabv3＋を自前のデータセットで学習させる

DeepLabv３＋をオリジナルデータセットでトレーニングできます。
TensorFlow公式モデルをつかいます。
DeepLabの使い方は基本的に公式リポジトリに書いてあります。
わからないところがあったらこの記事など読んでください。

　
画像はPASCAL VOCデータセットですが、自分のデータセットでもちゃんとトレーニングできます。

手順

１、モデル、モジュールを設定

１、公式モデルリポジトリをクローン

git clone https://github.com/tensorflow/models

２、researchディレクトリで作業します。

cd models/research

３、Colabで作業する場合はTensorFlow１を選択。

%tensorflow_version 1.x

４、tf_slimをインストール。

pip install tf_slim

必要なモジュールは、
Numpy
Pillow 1.0
tf Slim (which is included in the "tensorflow/models/research/" checkout)
Jupyter notebook
Matplotlib
Tensorflow
です。入れてない場合はインストールしてください。

５、tensorflow / models / research /ディレクトリをPYTHONPATHに追加して、ライブラリを使えるようにする。

# tensorflow/models/research/　から実行
export PYTHONPATH=$PYTHONPATH:`pwd`:`pwd`/slim

*ちなみに、Colabのパスの通し方がわからなかったので、僕はスクリプトのfrom deeplab.。。。import 。。。のdeeplabの部分を全部消しました。。。
あと、research/slimからnetをdeeplabに移動させました。。。

２、データを用意する

以下を用意します。

１、元画像

２、ラベル画像

輪郭線はあってもなくてもかまいません。

・画像要件

・それぞれディレクトリにまとめる。
　元画像のディレクトリ名：JPEGImage
　ラベル画像のディレクトリ名：SegmetationClass
・元画像とラベル画像の数と名前（拡張子抜きの名前）は同じにする
・画像サイズは任意。トレーニング時に513,513にクロップされることを念頭に。
（画像全体を収めたい場合は513,513以下を推奨）

セグメンテーションデータ作りにはlabelmeなどのツールが使えるそうです。

３、データを前処理する

１、ラベル画像がカラーの場合は白黒のラベル画像にします。

input_dir = "{データセットのパス}/SegmentationClass"
mkdir "{データセットのパス}/SegmentationClassRaw"
output_dir = "{データセットのパス}/SegmentationClassRaw" # SegmentationClassRawディレクトリを作っておく

python deeplab/datasets/remove_gt_colormap.py \
 --original_gt_folder=input_dir \
 --output_dir=output_dir

SegmentationClassRawディレクトリに白黒に変換した画像が保存されます。
小さなラベル値の白黒画像になるので、ほぼ真っ黒です。

２、画像ファイル名一覧のテキストファイルを作ります。

image0
image1
image3

という拡張子抜きの形式で、ファイル名一覧のテキストファイルを作ります。
以下の３テキストファイルを作ります。
１、trainval.txt ：全画像ファイル名
２、train.txt ：trainval.txtのうち、トレーニング用セット（全ファイルの９割くらいを割り当てる？）
３、val.txt ：trainval.txtのうち、内検証用セット（全ファイルの１割くらいを割り当てる？）

トレーニングと検証の割り振りは、最適なバランスは諸説あるようです。

テキストの抜き出し方はこちらの記事など参考にしてください。
ファイル名一覧をテキスト（.txt）ファイルに書き込む【Python】

３、データセットをTFRecord形式にします。

TensorFlowでのトレーニングで効率的に読み込めるTFRcord形式にします。

python deeplab/datasets/build_voc2012_data.py \
  --image_folder="{データセットのパス}/JPEGImages" \
  --semantic_segmentation_folder="{データセットのパス}/SegmentationClassRaw" \
  --list_folder="{データセットのパス}//ImageSets/Segmentation" \
  --image_format="jpg" \
  --output_dir="{データセットのパス}/tfrecord/"

４、トレーニングの設定をする

１、deeplab/datasets/data_generator.pyを自前データに合わせて書き換えます。

data_generator.py

_PASCAL_VOC_SEG_INFORMATION = DatasetDescriptor(
    splits_to_sizes={
        'train': 4552, # トレーニングイメージファイル数に書き換える
        'train_aug': 10582, #　そのままでOK
        'trainval': 5088, # 全イメージファイル数に書き換える
        'val': 536, # 検証イメージファイル数に書き換える
    },
    num_classes=21, #　そのままでOK
    ignore_label=255, #　そのままでOK
)

２、事前学習済みモデルをダウンロードします。

転移学習のために事前学習済みモデルの重みを使います。
公式リポジトリのリンクから好みのチェックポイントをダウンロードします。
MobileNetv２ベース（22MB）とXceptionベース（439MB）があります。

５、トレーニング

１、MobileNetv2バックボーンの場合

python deeplab/train.py --logtostderr \
   --training_number_of_steps=30000 \
   --train_split="train" \
   --model_variant="mobilenet_v2" \
   --output_stride=16 \
   --decoder_output_stride=4 \
   --train_crop_size="513,513" \
   --train_batch_size=1 \
   --dataset="pascal_voc_seg" \
   --tf_initial_checkpoint="{チェックポイントのパス}/deeplabv3_mnv2_pascal_train_aug/model.ckpt-30000" \
   --train_logdir="{データセットのパス}/checkpoint" \ # 先にチェックポイント書き込み先ディレクトリを作っておきます
   --dataset_dir="{データセットのパス}/tfrecord" \
   --fine_tune_batch_norm=false \ # CPUのみの場合は”true”
   --initialize_last_layer=true \
   --last_layers_contain_logits_only=false

２、Xception_65バックボーンの場合

python deeplab/train.py --logtostderr \
   --training_number_of_steps=30000 \
   --train_split="train" \
   --model_variant="xception_65" \
   --atrous_rates=6 \
   --atrous_rates=12 \
   --atrous_rates=18 \
   --output_stride=16 \
   --decoder_output_stride=4 \
   --train_crop_size="513,513" \
   --train_batch_size=1 \
   --dataset="pascal_voc_seg" \
   --tf_initial_checkpoint="{チェックポイントのパス}/deeplabv3_pascal_train_aug/model.ckpt" \
   --train_logdir="{データセットのパス}/checkpoint" \ # 先にチェックポイント書き込み先ディレクトリを作っておきます
   --dataset_dir="{データセットのパス}/tfrecord" \
   --fine_tune_batch_norm=false \ # CPUのみの場合は”true”
   --initialize_last_layer=true \
   --last_layers_contain_logits_only=false

1200秒（20分）に一回、チェックポイントが保存されます。save_interval_secsで保存インターバルを指定できます。トレーニング終了後も保存されます。

＊参考結果

５０００画像、３００００エポック、MobileNetv2バックボーン、ColabGPUでトレーニングしたところ、１時間ぐらいでした。
セグメントしたい物体が一つ（２ラベル）だったせいか、結果は大満足でした。

6、テスト

検証セットを使ってテストします。

元画像とカラーセグメンテーションマップ画像がログディレクトリに保存されます。

python deeplab/vis.py --logtostderr \
  --vis_split="val" \
  --model_variant="mobilenet_v2" \
  --output_stride=16 \
  --decoder_output_stride=4 \
  --vis_crop_size="513,513" \
  --dataset="pascal_voc_seg" \
  --checkpoint_dir="{チェックポイントのパス}/checkpoint" \
  --vis_logdir="{結果画像の書き込み先ディレクトリパス（作っておく）}" \
  --dataset_dir="{データセットのパス}/tfrecord" \
  --max_number_of_iterations=1 --eval_interval_secs=0

凍結グラフにして保存する

python export_model.py --model_variant="mobilenet_v2" --checkpoint_path="{チェックポイントを保存したディレクトリ}/model.ckpt-30000" --export_path="{凍結グラフを保存するディレクトリ}/frozen_graph.pb"

🐣

---

お仕事のご相談こちらまで
[e-mail]

Core MLを使ったアプリを作っています。
機械学習関連の情報を発信しています。

Twitter
Medium