Edited at

ESPNetで自作データセットを学習してセグメンテーション


初めに

この記事はまだまだ未完であることをご容赦ください。


ESPNetを使うことになった

深層学習を使ったsemantic segmentationの研究は多々あるんですが

ロボットなどに実装するためリアルタイム性を求めた軽量なモデルはそう多くないです。

有名なので言えばENet・ICNet・ERFNetなど。

その他FCN・SegNet・Deeplabなどはかなり重量級なモデル。

現在のところリアルタイム性と分類精度はトレードオフですが

今回使ったESPNetはそこそこの精度でかなり軽量なモデルになっています。

論文によればJetson TX2上で9fps(1024✕512)という数字を叩き出しています。

今までCaffe上でENetを使用していたが精度に限界を感じたため

上位互換のESPNetを使ってみることにした。

ESPNetはPytorchで実装されており、初めてPytorchを使うにあたって

備忘録として記事を書きました。CaffeもCaffe2に移行しつつあることや

使い方に癖があるためこれをいい機会にPyTorchに乗り換えようか。

github:ESPNet

論文:ESPNet: Efficient Spatial Pyramid of Dilated Convolutions for Semantic Segmentation


動作環境


  • Ubuntu 16.04

  • Intel(R) Core(TM) i7-6800K CPU @ 3.40GHz

  • GeForce GTX 1080

  • CUDA 8.0

  • cuDNN 7.0

  • Anaconda 3-5.3.1(python3.7)

  • PyTorch 0.3.0

  • OpenCV 3.3.0


環境構築

pyenv・Anaconda・Pytorchの導入は以下の記事と同じですが一応書いておきます

Ubuntu 16.10にPyTorchを導入してみた

PyTorch初心者なので記事に従っていますが、PyTorchを入れる段階で

Anaconda入れたくない人はやり方があるので他の記事でおねがいします。


1. pyenv


terminal

git clone https://github.com/yyuu/pyenv.git ~/.pyenv


PATHを通すため.profileもしくは.bashrcなどに以下を追加


.bashrc

#pyenv

export PYENV_ROOT=$HOME/.pyenv
export PATH=$PYENV_ROOT/bin:$PATH
eval "$(pyenv init -)"

設定を反映させましょう


terminal

source ~/.bashrc

もしくは
source ~/.profile


2. Anaconda

Anacondaの最新バージョンを確認してインストール


terminal

pyenv install -l | grep anaconda3

pyenv install anaconda3-5.0.1

PATHを通すため.profileもしくは.bashrcに以下を追加


.bashrc

#Anaconda

export PATH="$PYENV_ROOT/versions/anaconda3-5.0.1/bin:$PATH"

設定を反映させましょう


terminal

source ~/.bashrc

もしくは
source ~/.profile


3. PyTorch

公式HP (http://pytorch.org) のTopで自身の環境を指定すると、

それに応じたインストールコマンドが指定されます。私の場合は


terminal

conda install pytorch=0.3.0 cuda80 -c pytorch


間違えて異なるバージョンをインストールしてしまった場合は


terminal

conda uninstall pytorch torchvision cuda80 -c soumith


ついでに必要なパッケージもインストール


terminal

conda install -c conda-forge opencv python-graphviz



実際にESPNetを使ってみる

自分の好きなところにESPNetのリポジトリをクローン


terminal

git clone https://github.com/sacmehta/ESPNet.git


ディレクトリはこんな構造になってます

ESPNet/

├ pretrained/    cityscapseで学習済みの重み
│ ├ encoder/
│ └ decoder/
├ sample_video/
├ test/        test用のプログラム
│ ├ data/    テスト用のサンプル画像
│ └ results/ VisualizeResult.pyを実行すると作成され出力結果がここに入る
└ train/
└ city/    cityscapesで学習する際に使うテキストファイル
  ├ train.txt トレーニングに用いるRGB画像とラベル画像のパスが羅列されたテキスト
  └ val.txt バリーションに用いるRGB画像とラベル画像のパスが羅列されたテキスト


0. ちょっとテストしてみる


terminal

$ cd ESPNet/test

~/ESPNet/test$ python VisualizeResults.py --modelType 1 --p 2 --q 3

出力結果がtest/results以下に保存されていればオッケー


1. CityScapesで学習させる場合


1.1 train.txt/val.txtのパスを変更しなくて良い方法

Cityscapseのデータセットを下記のようにcity/以下に置く場合は1.2のパスを適応させる手順は必要ないです

ESPNet/

├ pretrained/
├ sample_video/
├ test/       
└ train/
└ city/
├ leftImg8bit/
│ └...
├ gtFine/
│ └...   
├ train.txt
└ val.txt


1.2 train.txt/val.txtのパスを変更する必要がある場合

容量不足等の理由で外部ストレージにデータセットを置く場合は

train.txtとval.txtを適切なパスを示すように変更する必要があります。

テキストエディタとかで置換してやってください。



変更前

/leftImg8bit/train/zurich/zurich_000022_000019_leftImg8bit.png,/gtFine/train/zurich/zurich_000022_000019_gtFine_labelTrainIds.png
変更後
/home/…/dataset/cityscapes/image/train/zurich/zurich_000022_000019_leftImg8bit.png,/home/…/dataset/cityscapes/label/gtFine/train/zurich/zurich_000022_000019_gtFine_labelIds.png

また、ラベル画像のファイル名が元々のtrain.txt/val.txtでは

〜gtFine_labelTrainIds.png だったのですが実際の画像ファイル名は

〜gtFine_labelIds.png    で「Train」があるかないかの微妙な違いがありました

これも自分の環境に合わせて置換してください

私の用にルートからのフルパスで置換した場合は


ESPNet/train/loadData.pyの58,59行目

img_file = ((self.data_dir).strip() + '/' + line_arr[0].strip()).strip()

label_file = ((self.data_dir).strip() + '/' + line_arr[1].strip()).strip()
から
img_file = (line_arr[0].strip()).strip()
label_file = (line_arr[1].strip()).strip()

に変更して保存してください。

そうでない場合は学習実行時に引数として

python main.py --data_dir /home/media/ssd

の用にパスを指定してうまく行くようにしてください

しかしcityscapes自体は30classであるはずが

label画像を見ると30以上の値が入っているため、これをどうにかしなければいけない

私が作成したプログラムがあるのでパスやlabelのマップを変更してもらえれば使えると思います


2. 自分のデータセットで学習させる場合


2.1 RGB画像とラベル画像の用意

とりあえずどこでも良いのでファイル名が対応したRGB画像とラベル画像をこんな感じで用意してください

dataset/

├ train/
│ ├ rgb/
│ │ ├ 000000.png
│ │ ├ 000001.png
│ │
│ └ label/
│ ├ 000000.png
│ ├ 000001.png
└ val/
├ rgb/
│ ├ 000005.png
│ ├ 000009.png

└ label/
├ 000005.png
├ 000009.png


2.2 train.txtとval.txtの作成

ESPNet/train/にcreate_dataset_txt.pyを保存してください。

実行する前にtrain_rgb_dirなどは各自のデータセットを置いた場所に書き換えてください

train_txtとval_txtはESPNet/train/jisakuのようにESPNetがある場所したほうがわかりやすいと思います。


create_dataset_txt.py

# -*- coding: utf-8 -*-

import os
import glob

def main():
Create_txt(train_rgb_dir, train_label_dir, train_txt)
Create_txt(val_rgb_dir, val_label_dir, val_txt)

def Create_txt(rgb_dir, label_dir, txt_path):
filelist_rgb = glob.glob(os.path.join(rgb_dir, '*.png'))
txt = os.path.basename(txt_path)
try:
with open(txt_path, mode='x') as f:
for file_path_rgb in filelist_rgb:
filename = os.path.basename(file_path_rgb)
file_path_label = os.path.join(label_dir, filename)
s = file_path_rgb + ',' + file_path_label + '\n'
f.write(s)

print ('create :', txt)

except FileExistsError:
print(txt, 'already exists')

if __name__ == '__main__':
#書き換えが必要です
train_rgb_dir = '/path/to/dataset/train/rgb'
train_label_dir = '/path/to/dataset/train/label'
val_rgb_dir = '/path/to/dataset/val/rgb'
val_label_dir = '/path/to/dataset/val/label'

#jisakuの部分は好きな名前で構いません
#※プログラムを実行する前にディレクトリを作成しておいてください
train_txt = '/path/to/ESPNet/train/jisaku/train.txt'
val_txt = '/path/to/ESPNettrain/jisaku/val.txt'

main()


実行します


terminal

$ cd /ESPNet/train

~/ESPNet/train$ python create_dataset_txt.py

恐らく、jisakuディレクトリにtrain.txtとval.txtができていると思うので確認してください

ESPNet/

├ pretrained/
├ sample_video/
├ test/       
└ train/
├ city/
└ jisaku/
├ train.txt
└ val.txt


2.3 ラベル画像のリマッピング※必要な人だけ

データセットに10クラスしか含まれていない場合でも

ラベル画像には20や30といったクラス数より大きい値が振られていることがある。

これに気づかずにクラス数を10に指定して学習させようとするよ

クラス数より大きい値がラベル画像中に出てきてしましまうため、エラーが起こる

これを解消するためにラベル画像を0から9(10クラスの場合)の値に収めるようにリマッピングする処理をする必要がある。

この先、追記予定

PS

別にリマッピングしなくてもラベル画像中に出てくる最大値をクラス数に指定すれば問題ないですが

私はなんか気持ち悪いのでラベル画像をリマッピングしています。


3. 学習スタート

学習は2段階あります。

①エンコーダーの学習

②デコーダーの学習


3.1 エンコーダーの学習


terminal

$ cd ESPNet/train

#Cityscapesでやる場合
ESPNet/train$ python main.py --scaleIn 8 --p 2 --q 8
#自作データセットの場合(--classesにはデータセットのクラス数を指定)
ESPNet/train$ python main.py --data_dir ./jisaku --classes 12 --scaleIn 8 --p 2 --q 8




[247/387] loss: 0.675 time:0.48
[248/387] loss: 0.651 time:0.48
[249/387] loss: 0.701 time:0.48
[250/387] loss: 0.536 time:0.48
[251/387] loss: 1.010 time:0.50
[252/387] loss: 0.862 time:0.48
[253/387] loss: 0.860 time:0.48
[254/387] loss: 0.758 time:0.48
[255/387] loss: 0.860 time:0.48
[256/387] loss: 0.794 time:0.48
[257/387] loss: 0.671 time:0.48
[258/387] loss: 0.738 time:0.48


重みのスナップショットは

ESPNet/train/results_enc__enc_2_8/

に保存されています。※2と8は学習時に指定したパラメータの値

収束してきたら、学習を止めて次はデコーダーの学習に移りましょう。


3.2 デコーダーの学習

エンコーダーの学習には引数は指定していませんでしたが、デフォルトでbatch_size=12を用いていました。

デコーダーではbatch_size=6なので指定してやりましょう。

環境によってはもっと小さい値をしてする必要もあるので、out of memoryと表示された場合は

コマンドラインで--batch_sizeを適切な値にしてやってください。


terminal

#Cityscapesでやる場合

ESPNet/train$ python main.py --batch_size 6 --scaleIn 1 --p 2 --q 8 --decoder True --pretrained ../pretrained/encoder/espnet_p_2_q_8.pth
#自作データセットの場合(--classesにはデータセットのクラス数を指定)
ESPNet/train$ python main.py --classes 12 --data_dir ./izunuma --batch_size 6 --scaleIn 1 --p 2 --q 8 --decoder True --pretrained ./results_enc__enc_2_8/model_81.pth

Encoder loaded!
Total network parameters: 354517
Data statistics
[131.36485 144.77011 134.60292] [76.74333 68.792656 71.73109 ]
[ 9.702483 9.986656 7.9102645 8.283521 3.7311087 10.492059
10.206483 4.453087 10.42091 10.340107 10.338798 1.9798005]
Learning rate: 0.0005
[0/773] loss: 2.491 time:2.24
[1/773] loss: 2.485 time:0.40
[2/773] loss: 2.478 time:0.40
[3/773] loss: 2.478 time:0.40
[4/773] loss: 2.467 time:0.40
[5/773] loss: 2.465 time:0.40



収束したらテストしてみてください。


参考サイト

https://github.com/sacmehta/ESPNet

Ubuntu 16.10にPyTorchを導入してみた

Tensorflow Lite (TensorflowLite) / Tensorflow+RaspberryPi+Python で超軽量 "Semantic Segmentation" モデル "UNet" "ENet" を実装する軽量モデルその2