TensorFlowセグメンテーションモデルの変換とOpenCVでの実行

目標

このチュートリアルでは、以下の方法を学ぶ:

• TensorFlow (TF) のセグメンテーションモデルを変換する
• 変換したTensorFlowモデルをOpenCVで実行する
• TensorFlowモデルとOpenCV DNNモデルの評価を行う

上記の項目を、DeepLabアーキテクチャを例に探求していく。

はじめに

OpenCV APIを用いた TensorFlow分類 モデルとセグメンテーションモデルの移行パイプラインに関わる主要な概念は、グラフ最適化の段階を除けばほぼ同じである。TensorFlowモデルを cv.dnn.Net に変換する最初のステップは、フリーズしたTFモデルグラフを取得することである。フリーズグラフは、モデルのグラフ構造と、必要な変数（例えば重み）の保持された値とを組み合わせて定義する。通常、フリーズグラフは protobuf (.pb) ファイルに保存される。生成されたセグメンテーションモデルの .pb ファイルを cv.dnn.readNetFromTensorflow で読み込むには、TFの graph transform tool でグラフを修正する必要がある。

実践

この章では、以下の点を扱う。

1. TF分類モデルの変換パイプラインを作成し、推論を提供する
2. TF分類モデルを評価しテストする

単に評価やモデルパイプラインのテストを実行したいだけであれば、「Model Conversion Pipeline」のチュートリアル部分は省略してもよい。

モデル変換パイプライン

本節のコードは dnn_model_runner モジュール内にあり、次の行で実行できる:

python -m dnn_model_runner.dnn_conversion.tf.segmentation.py_to_py_deeplab

TensorFlowのセグメンテーションモデルは TensorFlow Research Models のセクションにあり、ここには公開された研究論文に基づくモデルの実装が含まれている。学習済みのTF DeepLabV3のアーカイブを、以下のリンクから取得する:

http://download.tensorflow.org/models/deeplabv3_mnv2_pascal_trainval_2018_01_29.tar.gz

フリーズグラフを取得する完全なパイプラインは deeplab_retrievement.py に記述されている:

def get_deeplab_frozen_graph():
 # define model path to download
    models_url = 'http://download.tensorflow.org/models/'
    mobilenetv2_voctrainval = 'deeplabv3_mnv2_pascal_trainval_2018_01_29.tar.gz'
 
 # construct model link to download
    model_link = models_url + mobilenetv2_voctrainval
 
 try:
        urllib.request.urlretrieve(model_link, mobilenetv2_voctrainval)
 except Exception:
        print("TF DeepLabV3 was not retrieved: {}".format(model_link))
 return
 
    tf_model_tar = tarfile.open(mobilenetv2_voctrainval)
 
 # iterate the obtained model archive
 for model_tar_elem in tf_model_tar.getmembers():
 # check whether the model archive contains frozen graph
 if TF_FROZEN_GRAPH_NAME in os.path.basename(model_tar_elem.name):
 # extract frozen graph
            tf_model_tar.extract(model_tar_elem, FROZEN_GRAPH_PATH)
 
    tf_model_tar.close()

このスクリプトを実行すると:

python -m dnn_model_runner.dnn_conversion.tf.segmentation.deeplab_retrievement

deeplab/deeplabv3_mnv2_pascal_trainval に frozen_inference_graph.pb が得られる。

OpenCVでネットワークを読み込む前に、抽出した frozen_inference_graph.pb を最適化する必要がある。グラフを最適化するために、デフォルトの引数でTFの TransformGraph を使用する:

DEFAULT_OPT_GRAPH_NAME = "optimized_frozen_inference_graph.pb"
DEFAULT_INPUTS = "sub_7"
DEFAULT_OUTPUTS = "ResizeBilinear_3"
DEFAULT_TRANSFORMS = "remove_nodes(op=Identity)" \
 " merge_duplicate_nodes" \
 " strip_unused_nodes" \
 " fold_constants(ignore_errors=true)" \
 " fold_batch_norms" \
 " fold_old_batch_norms"
 
 
def optimize_tf_graph(
        in_graph,
        out_graph=DEFAULT_OPT_GRAPH_NAME,
        inputs=DEFAULT_INPUTS,
        outputs=DEFAULT_OUTPUTS,
        transforms=DEFAULT_TRANSFORMS,
        is_manual=True,
        was_optimized=True
):
 # ...
 
    tf_opt_graph = TransformGraph(
        tf_graph,
        inputs,
        outputs,
        transforms
    )

グラフ最適化処理を実行するには、次の行を実行する:

python -m dnn_model_runner.dnn_conversion.tf.segmentation.tf_graph_optimizer --in_graph deeplab/deeplabv3_mnv2_pascal_trainval/frozen_inference_graph.pb

その結果、deeplab/deeplabv3_mnv2_pascal_trainval ディレクトリに optimized_frozen_inference_graph.pb が含まれる。

モデルのグラフを取得したら、以下のステップを見ていこう:

1. TFの frozen_inference_graph.pb グラフを読み込む
2. 最適化したTFフリーズグラフをOpenCV APIで読み込む
3. 入力データを準備する
4. 推論を実行する
5. 予測結果から色付けされたマスクを取得する
6. 結果を可視化する

# get TF model graph from the obtained frozen graph
deeplab_graph = read_deeplab_frozen_graph(deeplab_frozen_graph_path)
 
# read DeepLab frozen graph with OpenCV API
opencv_net = cv2.dnn.readNetFromTensorflow(opt_deeplab_frozen_graph_path)
print("OpenCV model was successfully read. Model layers: \n", opencv_net.getLayerNames())
 
# get processed image
original_img_shape, tf_input_blob, opencv_input_img = get_processed_imgs("test_data/sem_segm/2007_000033.jpg")
 
# obtain OpenCV DNN predictions
opencv_prediction = get_opencv_dnn_prediction(opencv_net, opencv_input_img)
 
# obtain TF model predictions
tf_prediction = get_tf_dnn_prediction(deeplab_graph, tf_input_blob)
 
# get PASCAL VOC classes and colors
pascal_voc_classes, pascal_voc_colors = read_colors_info("test_data/sem_segm/pascal-classes.txt")
 
# obtain colored segmentation masks
opencv_colored_mask = get_colored_mask(original_img_shape, opencv_prediction, pascal_voc_colors)
tf_colored_mask = get_tf_colored_mask(original_img_shape, tf_prediction, pascal_voc_colors)
 
# obtain palette of PASCAL VOC colors
color_legend = get_legend(pascal_voc_classes, pascal_voc_colors)
 
cv2.imshow('TensorFlow Colored Mask', tf_colored_mask)
cv2.imshow('OpenCV DNN Colored Mask', opencv_colored_mask)
 
cv2.imshow('Color Legend', color_legend)

モデル推論を行うために、PASCAL VOC検証用データセットから以下の画像を使用する:

目標とするセグメンテーション結果は次のとおりである:

PASCAL VOCの色のデコードおよび予測されたマスクへのマッピングのために、PASCAL VOCクラスの完全な一覧と対応する色を含むpascal-classes.txtファイルも必要となる。

学習済みのTF DeepLabV3 MobileNetV2を例に、各ステップをより詳しく見ていこう:

• TFの frozen_inference_graph.pb グラフを読み込む:

# init deeplab model graph
model_graph = tf.Graph()
 
# obtain
with tf.io.gfile.GFile(frozen_graph_path, 'rb') as graph_file:
    tf_model_graph = GraphDef()
tf_model_graph.ParseFromString(graph_file.read())
 
with model_graph.as_default():
    tf.import_graph_def(tf_model_graph, name='')

• 最適化したTFフリーズグラフをOpenCV APIで読み込む:

# read DeepLab frozen graph with OpenCV API
opencv_net = cv2.dnn.readNetFromTensorflow(opt_deeplab_frozen_graph_path)

• cv2.dnn.blobFromImage 関数で入力データを準備する:

# read the image
input_img = cv2.imread(img_path, cv2.IMREAD_COLOR)
input_img = input_img.astype(np.float32)
 
# preprocess image for TF model input
tf_preproc_img = cv2.resize(input_img, (513, 513))
tf_preproc_img = cv2.cvtColor(tf_preproc_img, cv2.COLOR_BGR2RGB)
 
# define preprocess parameters for OpenCV DNN
mean = np.array([1.0, 1.0, 1.0]) * 127.5
scale = 1 / 127.5
 
# prepare input blob to fit the model input:
# 1. subtract mean
# 2. scale to set pixel values from 0 to 1
input_blob = cv2.dnn.blobFromImage(
    image=input_img,
    scalefactor=scale,
    size=(513, 513),  # img target size
    mean=mean,
    swapRB=True,  # BGR -> RGB
    crop=False # center crop
)

cv2.dnn.blobFromImage 関数における前処理の順序に注意してほしい。まず平均値が減算され、その後でピクセル値が指定したスケールで乗算される。したがって、TFの画像前処理パイプラインを再現するには、mean に 127.5 を掛ける。もう一つの重要な点はTF DeepLabの画像前処理である。画像をTFモデルに渡すには適切な形状を構築するだけでよく、残りの画像前処理は feature_extractor.py に記述されており、自動的に呼び出される。

• OpenCVの cv.dnn_Net による推論を行う:

# set OpenCV DNN input
opencv_net.setInput(preproc_img)
 
# OpenCV DNN inference
out = opencv_net.forward()
print("OpenCV DNN segmentation prediction: \n")
print("* shape: ", out.shape)
 
# get IDs of predicted classes
out_predictions = np.argmax(out[0], axis=0)

上記のコードを実行すると、次の出力が得られる:

OpenCV DNN segmentation prediction:
* shape:  (1, 21, 513, 513)

21個の予測チャンネルのそれぞれ（21はPASCAL VOCクラスの数を表す）には確率が含まれており、ピクセルがそのPASCAL VOCクラスにどれだけ対応しやすいかを示す。

• TFモデルによる推論を行う:

preproc_img = np.expand_dims(preproc_img, 0)
 
# init TF session
tf_session = Session(graph=model_graph)
 
input_tensor_name = "ImageTensor:0",
output_tensor_name = "SemanticPredictions:0"
 
# run inference
out = tf_session.run(
    output_tensor_name,
    feed_dict={input_tensor_name: [preproc_img]}
)
 
print("TF segmentation model prediction: \n")
print("* shape: ", out.shape)

TFの推論結果は次のとおりである:

TF segmentation model prediction:
* shape:  (1, 513, 513)

TensorFlowの予測には、対応するPASCAL VOCクラスのインデックスが含まれる。

• OpenCVの予測を色付きマスクに変換する:

mask_height = segm_mask.shape[0]
mask_width = segm_mask.shape[1]
 
img_height = original_img_shape[0]
img_width = original_img_shape[1]
 
# convert mask values into PASCAL VOC colors
processed_mask = np.stack([colors[color_id] for color_id in segm_mask.flatten()])
 
# reshape mask into 3-channel image
processed_mask = processed_mask.reshape(mask_height, mask_width, 3)
processed_mask = cv2.resize(processed_mask, (img_width, img_height), interpolation=cv2.INTER_NEAREST).astype(
    np.uint8)
 
# convert colored mask from BGR to RGB
processed_mask = cv2.cvtColor(processed_mask, cv2.COLOR_BGR2RGB)

このステップでは、セグメンテーションマスクの確率を、予測されたクラスの適切な色にマッピングする。結果を見てみよう:

• TFの予測を色付きマスクに変換する:

colors = np.array(colors)
processed_mask = colors[segm_mask[0]]
 
img_height = original_img_shape[0]
img_width = original_img_shape[1]
 
processed_mask = cv2.resize(processed_mask, (img_width, img_height), interpolation=cv2.INTER_NEAREST).astype(
    np.uint8)
 
# convert colored mask from BGR to RGB for compatibility with PASCAL VOC colors
processed_mask = cv2.cvtColor(processed_mask, cv2.COLOR_BGR2RGB)

結果は次の通りである。

その結果、2つの等しいセグメンテーションマスクが得られる。

モデルの評価

dnn/samples で提供されている dnn_model_runner モジュールを使うと、PASCAL VOCデータセット上で完全な評価パイプラインを実行し、DeepLab MobileNetモデルの実行テストを行える。

評価モード

以下の行は、評価モードでモジュールを実行する例を表す:

python -m dnn_model_runner.dnn_conversion.tf.segmentation.py_to_py_segm

モデルはOpenCVの cv.dnn_Net オブジェクトに読み込まれる。TFモデルとOpenCVモデルの評価結果（ピクセル精度、平均IoU、推論時間）はログファイルに書き込まれる。推論時間の値はチャートにも描画され、得られたモデル情報を一般化する。

必要な評価設定は test_config.py で定義されている:

@dataclass
class TestSegmConfig:
    frame_size: int = 500
    img_root_dir: str = "./VOC2012"
    img_dir: str = os.path.join(img_root_dir, "JPEGImages/")
    img_segm_gt_dir: str = os.path.join(img_root_dir, "SegmentationClass/")
 # reduced val: https://github.com/shelhamer/fcn.berkeleyvision.org/blob/master/data/pascal/seg11valid.txt
    segm_val_file: str = os.path.join(img_root_dir, "ImageSets/Segmentation/seg11valid.txt")
    colour_file_cls: str = os.path.join(img_root_dir, "ImageSets/Segmentation/pascal-classes.txt")

これらの値は、選択したモデルパイプラインに合わせて変更できる。

テストモード

以下の行は、モジュールをテストモードで実行する例であり、モデル推論の手順を提供する:

python -m dnn_model_runner.dnn_conversion.tf.segmentation.py_to_py_segm --test True --default_img_preprocess <True/False> --evaluate False

ここで default_img_preprocess キーは、scale、mean、std などの特定の値でモデルのテスト処理をパラメータ化するか、デフォルト値を使用するかを定義する。

テスト設定は test_config.py の TestSegmModuleConfig クラスで表現される:

@dataclass
class TestSegmModuleConfig:
    segm_test_data_dir: str = "test_data/sem_segm"
    test_module_name: str = "segmentation"
    test_module_path: str = "segmentation.py"
    input_img: str = os.path.join(segm_test_data_dir, "2007_000033.jpg")
    model: str = ""
 
    frame_height: str = str(TestSegmConfig.frame_size)
    frame_width: str = str(TestSegmConfig.frame_size)
    scale: float = 1.0
    mean: List[float] = field(default_factory=lambda: [0.0, 0.0, 0.0])
    std: List[float] = field(default_factory=list)
    crop: bool = False
    rgb: bool = True
    classes: str = os.path.join(segm_test_data_dir, "pascal-classes.txt")

デフォルトの画像前処理オプションは default_preprocess_config.py で定義されている:

tf_segm_input_blob = {
 "scale": str(1 / 127.5),
 "mean": ["127.5", "127.5", "127.5"],
 "std": [],
 "crop": "False",
 "rgb": "True"
}

モデルテストの基盤は samples/dnn/segmentation.py で表現される。segmentation.py は、--input に変換済みモデルを指定し、cv2.dnn.blobFromImage 用の引数を設定すれば、単独で実行できる。

「Model Conversion Pipeline」で説明したOpenCVのステップを、dnn_model_runner を使って一から再現するには、以下の行を実行する:

python -m dnn_model_runner.dnn_conversion.tf.segmentation.py_to_py_segm --test True --default_img_preprocess True --evaluate False