Edited at

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


やりたいこと

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


なぜ「人物」と「テニスラケット」をセマンティックセグメンテーションするのか

↓のような感じでテニスのサーブフォームで残像モーション動画を生成したかったからです。

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

残像モーション動画生成システムについてはこちら

【テニスのフォーム解析】残像モーション動画生成システム


環境

Windows10 GPU1080ti

Tensorflow 1.14.0

DeepLab v3+


DeepLab v3+のインストールはこちらから

DeepLab v3+のgithubページ

https://github.com/tensorflow/models/tree/master/research/deeplab


学習データの作成1 アノテーション

アノテーションは、Labelmeというツールを使用しました。

https://github.com/wkentaro/labelme

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

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

labelme (2).png


学習データの作成2 アノテーションデータとして作成したjsonファイルからラベル画像を作成

人物とラケットのセグメンテーション領域を色別で表示するラベル画像を作成します。

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

キャプチャ.PNG

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

serve01-01_0000.png


#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)


学習データの作成3 グレースケール画像に変換

↑で作成したラベル画像をグレースケール画像に変換します。

グレースケールの0~255の中で、0~21の番地をラベル情報として活用することになります。

0(black)⇒背景

15⇒person

21⇒racket

255(white)⇒ラベル情報なし

1~14、16~20は省略

serve01-01_0000.png

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

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

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

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

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


学習データの作成4 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


推論用のファイルを用いてセマンティックセグメンテーション動画を作成する

↓のような動画を作成します。

動画作成用の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")