Help us understand the problem. What is going on with this article?

誰でも出来る!DeepLab v3+でGPUを使って自作データセットで学習・推論する

はじめに

DeepLab v3+はセマンティックセグメンテーションのための最先端のモデルです。
この記事では、DeepLab v3+のgithubを使って、公開されたデータセットまたは自分で用意したデータセットで学習・推論までをおこなう方法を紹介します。

DeepLabv3+をやってみたという記事は検索すれば多く見つかります。
しかし、どれもこれも、結局どうやってデータ準備したの?とか、手順を省略しすぎて分からん!というものが多かった・・・。
今回の記事では、全ての手順を丁寧に説明しながら誰でもできるように紹介していきたいと思います。
※誰でもと書きましたが、tensorflowとGPUが絡む部分は各自で気を付けて頂かなければ失敗します。

環境構築

OSは Windows10 が対象です。
とはいえLinuxでもコマンドの書き方などを変えれば動くと思います。

学習にはGPUを使います。
tensorflowでGPUを使用するときに気をつけなければならないことは、
python, CUDA, cuDNN, tensorflowのバージョンをそれぞれ対応するものに合わせなければならないことです。
これらのバージョンは思った以上にセンシティブなので、ちょっとくらい違うバージョンでもいいだろ、と思ってると面倒なことになります。
Anaconda, CUDA, cuDNN, tensorflowのインストールは検索すれば多くの記事があるので、ここでは割愛します。

tensorflow+GPUで成功した環境

項目 バージョン等
OS Windows10
GPU RTX2060
GPUのドライバー 436.48
python Python 3.6.9 :: Anaconda, Inc.
CUDA CUDA Toolkit 9.0
cuDNN v7.4.1.5
tensorflow_gpu v1.13.1
keras 2.3.1

似たような記事を色々調べた結果、
- python=3.6
- CUDA=9.0
- cuDNN=7.4
- tensorflow=1.13

が一番安定してる?気がします (2019年10月20日時点)。

学習データセットの準備

今回は以下の2通りの方法を紹介します。

  • (A) Pascal VOCデータセットで学習する場合
  • (B) 自作データセットで学習する場合

Pascal VOCデータセットは一般に公開されているsemantic segmentationの有名なデータセットです。

(A) Pascal VOCデータセットで学習する場合

Pascal VOCデータセットをこちらからダウンロードする。
tar形式であるため、7-Zip等をインストールして解凍してください。

展開すると以下のようなディレクトリ構成となっています。

VOCdevkit
  └ VOC2012
      ├ ImageSets
      ├ JPEGImages
      ├ SegmentationClass
      └ SegmentationObject
  • ImageSets: 学習/検証用データのリストが格納されている
  • JPEGImages: 学習データがjpg形式で格納されている
  • SegmentationClass: インデックスカラーでsemantic segmentationされたデータ
  • SegmentationObject: インデックスカラーでinstance segmentationされたデータ(今回は使いません)

完了。

(B) 自作データセットで学習する場合

画像データの収集

まずは元となる画像を収集します。
例えば、犬や猫の画像をネット上から収集してきましょう。
ディープラーニングにおいて、データセット収集・設計が最も大変で最も大切であると言っても過言ではありません。
そこで、今回はPascalVOCのデータセットに準拠してデータを揃えます。

  • 画像形式はjpgであること
    =>画像形式をjpgに変換すれば問題ありません
  • 画像の大きさは横幅最大500, 縦幅最大500で揃えること
    例:(500x375), (450x500), ...
    =>画像をリサイズすれば問題ありません

pngをjpgに変換するには以下を使用してください。

img2jpg.py
# coding: utf-8

import cv2
import os
import glob
import numpy as np


def MyMkdir(dr, disp=True):
    if not os.path.exists(dr):
        os.makedirs(dr)
        if (disp):
            print("mkdir >> {0}".format(dr))


def get_filepath_param(filepath):
    basename = os.path.basename(filepath)
    dr = os.path.dirname(filepath)
    if (dr == ""):
        dr = "."
    name, ext = os.path.splitext(basename)
    ext = ext.replace(".", "")
    return basename, dr, name, ext


def get_files_path_from_dir(path, sort=True):
    files = []
    for file in glob.glob(path):
        file = file.replace('\\', '/')
        files.extend([file])
    if sort:
        files.sort()
    return files


def main():
    # 説明
    # in_dir内の"png"を全てout_dirに"jpg"形式で出力する

    in_dir = "images"
    in_ext = "png"
    out_dir = "out"
    out_ext = "jpg"

    MyMkdir(out_dir)

    files = get_files_path_from_dir(in_dir + "/*." + in_ext)

    for f in files:
        basename, dr, name, ext = get_filepath_param(f)
        out_path = "{dir}/{name}.{ext}".format(dir=out_dir, name=name, ext=out_ext)

        img = cv2.imread(f)
        cv2.imwrite(out_path, img)
        print("output >> {0}".format(out_path))


if __name__ == "__main__":
    main()

画像の大きさを横幅最大500, 縦幅最大500で揃えるには以下を使用してください。

max_resize.py
# coding: utf-8

import cv2
import os
import glob
import numpy as np


def MyMkdir(dr, disp=True):
    if not os.path.exists(dr):
        os.makedirs(dr)
        if (disp):
            print("mkdir >> {0}".format(dr))


def get_filepath_param(filepath):
    basename = os.path.basename(filepath)
    dr = os.path.dirname(filepath)
    if (dr == ""):
        dr = "."
    name, ext = os.path.splitext(basename)
    ext = ext.replace(".", "")
    return basename, dr, name, ext


def get_files_path_from_dir(path, sort=True):
    files = []
    for file in glob.glob(path):
        file = file.replace('\\', '/')
        files.extend([file])
    if sort:
        files.sort()
    return files


def scale_to_width(img, width):
    scale = width / img.shape[1]
    return cv2.resize(img, dsize=None, fx=scale, fy=scale)


def scale_to_height(img, height):
    scale = height / img.shape[0]
    return cv2.resize(img, dsize=None, fx=scale, fy=scale)


def main():
    # 説明
    # in_dir内の全ての"jpg"に対して、
    # 横長の画像であれば(500x???), 縦長の画像であれば(???x500)にリサイズしてout_dirに出力する

    in_dir = "images"
    in_ext = "jpg"
    out_dir = "out"
    out_ext = "jpg"

    max_height = 500
    max_width = 500

    MyMkdir(out_dir)

    files = get_files_path_from_dir(in_dir + "/*." + in_ext)

    for f in files:
        basename, dr, name, ext = get_filepath_param(f)
        out_path = "{dir}/{name}.{ext}".format(dir=out_dir, name=name, ext=out_ext)

        img = cv2.imread(f)

        if img.shape[0] > img.shape[1]:
            # 縦長の画像の場合
            img = scale_to_height(img, max_height)
        else:
            # 横長の画像の場合
            img = scale_to_width(img, max_width)

        cv2.imwrite(out_path, img)
        print("output >> {0} => ({1}, {2})".format(out_path, img.shape[1], img.shape[0]))


if __name__ == "__main__":
    main()

labelmeの概要

今回はアノテーションツール「labelme」を使ってアノテーション情報を作成します。
githubはこちら
labelmeはセマンティックセグメンテーションのアノテーションツールです。
labelmeを使ってデータをアノテーションし、labelmeの形式変換機能を使ってPascalVOC形式に変換します。
annotation.jpg

labelmeのインストール

Windows10でAnaconda環境の方は以下を実行してください。

conda install pillow=4.0.0
pip install labelme

それ以外の方はgithubのインストール方法に従ってください。
あとは、コマンドプロンプトで以下を実行すればlabelmeが起動します。

labelme.exe

データのアノテーション

labelmeの使い方はこちらを参考にしてください。
画像データのアノテーションを行うと、例えば以下のようなdatasetディレクトリ構成になります。

dataset/
 ├ 001.jpg
 ├ 001.json
 ├ 002.jpg
 ├ 002.json
 ├ 003.jpg
 └ 003.json

jpg画像に対してアノテーション情報jsonが付与されています。

次に、ラベル情報を書いたファイルlabels.txtを作成します。
例えば、[dog、cat、hourse]の3種類をsegmentationする場合、以下のようなファイルになります。

labels.txt
__ignore__
_background_
dog
cat
hourse

以下の2つは必ず入れてください。

  • __ignore__: 無視ラベル
  • _background_: 背景ラベル

PascalVOCデータセットの形式に変換

git cloneでlabelmeのソースを入手します。

git clone https://github.com/wkentaro/labelme

次に、labelme\examples\semantic_segmentation\labelme2voc.pyを使ってアノテーション済データセットをPascalVOC形式に変換します。
以下を実行する。

python labelme2voc.py "dataset" "VOC2012" --labels labels.txt

labelme2voc.py [input_dir] [output_dir] --labels [labels]

  • input_dir: labelmeでアノテーションした画像とjsonセットのディレクトリへのパス
  • output_dir: 出力ディレクトリへのパス
  • labels: ラベル情報テキストへのパス

実行後、以下のような以下のようなディレクトリが作成される。

VOC2012/
 ├ JPEGImages
 ├ SegmentationClass
 ├ SegmentationClassPNG
 ├ SegmentationClassVisualization
 └ class_names.txt

実行後のディレクトリ名はPascalVOCとは異なるため、修正する。
以下のようにフォルダ名変更する。

SegmentationClass => SegmentationClass_npy(使わないので適当な名前でよい)
SegmentationClassPNG => SegmentationClass

学習/検証データリストの作成

VOC2012ディレクトリ直下に以下のようにディレクトリ・ファイル作成

mkdir VOC2012/ImageSets/Segmentation

VOC2012/ImageSets/Segmentation/train.txt
VOC2012/ImageSets/Segmentation/trainval.txt
VOC2012/ImageSets/Segmentation/val.txt

ここで、学習データ数、検証データ数を記述します。
例えば、001.jpg, 002.jpg, 003.jpgがあった時、
学習データ=[001.jpg, 002.jpg]
検証データ=[003.jpg]
とする場合は以下のように記述する。

train.txtには学習データのリストを記述する。
※ファイル名のみ記述

train.txt
001
002

val.txtには検証データのリストを記述する。

val.txt
003

trainval.txtには学習+検証データのリストを記述する。

trainval.txt
001
002
003

最終的に、以下のようなディレクトリ構成になり、自作データセットの準備は完了。

VOC2012/
├─class_names.txt
│
├─ImageSets
│  └─Segmentation
│          train.txt
│          trainval.txt
│          val.txt
│
├─JPEGImages
│      001.jpg
│      002.jpg
│      003.jpg
│
├─SegmentationClass
│      001.png
│      002.png
│      003.png
│
├─SegmentationClassVisualization
│      001.jpg
│      002.jpg
│      003.jpg
│
└─SegmentationClass_npy
        001.npy
        002.npy
        003.npy

Pascal VOC形式のデータセットをTFRecordに変換

これまで、以下の2通りのデータセットの準備方法を説明しましたが、

  • (A) Pascal VOCデータセットで学習する場合
  • (B) 自作データセットで学習する場合

どちらの方法でも最終的にVOC2012ディレクトリが作成されたと思います。
以下では、このVOC2012TFRecordに変換します。

deeplabのソース入手とディレクトリ構築

以下を実行

git clone https://github.com/tensorflow/models
cd models/research/deeplab/datasets
mkdir pascal_voc_seg/VOCdevkit

pascal_voc_seg/VOCdevkitの下にVOC2012を配置

pascal_voc_seg/
    └ VOCdevkit/
        └ VOC2012/

SegmentationClassRawを生成

以下を実行

python ./remove_gt_colormap.py --original_gt_folder="./pascal_voc_seg/VOCdevkit/VOC2012/SegmentationClass" --output_dir="./pascal_voc_seg/VOCdevkit/VOC2012/SegmentationClassRaw"

実行後、特に表示はありませんが「SegmentationClassRaw」ディレクトリが作成されます。

余談
「SegmentationClass」にある画像は、インデックスカラーでセグメンテーションされています。
「SegmentationClassRaw」にある画像は「SegmentationClass」をグレースケール画像(RGBそれぞれに同じインデックス値を設定)に変換したものです。
(B) 自作データセットで学習する場合では「SegmentationClassRaw」は真っ黒な画像が生成されるかもしれませんが、よく見るとセグメンテーション部分が見えます。
SegmentationClassRawにbackground=0, dog=1, cat=2, ...のようにグレイスケール値が割当てられると思います。

TFRecordを生成

以下を実行

mkdir pascal_voc_seg/tfrecord
python ./build_voc2012_data.py ^
    --image_folder="./pascal_voc_seg/VOCdevkit/VOC2012/JPEGImages" ^
    --semantic_segmentation_folder="./pascal_voc_seg/VOCdevkit/VOC2012/SegmentationClassRaw" ^
    --list_folder="./pascal_voc_seg/VOCdevkit/VOC2012/ImageSets/Segmentation" ^
    --image_format="jpg" ^
    --output_dir="./pascal_voc_seg/tfrecord/"

※Linuxの場合は^\に変換
実行後、tfrecordディレクトリにTFRecordデータセットが出力される。

学習部

ディレクトリ構築

以下を実行してdeeplab公式が推奨しているディレクトリ構築をおこないます。

cd models/research/deeplab
mkdir ./datasets/pascal_voc_seg/init_models

※以下ではこのディレクトリmodels/research/deeplabをカレントディレクトリとして説明しています。

事前学習済モデルの入手

いわゆる転移学習(fine tuning)というやつです。
既に学習済みのモデルを使うことで、少ないデータ数で短時間でも高い精度を得られます。
事前学習済みモデルデータを以下からダウンロード
http://download.tensorflow.org/models/deeplabv3_pascal_train_aug_2018_01_04.tar.gz
解凍して出来たdeeplabv3_pascal_train_augを以下に置く

models/research/deeplab/datasets/pascal_voc_seg/init_models/deeplabv3_pascal_train_aug/

学習/検証データ設定の変更

注意(2019/10/22) segmentation_dataset.pyは非推奨になりました

models/research/deeplab/deprecated/segmentation_dataset.pyを開く
多くの関連記事で、segmentation_dataset.pyを編集するように紹介されていますが、これは非推奨になっています。
最新版のdeeplabv3ではそもそもsegmentation_dataset.pyがimportされていないので、値を書き換えても何の意味もないです。

正しくは、
models/research/deeplab/datasets/data_generator.py
を開きます。

data_generator.pyには以下のようにPascalVOCデータセットの設定が記述されている。

data_generator.py
_PASCAL_VOC_SEG_INFORMATION = DatasetDescriptor(
    splits_to_sizes={
        'train': 1464,
        'train_aug': 10582,
        'trainval': 2913,
        'val': 1449,
    },
    num_classes=21,
    ignore_label=255,
)

(A) Pascal VOCデータセットで学習する場合では、特に書き換えは必要ない。
(B) 自作データセットで学習する場合では、下記を参考にtrain, trainval, valを書き換える必要がある。

項目 内容
'train' 学習データの数
(VOC2012/ImageSets/Segmentation/train.txtの数)
'trainval' 学習+検証データの数
(VOC2012/ImageSets/Segmentation/trainval.txtの数)
'val' 検証データの数
(VOC2012/ImageSets/Segmentation/val.txtの数)
'train_aug' image augumentation用データの数、論文ではこれを使ってたのでその名残?無視して良い
num_classes 分類クラス数
注意:background(1) + 分類クラス数になります。
しかし、今回は転移学習を使うので、(B) 自作データセットで学習する場合であっても21クラスに固定してください
ignore_label 無視ラベルの画素値
この画素値のデータは学習時に一切考慮されなくなります。
(A) Pascal VOCデータセットで学習する場合ではSegmentationClassRawに白枠線が見えますが、255というのはこの白枠を指しています。
(B) 自作データセットで学習する場合では白値はありませんが、255で問題ありません。

結果として、(B) 自作データセットで学習する場合で[dog, cat, hourse]の3クラスのsegmentationを行う場合は、例えば以下のように記述します。
注意:本来は[background, dog, cat, hourse]の4クラス分類になりますが、転移学習を使うので21クラスのままでOKです
元のモデルの分類のindex=1のaeroplaneをdogで上書き学習、index=2のbicycleをcatで上書き学習するというイメージです。

data_generator.py
_PASCAL_VOC_SEG_INFORMATION = DatasetDescriptor(
    splits_to_sizes={
        'train': 200,
        'trainval': 300,
        'val': 100,
    },
    num_classes=21, # fix 21 class
    ignore_label=255,
)

パスを通す

パスの追加

「deeplab」と「slim」のモジュールを使用するため、PATHを設定する必要がある。
[システムの詳細設定]->[環境変数]->[システム環境変数]の[PYTHONPATH]に以下を追加する。

C:\Users\xxxx\Desktop\models\research
C:\Users\xxxx\Desktop\models\research\slim
  • PYTHONPATHが無ければ新しく作る
  • フルパスで追加

設定後、PCを再起動すると適用される。

デバッグ

ここで以下のエラー発生

ModuleNotFoundError: No module named 'nets'

models\research\slimをパスに追加していれば問題ないと思ったのだが
仕方がないので、
models\research\slim\nets

models\research\deeplab
の直下にコピーする

PYTHONPATHの確認

以下を実行すると、正しくPATHの設定が出来ているかどうかを確認できる。

python model_test.py -v

実行後、以下のように表示されれば、問題ない。

Ran 5 tests in 12.709s

OK (skipped=1)

学習部

学習の実行

以下を実行

python train.py ^
    --logtostderr ^
    --training_number_of_steps=300 ^
    --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="./datasets/pascal_voc_seg/init_models/deeplabv3_pascal_train_aug/model.ckpt" ^
    --train_logdir="./datasets/pascal_voc_seg/exp/train_on_trainval_set/train" ^
    --dataset_dir="./datasets/pascal_voc_seg/tfrecord" ^
    --fine_tune_batch_norm=false ^
    --initialize_last_layer=true ^
    --last_layers_contain_logits_only=false

※Linux等の場合は^\に変換

引数名 意味 値(例)
--logtostderr ログを標準エラー出力へ出力
--training_number_of_steps 学習回数 300
※お試し程度で300にしていますが、実際はもっと増やしてください
--train_split 使用データ [train, val, trailval]
--model_variant 識別モデル種類 [xception_65, mobilenet_v2]
--atrous_rates Atrous畳み込みの比率
※複数回設定可能
6, 12, 18
--output_stride 出力ストライド(atrous_rateとの組み合わせ) 16
--decoder_output_stride 入出力の空間解像度の比率 4
--train_crop_size 画像の切り出しサイズ "513,513" (width, height)
※「train_crop_size」は「output_stride*k+1 (k>=1)」の値を指定してください。
--train_batch_size ミニバッチサイズ 1
--dataset データセット名 [cityscapes, pascal_voc_seg, ade20k]
--tf_initial_checkpoint 学習済みモデルへのパス "./datasets/pascal_voc_seg/init_models/deeplabv3_pascal_train_aug/model.ckpt"
※ファイルへのパスではないので、該当ファイルが無くても問題ない
--train_logdir 学習ログ出力フォルダへのパス "./datasets/pascal_voc_seg/exp/train_on_trainval_set/train"
--dataset_dir tfrecordデータセットフォルダへのパス "./datasets/pascal_voc_seg/tfrecord"
--fine_tune_batch_norm Batch Normalizationの実行 [true, false]
※GPUで学習するときはfalse
--initialize_last_layer 最後のレイヤーの初期化 true

実際にはもっと多くの引数があります。
train.pyのコードを参照してください。

デバッグ

GPUメモリエラー

以下のようなエラーが出た場合、GPUメモリが足りていません。

ResourceExhaustedError (see above for traceback): OOM when allocating tensor wit
h shape[1,64,257,257] and type float on /job:localhost/replica:0/task:0/device:G
PU:0 by allocator GPU_0_bfc

train_batch_sizetrain_crop_sizeを小さくしてみてください。

学習ログの出力

以下に設定したディレクトリに学習結果ログが出力される。
--train_logdir="./datasets/pascal_voc_seg/exp/train_on_trainval_set/train"
TensorFlow saveメソッドは以下の4種類のファイルを保存します。
- checkpoint
- model.ckpt-16272.data-00000-of-00001
- model.ckpt-16272.index
- model.ckpt-16272.meta

ファイル 意味
checkpoint 学習の途中記録
.data-????-of-???? モデルの重みの断片(複数)
.index どの重みがどの断片に保存されているかを示すインデックスファイル
.meta 保存されたグラフ構造が記述されている

tensorboardで学習進捗の確認

以下を実行してtensorboardの起動

pip install tensorboard
tensorboard --logdir="./datasets/pascal_voc_seg/exp/train_on_trainval_set"

ブラウザでhttp://localhost:6006にアクセス

推論部

推論結果の可視化

以下を実行

python vis.py ^
    --logtostderr ^
    --vis_split="val" ^
    --model_variant="xception_65" ^
    --atrous_rates=6 ^
    --atrous_rates=12 ^
    --atrous_rates=18 ^
    --output_stride=16 ^
    --decoder_output_stride=4 ^
    --vis_crop_size="513,513" ^
    --dataset="pascal_voc_seg" ^
    --checkpoint_dir="./datasets/pascal_voc_seg/exp/train_on_trainval_set/train" ^
    --vis_logdir="./datasets/pascal_voc_seg/exp/train_on_trainval_set/vis" ^
    --dataset_dir="./datasets/pascal_voc_seg/tfrecord" ^
    --max_number_of_iterations=1 ^
    --eval_interval_secs=0
引数名 意味 値(例)
--logtostderr ログを標準エラー出力へ出力
--vis_split 使用データ [train, val, trailval]
--model_variant 識別モデル種類 [xception_65, mobilenet_v2]
※学習時と同じ設定にする必要あり
--atrous_rates Atrous畳み込みの比率
※複数回設定可能
6, 12, 18
※学習時と同じ設定にする必要あり
--output_stride 出力ストライド(atrous_rateとの組み合わせ) 16
※学習時と同じ設定にする必要あり
--decoder_output_stride 入出力の空間解像度の比率 4
※学習時と同じ設定にする必要あり
--vis_crop_size 画像の切り出しサイズ "513,513" (width, height)
※「train_crop_size」は「output_stride*k+1 (k>=1)」の値を指定してください。
※この設定より大きい画像を入力とするとエラーが出ます。
--dataset データセット名 [cityscapes, pascal_voc_seg, ade20k]
--checkpoint_dir 学習済みモデルへのパス "./datasets/pascal_voc_seg/exp/train_on_trainval_set/train"
--vis_logdir 学習ログ出力フォルダへのパス "./datasets/pascal_voc_seg/exp/train_on_trainval_set/vis"
--dataset_dir tfrecordデータセットフォルダへのパス "./datasets/pascal_voc_seg/tfrecord"
--max_number_of_iterations 最大繰り返し数
--eval_interval_secs 繰り返し待ち時間(秒)

実行結果

--vis_logdirに推論結果の画像が出力されます。
元画像が0001.jpgだとすると、セグメンテーションされた画像が0001_prediction.jpgとして出力されます。

デバッグ

以下のようなエラーが出た場合、

tensorflow.python.framework.errors_impl.InvalidArgumentError: padded_shape[1]=160 is not divisible by block_shape[1]=18
         [[node aspp3_depthwise/depthwise/SpaceToBatchND (defined at C:\Users\xxxx\Desktop\models\research\deeplab\core\utils.py:92) ]]
         [[node ArgMax (defined at C:\Users\xxxx\Desktop\models\research\deeplab\model.py:197) ]]

原因は入力画像の大きさにあると考えられます。
--vis_crop_sizeを最大入力画像サイズより大きく設定しましょう。

モデルのエクスポート

学習済みのモデルをProtocol Buffer(.pbファイル)で出力します。
.pbファイルはAndroid, iOS等で使う場合に必要になるようです。

以下を実行します。

python export_model.py ^
    --checkpoint_path="./datasets/pascal_voc_seg/exp/train_on_trainval_set/train/model.ckpt-300" ^
    --export_path="./datasets/pascal_voc_seg/exp/train_on_trainval_set/export/frozen_inference_graph.pb" ^
    --num_classes=21 ^
    --model_variant="xception_65" ^
    --atrous_rates=6 ^
    --atrous_rates=12 ^
    --atrous_rates=18 ^
    --output_stride=16 ^
    --decoder_output_stride=4

--checkpoint_pathは任意の学習回数のモデルを指定してください。
--num_classes=21について、学習時のクラス数とは関係なしに「21」に固定しておいてください。
理由としては、事前学習モデルとして学習部において--tf_initial_checkpoint="./datasets/pascal_voc_seg/init_models/deeplabv3_pascal_train_aug/model.ckpt"を使っており、このモデルは21クラス分類であるためです。
学習時のクラス数と初期モデルのクラス数が違っても問題ないので、安心してください。

実行後、--export_pathで指定したフォルダにfrozen_inference_graph.pbが出力されると思います。

エクスポートしたモデルで推論

概要

エクスポートしたfrozen_inference_graph.pbを使って推論を行います。
ここでは、使いやすいように画像とモデルのパスを指定して推論を行うプログラムを紹介します。
基本的には下記を参照して作っています。
https://github.com/tensorflow/models/blob/master/research/deeplab/deeplab_demo.ipynb

画像とモデルを指定して推論と結果表示を行う

以下のプログラムを使用してください。
img_pathに推論する画像を、model_pathにfrozen_inference_graph.pbを指定します。

prediction.py
# coding: utf-8

import os
from io import BytesIO
import tarfile
import tempfile
from six.moves import urllib

from matplotlib import gridspec
from matplotlib import pyplot as plt
import numpy as np
from PIL import Image

import tensorflow as tf


class DeepLabModel(object):
    INPUT_TENSOR_NAME = 'ImageTensor:0'
    OUTPUT_TENSOR_NAME = 'SemanticPredictions:0'
    INPUT_SIZE = 513
    FROZEN_GRAPH_NAME = 'frozen_inference_graph'

    def __init__(self, frozen_path):
        self.graph = tf.Graph()
        graph_def = None
        with open(frozen_path, 'rb') as f:
            graph_def = tf.GraphDef()
            graph_def.ParseFromString(f.read())
        if graph_def is None:
            raise RuntimeError('Cannot find inference graph in tar archive.')
        with self.graph.as_default():
            tf.import_graph_def(graph_def, name='')
        self.sess = tf.Session(graph=self.graph)

    def run(self, image):
        width, height = image.size
        resize_ratio = 1.0 * self.INPUT_SIZE / max(width, height)
        target_size = (int(resize_ratio * width), int(resize_ratio * height))
        resized_image = image.convert('RGB').resize(target_size, Image.ANTIALIAS)
        batch_seg_map = self.sess.run(
            self.OUTPUT_TENSOR_NAME,
            feed_dict={self.INPUT_TENSOR_NAME: [np.asarray(resized_image)]})
        seg_map = batch_seg_map[0]
        return resized_image, seg_map

    def returnSize(self,image):
        width, height = image.size
        resize_ratio = 1.0 * self.INPUT_SIZE / max(width, height)
        target_size = (int(resize_ratio * width), int(resize_ratio * height))
        return target_size


def create_pascal_label_colormap():
    """Creates a label colormap used in PASCAL VOC segmentation benchmark.

    Returns:
      A Colormap for visualizing segmentation results.
    """
    colormap = np.zeros((256, 3), dtype=int)
    ind = np.arange(256, dtype=int)

    for shift in reversed(range(8)):
        for channel in range(3):
            colormap[:, channel] |= ((ind >> channel) & 1) << shift
        ind >>= 3

    return colormap


def label_to_color_image(label):
    """Adds color defined by the dataset colormap to the label.

    Args:
      label: A 2D array with integer type, storing the segmentation label.

    Returns:
      result: A 2D array with floating type. The element of the array
        is the color indexed by the corresponding element in the input label
        to the PASCAL color map.

    Raises:
      ValueError: If label is not of rank 2 or its value is larger than color
        map maximum entry.
    """
    if label.ndim != 2:
        raise ValueError('Expect 2-D input label')

    colormap = create_pascal_label_colormap()

    if np.max(label) >= len(colormap):
        raise ValueError('label value too large.')

    return colormap[label]


def vis_segmentation(image, seg_map):
    """Visualizes input image, segmentation map and overlay view."""
    plt.figure(figsize=(15, 5))
    grid_spec = gridspec.GridSpec(1, 4, width_ratios=[6, 6, 6, 1])

    plt.subplot(grid_spec[0])
    plt.imshow(image)
    plt.axis('off')
    plt.title('input image')

    plt.subplot(grid_spec[1])
    seg_image = label_to_color_image(seg_map).astype(np.uint8)
    plt.imshow(seg_image)
    plt.axis('off')
    plt.title('segmentation map')

    plt.subplot(grid_spec[2])
    plt.imshow(image)
    plt.imshow(seg_image, alpha=0.7)
    plt.axis('off')
    plt.title('segmentation overlay')

    unique_labels = np.unique(seg_map)
    ax = plt.subplot(grid_spec[3])
    plt.imshow(
        FULL_COLOR_MAP[unique_labels].astype(np.uint8), interpolation='nearest')
    ax.yaxis.tick_right()
    plt.yticks(range(len(unique_labels)), LABEL_NAMES[unique_labels])
    plt.xticks([], [])
    ax.tick_params(width=0.0)
    plt.grid('off')
    plt.show()


# label setting
LABEL_NAMES = np.asarray([
    'background', 'dog', 'cat', 'hourse'
])

FULL_LABEL_MAP = np.arange(len(LABEL_NAMES)).reshape(len(LABEL_NAMES), 1)
FULL_COLOR_MAP = label_to_color_image(FULL_LABEL_MAP)


def main():
    img_path = "img_001.jpg"
    model_path = "./datasets/pascal_voc_seg/exp/train_on_trainval_set/export/frozen_inference_graph.pb"

    # load model
    model = DeepLabModel(model_path)

    # read image
    original_im = Image.open(img_path)

    # inferences DeepLab model
    resized_im, seg_map = model.run(original_im)

    # show inference result
    vis_segmentation(resized_im, seg_map)


if __name__ == '__main__':
    main()

参考資料

まとめ

いかがでしたでしょうか。
誰でも出来るとは書きましたが、どこか1つでも間違ったり、GPUでエラーが出たり、日が経てばバージョン変更で仕様が変わったりするとやはりうまくいかない可能性もあります。
しかし、出来るだけ分かりやすく?書いたつもりなので、誰かの助けになると幸いです。

Why not register and get more from Qiita?
  1. We will deliver articles that match you
    By following users and tags, you can catch up information on technical fields that you are interested in as a whole
  2. you can read useful information later efficiently
    By "stocking" the articles you like, you can search right away
Comments
Sign up for free and join this conversation.
If you already have a Qiita account
Why do not you register as a user and use Qiita more conveniently?
You need to log in to use this function. Qiita can be used more conveniently after logging in.
You seem to be reading articles frequently this month. Qiita can be used more conveniently after logging in.
  1. We will deliver articles that match you
    By following users and tags, you can catch up information on technical fields that you are interested in as a whole
  2. you can read useful information later efficiently
    By "stocking" the articles you like, you can search right away