GluonとONNX

はじめに

touch-sp.hatenablog.com
以前このようなものを書いた。
その後、ONNX形式で保存されたモデルをGluonモデルとして読み込めたら便利だなと思っていたら、実は簡単にできたという話。

環境

Windows7 64bit
Python 3.6.6
GPUなし

バージョン確認(pip freeze)

astroid==2.0.4
certifi==2018.10.15
chardet==3.0.4
colorama==0.4.0
graphviz==0.8.4
idna==2.6
isort==4.3.4
lazy-object-proxy==1.3.1
mccabe==0.6.1
mxnet==1.3.1b20181018
numpy==1.14.6
onnx==1.3.0
Pillow==5.3.0
protobuf==3.6.1
pylint==2.1.1
requests==2.18.4
six==1.11.0
typed-ast==1.1.0
typing==3.6.6
typing-extensions==3.6.6
urllib3==1.22
wrapt==1.10.11

学習済みモデルのダウンロード

  • こちら』から学習済みモデルをダウンロード

サンプル画像のダウンロード

  • こちら』からサンプル画像をダウンロード

実行ファイル

import numpy as np
from PIL import Image
import mxnet as mx
import mxnet.contrib.onnx as onnx_mxnet

#データの読み込み
input_image_dim = 224
img = Image.open('super_res_input.jpg').resize((input_image_dim, input_image_dim))
img_ycbcr = img.convert("YCbCr")
img_y, IMG_CB, IMG_CR = img_ycbcr.split()
INPUT_IMG = np.array(img_y)[np.newaxis, np.newaxis, :, :]

#モデルの読み込み
model = onnx_mxnet.onnx2mx.import_to_gluon.import_to_gluon('super_resolution.onnx', ctx=mx.cpu())

#モデルの適応
output = model(mx.nd.array(INPUT_IMG))

#画像の保存
output = mx.nd.squeeze(output)
img_out_y = Image.fromarray(output.asnumpy().clip(0,255).astype(np.uint8), mode='L')
result_img = Image.merge(
        "YCbCr", [img_out_y,
                  IMG_CB.resize(img_out_y.size, Image.BICUBIC),
                  IMG_CR.resize(img_out_y.size, Image.BICUBIC)]).convert("RGB")
result_img.save("super_res_output.jpg")

おまけ

どのようなデータを入力するべきかを知るには以下の1行を実行すれば良い

>>> onnx_mxnet.get_model_metadata('super_resolution.onnx')
{'input_tensor_data': [('1', (1, 1, 224, 224))], 'output_tensor_data': [('27', (1, 1, 672,
672))]}