TensorFlow Object Detecion API のモデルを nnoir へ変換する例

本記事の内容は TensorFlow Object Detecion API のモデルを nnoir-onnx を利用して nnoir へ変換する例についてです。

nnoir

ONNX から変換することのできるニューラルネットの表現です。 読みは「ノワール」で
NN Optimization IR の略です。

github.com

nnoir-onnx

ONNX モデルから計算グラフを抽出し nnoir に変換する機能を提供します。
nnoir-onnx の README.md に記載されている ONNX オペレーターを変換可能です。

github.com

Google Colaboratory で実行可能な変換サンプルコード

github.com

https://github.com/Idein/tensorflow-object-detection-api-to-nnoir にあります。
このブログで例として変換しているのは、ssdlite_mobilenet_v2_coco_2018_05_09 です。
モデルの後処理に nnoir-onnx 1.0.13 で非対応のオペレーター (Slice) が後処理に含まれていることから、非対応オペレーターが含まれない位置までを nnoir に変換し、後処理を別途記述します。

環境

Google Colaboratory に以下のパッケージをインストール後、TensorFlow Object Detection API をセットアップしています。

!pip install -U tf2onnx==1.8.5
!pip install nnoir-onnx==1.0.13
!pip install ensemble-boxes==1.0.6

後処理を取り除く

モデルの後処理に nnoir-onnx 1.0.13 で非対応のオペレーターが含まれているため、add_postprocessing_op=False を設定し取り除いています。

!python /content/models/research/object_detection/export_tflite_ssd_graph.py \
--pipeline_config_path=/content/ssdlite_mobilenet_v2_coco_2018_05_09/pipeline.config \
--trained_checkpoint_prefix=/content/ssdlite_mobilenet_v2_coco_2018_05_09/model.ckpt \
--output_directory=/content \
--add_postprocessing_op=False

tflite_graph.pb の input name と output name は Netron を利用することで確認可能です。

!tflite_convert \
--graph_def_file=/content/tflite_graph.pb \
--output_file=/content/ssdlite_mobilenet_v2.tflite \
--input_format=TENSORFLOW_GRAPHDEF \
--output_format=TFLITE \
--inference_type=FLOAT \
--input_shapes=1,300,300,3 \
--input_arrays=normalized_input_image_tensor \
--output_arrays=raw_outputs/box_encodings,raw_outputs/class_predictions

Anchor を作成

tflite_graph.pb から nnoir 推論後の後処理に利用する anchors.npy を生成します。

# https://qiita.com/PINTO/items/1312d308b553362a8ebf#%EF%BC%96appendix
from tensorflow.python.platform import gfile
from tensorflow.python.framework import tensor_util

GRAPH_PB_PATH = '/content/tflite_graph.pb'

with tf.Session() as sess:
  with tf.gfile.GFile(GRAPH_PB_PATH,'rb') as f:
    graph_def = tf.GraphDef()
    graph_def.ParseFromString(f.read())
    sess.graph.as_default()
    tf.import_graph_def(graph_def, name='')
    graph_nodes=[n for n in graph_def.node]
    wts = [n for n in graph_nodes if n.op=='Const']
    for n in wts:
        if n.name == 'anchors':
            print("Name of the node - %s" % n.name)
            print("Value - ")
            anchors = tensor_util.MakeNdarray(n.attr['value'].tensor)
            print("anchors.shape =", anchors.shape)
            print(anchors)
            np.save('/content/anchors.npy', anchors)
            np.savetxt('/content/anchors.csv', anchors, delimiter=',')
            break

ONNX へ変換

!python -m tf2onnx.convert --tflite /content/ssdlite_mobilenet_v2.tflite \
--output /content/ssdlite_mobilenet_v2.onnx --opset 13

nnoir へ変換

元の onnx モデルと変換後の nnoirファイルパスを指定し変換を行います。

!onnx2nnoir -o /content/tf1_ssdlite_mobilenet_v2.nnoir /content/ssdlite_mobilenet_v2.onnx \
--graph_name ssdlite_mobilenet_v2

nnoir で推論

nnoir ファイルをロードし、前処理した画像を渡します。

nnoir_model = nnoir.load('/content/tf1_ssdlite_mobilenet_v2.nnoir')

# 前処理
image = cv2.imread("/content/models/research/object_detection/test_images/image1.jpg")
image = cv2.resize(image, (300, 300))
image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
image_expanded = np.expand_dims(image_rgb.astype(np.float32), axis=0)
image_normalized =  (np.float32(image_expanded) - 127.5) / 127.5

box_encodings, class_predictions = nnoir_model.run(image_normalized)

Post Process を取り除いているため bounding box をデコードするコードが必要です。

参考：tensorflow/tensorflow/lite/detection_postprocess.cc

# pipeline.config に記載されている値
y_scale = 10.0
x_scale = 10.0
h_scale = 5.0
w_scale = 5.0

def decode_box_encodings(box_encoding, anchors, num_boxes):
    decoded_boxes = np.zeros((num_boxes, 4), dtype=np.float32)
    for i in range(num_boxes):
        ycenter = box_encoding[i][0] / y_scale * anchors[i][2] + anchors[i][0]
        xcenter = box_encoding[i][1] / x_scale * anchors[i][3] + anchors[i][1]
        half_h = 0.5 * math.exp((box_encoding[i][2] / h_scale)) * anchors[i][2]
        half_w = 0.5 * math.exp((box_encoding[i][3] / w_scale)) * anchors[i][3]
        decoded_boxes[i][0] = (xcenter - half_w) # xmin
        decoded_boxes[i][1] = (ycenter - half_h) # ymin
        decoded_boxes[i][2] = (xcenter + half_w) # xmax
        decoded_boxes[i][3] = (ycenter + half_h) # ymax
    return decoded_boxes

# 作成した anchors.npy を読み込む
anchors = np.load('/content/anchors.npy')
box_decoded = decode_box_encodings(box_encodings[0], anchors, 1917)
box_not_background = np.take(box_decoded, indexes_not_background, axis=0)
box_clipped = box_not_background.clip(0, 1)

Non-maximum Suppression は ZFTurbo/Weighted-Boxes-Fusion のものを利用しています。

github.com

nms_boxes, nms_scores, nms_labels = nms(box_clipped.tolist(), scores_not_background.tolist(), labels_not_background.tolist(), iou_thr=0.5)

Bounding Box と label の描画

image_pil = Image.fromarray(image_rgb)

# 0が背景
coco_labels = ['background']
# https://raw.githubusercontent.com/amikelive/coco-labels/master/coco-labels-paper.txt
with open("/content/coco-labels-paper.txt",'r') as f:
  for line in f:
    coco_labels.append(line.rstrip())

coco_labels_length = len(coco_labels)
draw = ImageDraw.Draw(image_pil)
for box, scores, label in zip(nms_boxes, nms_scores, nms_labels):
  if label > 0 and label < coco_labels_length:
    xmin = box[0]*300
    ymin = box[1]*300
    xmax = box[2]*300
    ymax = box[3]*300
    draw.rectangle((xmin, ymin, xmax, ymax), outline=(0, 0, 255), width=2)
    draw.text([xmin+2, ymin+2], text=coco_labels[label])

image_pil

nnoir を用いた推論