TensorFlowのモデルファイル.pbをreadNetFromTensorflow()で読み込ませる。 #Python

本記事は自前のデータセットで学習した物体検出モデルを高速に処理するためにOpenCVのreadNetFromTensorflow()を使おうとしてはまった話です。

その対策についてまとめます。

動作環境

Python: 3.7.6
Tensorflow: 1.14.0
OpenCV: 4.2.0

readNetFromTensorflow()について

学習済みモデルをopencvで動かす為のコードが以下。

# How to load a Tensorflow model using OpenCV
# Jean Vitor de Paulo Blog - https://jeanvitor.com/tensorflow-object-detecion-opencv/

import cv2

# Load a model imported from Tensorflow
tensorflowNet = cv2.dnn.readNetFromTensorflow('frozen_inference_graph.pb', 'graph.pbtxt')

# Input image
img = cv2.imread('img.jpg')
rows, cols, channels = img.shape

# Use the given image as input, which needs to be blob(s).
tensorflowNet.setInput(cv2.dnn.blobFromImage(img, size=(300, 300), swapRB=True, crop=False))

# Runs a forward pass to compute the net output
networkOutput = tensorflowNet.forward()

# Loop on the outputs
for detection in networkOutput[0,0]:

    score = float(detection[2])
    if score > 0.2:

        left = detection[3] * cols
        top = detection[4] * rows
        right = detection[5] * cols
        bottom = detection[6] * rows

        #draw a red rectangle around detected objects
        cv2.rectangle(img, (int(left), int(top)), (int(right), int(bottom)), (0, 0, 255), thickness=2)

# Show the image with a rectagle surrounding the detected objects 
cv2.imshow('Image', img)
cv2.waitKey()
cv2.destroyAllWindows()

詳しくはHow to load Tensorflow models with OpenCVが参考になります。

モデルの呼び出しに必要なファイル

tensorflowNet = cv2.dnn.readNetFromTensorflow('frozen_inference_graph.pb', 'graph.pbtxt')

frozen_inference_graph.pb（凍結グラフ）
graph.pbtxt（モデル構造）

opencvでtensorflowモデルを読み込むには上記の２つのファイルが必要になります。

ファイルの作成

Step１学習の実行

学習にはtrain.pyを使用しました。

$ python train.py \
      --logtostderr \
      --train_dir=log \
      --pipeline_config_path=ssd_mobilenet_v1_coco.config

ここで、logというフォルダができ、その中に

model.ckpt
pipeline.config

が作成されます。この２つのファイルを次のステップで使います。

Step2 frozen_inference_graph.pb（凍結グラフ）作成

こちらexport_inference_graph.pyを使用します。

$ python export_inference_graph.py \
    --input_type image_tensor \
    --pipeline_config_path log/pipeline.config \
    --trained_checkpoint_prefix log/model.ckpt-3000 \ #3000はstep数
    --output_directory model_data

実行すると、model_dataフォルダの中に

frozen_inference_graph.pb
pipeline.config

が見つかります。これらを次に使います。

Step3 graph.pbtxt（モデル構造）作成

こちらtf_text_graph_ssd.pyを使用します。

$ python tf_text_graph_ssd.py \
    --input model_data/frozen_inference_graph.pb \
    --output graph_data/graph.pbtxt \
    --config model_data/pipeline.config

実行すると、graph_dataフォルダの中にgraph.pbtxtができています。

これで、凍結グラフとモデル構造が揃いました。

モデルの読み込み

# モデルの読み込み
model = cv2.dnn.readNetFromTensorflow('graph_data/frozen_inference_graph.pb', 
                                      'graph_data/graph.pbtxt')
print('読み込み完了')

すると・・・

---------------------------------------------------------------------------
error                                     Traceback (most recent call last)
<ipython-input-44-afa605a67bd6> in <module>
      1 # モデルの読み込み
      2 model = cv2.dnn.readNetFromTensorflow('graph_data/frozen_inference_graph.pb', 
----> 3                                       'graph_data/graph.pbtxt')
      4 print('読み込み完了')

error: OpenCV(4.2.0) /Users/travis/build/skvark/opencv-python/opencv/modules/dnn/src/tensorflow/tf_importer.cpp:544: error: (-2:Unspecified error) Input layer not found: FeatureExtractor/MobilenetV1/Conv2d_0/weights/read/_166__cf__169 in function 'connect'

Input layer not found: /// in function 'connect'
入力レイヤーがみつかりません！
このerrorに四苦八苦した結果、次のような方法で落ち着きました。

Input layer not foundが出た時の対策。

【改善前】

node {
  name: "image_tensor"
  op: "Placeholder"
  attr {
    key: "dtype"
    value {
      type: DT_UINT8
    }
  }
  attr {
    key: "shape"
    value {
      shape {
        dim {
          size: -1
        }
        dim {
          size: 300
        }
        dim {
          size: 300
        }
        dim {
          size: 3
        }
      }
    }
  }
}
node {
  name: "FeatureExtractor/MobilenetV1/MobilenetV1/Conv2d_0/BatchNorm/batchnorm/mul_1"
  op: "Conv2D"
  input: "FeatureExtractor/MobilenetV1/Conv2d_0/weights/read/_166__cf__169"
  attr {
    key: "data_format"
    value {
      s: "NHWC"
    }
  }

graph.pbtxtの先頭をみてみると、確かに入力レイヤーが存在しないことがわかります。
ならば入力レイヤーを作ってやれば良いということで、

【改善後】

node {
  name: "image_tensor"
  op: "Placeholder"
  attr {
    key: "dtype"
    value {
      type: DT_UINT8
    }
  }
  attr {
    key: "shape"
    value {
      shape {
        dim {
          size: -1
        }
        dim {
          size: 300
        }
        dim {
          size: 300
        }
        dim {
          size: 3
        }
      }
    }
  }
}
node {                           #こっから
  name: "Preprocessor/mul"
  op: "Mul"
  input: "image_tensor"
  input: "Preprocessor/mul/x"
}
node {
  name: "Preprocessor/sub"
  op: "Sub"
  input: "Preprocessor/mul"
  input: "Preprocessor/sub/y"
}                                #ここまで追加
node {
  name: "FeatureExtractor/MobilenetV1/MobilenetV1/Conv2d_0/BatchNorm/batchnorm/mul_1"
  op: "Conv2D"
  input: "Preprocessor/sub"      #追加
  input: "FeatureExtractor/MobilenetV1/Conv2d_0/weights/read/_166__cf__169"
  attr {
    key: "data_format"
    value {
      s: "NHWC"
    }
  }

このようにgraph.pbtxtに直接変更を加えました。
そして

無事読み込むことができました。

そもそもなぜ入力レイヤーが定義されていなかったのか、その辺りの原因究明はできておりません。
もしご存知の方がおられましたらコメントいただけると幸いです。

TensorFlowのモデルファイル.pbをreadNetFromTensorflow()で読み込ませる。