Python ONNX Runtime으로 Resnet 테스트 하기

목적

ONNX Model의 검증이 필요할 때에 간략하게 테스트 하는 방법

Dependency package

Pillow가 없을 경우 pip 설치가 필요함

출력

해당 모델은 customized된 Resnet 모델로 Output이 10개 이하임.
그래서 수치의 결과물을 보기 위한 것으로, 모든 값을 한 행에 출력을 하게끔 작성함.
필요에 따라 적절한 수정이 필요함.

Python Code

import os
import onnxruntime as rt
from PIL import Image
import numpy as np

model_path = 'model path str'
image_dir = 'image path str'
width = 224
height = 224

sess = rt.InferenceSession(model_path)

for filename in os.listdir(image_dir):
    if filename.endswith('.jpg') or filename.endswith('.jpeg'):
        img = Image.open(os.path.join(image_dir, filename))

        img = img.resize((width, height))
        img = np.array(img).astype('float32')

        img = np.expand_dims(img, axis=0)
        img = np.transpose(img, (0, 3, 2, 1))

        input_name = sess.get_inputs()[0].name
        result = sess.run(None, {input_name: img})

        result = np.reshape(result, (-1))
        print("{}: {}".format(filename, ', '.join(np.array2string(r.flatten(), separator=', ', max_line_width=np.inf) for r in result)))