【Python/OpenCV4】ONNX形式の学習済みモデルをロード

Python版OpenCVでONNX形式の学習済みモデルをロードする方法をソースコード付きでまとめました。

【OpenCV4】ONNX形式の学習済みモデルをロード

ONNX（Open Neural Network eXchange：オニキス）とは、様々な機械学習フレームワークの間で、ディープラーニングの学習モデルを効率的にやり取りするための共通フォーマットです。
例えばKerasで作成した学習済みモデルをONNX形式で保存すれば、OpenCVなど他のフレームワークでもそのまま利用できます。

OpenCVではバージョン4から「cv2.dnn.readNetFromONNX()」でONNXをロードできるようになりました。

サンプルコード

サンプルプログラムのソースコードです。

# -*- coding: utf-8 -*-
import numpy as np
import matplotlib.pyplot as plt
import cv2
import sklearn
from sklearn import datasets, model_selection, svm, metrics

# データセットのロード
mnist = datasets.fetch_mldata('MNIST original', data_home='data/src/download/')

N = np.arange(len(mnist['data']))

n = np.random.choice(N, 5, replace=False);
for i in n:
    img = mnist['data'].reshape((28, 28))
    img = img / 255.
    img = img.reshape((-1, 1, 28, 28))

    # ONNX形式の学習済みモデルをロード
    onnx = cv2.dnn.readNetFromONNX('model.onnx')

    # 入力
    onnx.setInput(img)

    # 検証
    history = onnx.forward()
    print('予測: ', np.argmax(history))
    print('正解: ', int(mnist['target']))

'''
予測:4
正解:4
'''

model.onnxは「https://github.com/onnx/models/tree/master/mnist」から入手できます。

NN によって予想された結果と、正解が出力されます。