DeepLab v3+でオリジナルデータを学習してセマンティックセグメンテーションする

#やりたいこと
オリジナルデータを学習させてDeepLab v3+で「人物」と「テニスラケット」をセメンティックセグメンテーションできるようにします。DeepLab v3+でのオリジナルデータの学習はやり方が特殊で、調べながらやるのに苦労しました。ですので、自分への備忘録も兼ねて記事を作成しました。

#なぜ「人物」と「テニスラケット」をセマンティックセグメンテーションするのか
↓のような感じでテニスのサーブフォームで残像モーション動画を生成したかったからです。

サーブフォーム分析用の残像動画は、セマンティックセグメンテーション（深層学習技術）で背景と人間を分離したマスク画像を生成し、フレーム毎に重ね合わせて作成。動画をクラウドに投げ込んで、自動生成するようにしました。AWS LambdaとS3でシステム構築。クラウド破産するので公開はしてませんが。 pic.twitter.com/gZFI3EYQxf

— おたこ (@otakoma) August 3, 2019

セマンティックセグメンテーションで、人物とラケットとして抜き出した画像を前の画像に重ね合わせていくことで残像モーション動画を作成します。

残像モーション動画生成システムについてはこちら
【テニスのフォーム解析】残像モーション動画生成システム

#環境
Windows10 GPU1080ti
Tensorflow 1.14.0
DeepLab v3+

#DeepLab v3+のインストールはこちらから
DeepLab v3+のgithubページ
https://github.com/tensorflow/models/tree/master/research/deeplab

#学習データの作成１　アノテーション
アノテーションは、Labelmeというツールを使用しました。
https://github.com/wkentaro/labelme

参考ページ:セマンティックセグメンテーションのアノテーションツールlabelme

Labelmeを立ち上げ、Create Polygonsで人物とラケットの周りを点で囲ってpersonとracketとラベリングしてアノテーションデータを作成していきます。アノテーションデータはjsonファイルで保存されます。

#学習データの作成２　アノテーションデータとして作成したjsonファイルからラベル画像を作成
人物とラケットのセグメンテーション領域を色別で表示するラベル画像を作成します。

もともと、DeepLabには、20個のラベルが準備されており、15番目にPERSONが割り当てられているのでそれを流用します。また21番目のラベルとしてRACKETを追加します。色とラベルの対応表は↓のようになっています。

下記pythonスクリプトで、personとracketの領域を色別に表現する↓のようなラベル画像を作成します。

#labelmeでアノテーションしたjsonファイルからラベル画像を作成
import json
import glob
from PIL import Image,ImageDraw
import os

lbl_dir_gen = "./lbl/"

pil_img = Image.open('2007_000032.png')
palette = pil_img.getpalette()

path=('***/img_json/*.json')
fileList=glob.glob(path)

for i,f in enumerate(fileList):
    file = open(f, 'r')
    
    jsonData = json.load(file)
    h=jsonData['imageHeight']
    w=jsonData['imageWidth']
    points=jsonData['shapes'][0]['points']
    array=[]
    for j,p in enumerate(points):
        array.append(p[0])
        array.append(p[1])

    if(len(jsonData['shapes'])>1):
        points2=jsonData['shapes'][1]['points']
        array2=[]
        for j,p in enumerate(points2):
            array2.append(p[0])
            array2.append(p[1])
        
    basename=os.path.basename(f)
    root_ext_pair = os.path.splitext(basename)
    file.close()
    
    im = Image.new('P', (w, h))
    draw = ImageDraw.Draw(im)
    draw.polygon(array, fill=15, outline=255)#fill=21
    if(len(jsonData['shapes'])>1):
        draw.polygon(array2, fill=21, outline=255)#fill=21
    im.putpalette(palette)
    im.save(lbl_dir_gen + root_ext_pair[0] + ".png", quality = 100)

#学習データの作成３　グレースケール画像に変換
↑で作成したラベル画像をグレースケール画像に変換します。
グレースケールの0～255の中で、0～21の番地をラベル情報として活用することになります。
0(black)⇒背景
15⇒person
21⇒racket
255(white)⇒ラベル情報なし
1～14、16～20は省略

/models/research/deeplab/datasets/VOCdevkit/VOC2012/SegmentationClass

グレースケール画像に変換するためには、↑のフォルダにラベル画像を格納した状態にして、↓のコマンドを実行します。

python /models/research/deeplab/datasets/remove_gt_colormap.py

結果として、下記フォルダにグレースケール画像が作成されて保存されます。

/models/research/deeplab/datasets/VOCdevkit/VOC2012/SegmentationClassRaw

#学習データの作成４　TFRecords形式に変換する
作成したグレースケール画像をTFRecords形式のデータに変換します。

TFRecordsとはTensorFlowが推奨フォーマットとして提供しているもので、画像データなどメモリに収まらないような大きなデータを処理できるようバイナリ化したフォーマットのことです。
参考：TensorFlow推奨フォーマット「TFRecord」の作成と読み込み方法

/models/research/deeplab/datasets/VOCdevkit/VOC2012/ImageSets/Segmentation

に
train.txt
trainval.txt
val.txt
を作成します。

train.txtは、学習に用いるファイル名（拡張子は入れない）を列挙したリストファイルになります。
trainval.txt、val.txtも同様。

下記pythonスクリプトを使用すると、フォルダに格納されている全てのjpgファイルのファイル名を抽出してtrain.txtに転記します。ご活用ください。

#train.txtファイルの作成
import json
import glob
from PIL import Image,ImageDraw
import os

path='***/*.jpg'

fileList=glob.glob(path)
file_list = [os.path.basename(r) for r in fileList]
file = open('train.txt', 'w')
for i,f in enumerate(file_list):
    file.write(os.path.splitext(f)[0]+"\n")
file.close()

次に、data_generator.pyの中身を書き換えます。

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

trainには学習データの数を設定します。ここでは1464としてますが、train.txtで記入しているファイル数に合わせてください。trainval valも同様です。
num_classesにはラベルのクラス数を設定します。21個のクラス数なので21とします。
ignore_labelとは、無視する色を設定します。グレースケール画像で255(white）はラベルを割り当てていませんので、ignore_labelとして設定します。

こで、下記コマンドを実行させます。

python build_voc2012_data.py  --image_format="jpg" 

train.txt、trainval.txt、val.txtのリストに則したTFRecordsファイルが下記フォルダに生成されます。

/models/research/deeplab/datasets/tfrecord

はい、これでやっと学習データの作成は終了し、学習作業に移行します。

#学習
train.pyを実行すれば学習が開始されます。

python train.py   --logtostderr   --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=512,512   --train_batch_size=2  --training_number_of_steps=10000   --fine_tune_batch_norm=false   --train_logdir="./datasets/pascal_voc_seg/exp/train_on_trainval_set/train"  --dataset_dir="./datasets/tfrecord"  --dataset="original" --initialize_last_layer=true

引数の意味は、DeepLab v3+でオリジナルデータを学習してセグメンテーションできるようにするで詳しい説明があるので参考にしてください。

--training_number_of_stepsは学習ステップ数です。これが少ないと学習不足で推論画像が真っ黒、またはちゃんとセグメンテーションされていない画像になります。時間はかかりますが十分なステップ数で学習しましょう。

あと、ハマりポイントとしては、train_crop_size=512,512　のところですかね。学習に用いる画像が512x512より大きいとエラーが発生します。学習時に画像を切り抜いて学習するらしく、そのサイズを設定するのですが、
train_crop_size＞学習に用いる画像のサイズ
である必要があります。

ただ、train_crop_sizeをむやみに大きくすると、メモリーエラーが発生するので、学習に用いる画像を不要に大きくしないことが大切かと。

しかし、なぜ切り抜くサイズが画像サイズより大きい必要があるんですかね。
ふつうに考えたら大小関係が逆になる気がするのですが。。よくわかっていません。

学習ログ（ckptファイルなど）は↓に出力されるのでご確認ください。

/models/research/deeplab/datasets/pascal_voc_seg/exp/train_on_trainval_set/train

#評価
vis.pyを実行すると、学習したモデルで評価用の画像でセマンティックセグメンテーションします。

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="512,512"   --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/tfrecord"    --max_number_of_iterations=1   --eval_interval_secs=0 --dataset="original"

#推論用のファイルを作成
学習により、
model.ckpt-10000
ファイルが生成されています。
（10000の部分は、ステップ数により変わります）

このckptファイルを推論用のファイル frozen_inference_graph.pb に変換するよう、下記コマンドを実行します。

python export_model.py --logtostderr --checkpoint_path="*/models/research/deeplab/datasets/pascal_voc_seg/exp/train_on_trainval_set/train/model.ckpt-10000" --export_path="*/models/research/deeplab/datasets/pascal_voc_seg/exp/train_on_trainval_set/export/frozen_inference_graph.pb" --model_variant="xception_65" --atrous_rates=6 --atrous_rates=12 --atrous_rates=18 --output_stride=16 --decoder_output_stride=4 --num_classes=23 --crop_size=512 --crop_size=512 --inference_scales=1.0

#推論用のファイルを用いてセマンティックセグメンテーション動画を作成する
↓のような動画を作成します。

↑の残像動画は、セマンティックセグメンテーションで背景分離することで生成した「人＋ラケット」画像を重ね合わせてつくっています。 pic.twitter.com/T5MHOGZRLt

— おたこ (@otakoma) July 27, 2019

動画作成用のpythonコードは↓です。
これは　deeplab_demo.ipynbを改変してつくったものです。

PIL画像とOpenCV画像の変換の行き来があって計算効率の悪いコードになっていますがお許しを。

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,ImageDraw,ImageFont

import tensorflow as tf
import cv2

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():
    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):
    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 cv2pil(image):
    new_image = image.copy()
    #print(new_image.ndim)
    if new_image.ndim == 2:  # モノクロ
        pass
    elif new_image.shape[2] == 3:  # カラー
        new_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    elif new_image.shape[2] == 4:  # 透過
        new_image = cv2.cvtColor(image, cv2.COLOR_BGRA2RGBA)
    new_image = Image.fromarray(new_image)
    im=new_image
    return new_image

def pil2cv(image):
    ''' PIL型 -> OpenCV型 '''
    new_image = np.array(image)
    if new_image.ndim == 2:  # モノクロ
        pass
    elif new_image.shape[2] == 3:  # カラー
        new_image = new_image[:, :, ::-1]
    elif new_image.shape[2] == 4:  # 透過
        new_image = new_image[:, :, [2, 1, 0, 3]]
    return new_image

def seg2binary(image):
    img=np.asarray(image)
    im_f=np.where(img == 0, 0, 255)
    new_image = Image.fromarray(im_f.astype(np.uint8))
    return new_image
    
print("start")

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

frozen_path='./frozen_inference_graph.pb'
MODEL = DeepLabModel(frozen_path)

VIDEO_DATA = "serve-pose.mp4"
ESC_KEY = 0x1b
DURATION = 1.0
cv2.namedWindow("motion")
video = cv2.VideoCapture(VIDEO_DATA)
fps    = video.get(cv2.CAP_PROP_FPS)

#fourcc = cv2.VideoWriter_fourcc(*'MJPG')# 形式はMP4Vを指定
fourcc = cv2.VideoWriter_fourcc(*'DIVX')# 形式はMP4Vを指定

end_flag, frame_next = video.read()
height, width, channels = frame_next.shape
frame_pre = frame_next.copy()
original_im=cv2pil(frame_pre)
add_im, seg_map = MODEL.run(original_im)
#width, height = MODEL.returnSize(first_im)
width, height =add_im.size
out = cv2.VideoWriter('output.mp4',fourcc,fps, (int(width)*2, int(height)*2))

while(end_flag): 
 
    original_im=cv2pil(frame_pre)
    resized_im, seg_map = MODEL.run(original_im)
    seg_image = label_to_color_image(seg_map).astype(np.uint8)
    
    im = Image.fromarray(seg_image)
    image=Image.blend(resized_im,im,0.7)
    
    img=pil2cv(image)#blend
    
    cv2.imshow("motion", img)# モーション画像を表示 
    out.write(img)
    
    if cv2.waitKey(20) == ESC_KEY:# Escキー押下で終了
        break

    frame_pre = frame_next.copy()
    end_flag, frame_next = video.read()


cv2.destroyAllWindows()
out.release()
video.release()

print("end")