TensorRT-Demo/modnet/test_onnx.py

30 lines
797 B
Python
Raw Normal View History

2023-03-06 20:44:29 +08:00
"""run_onnx.py
A simple script for verifying the modnet.onnx model.
I used the following image for testing:
$ gdown --id 1fkyh03NEuSwvjFttYVwV7TjnJML04Xn6 -O image.jpg
"""
import numpy as np
import cv2
import onnx
import onnxruntime
img = cv2.imread('image.jpg')
img = cv2.resize(img, (512, 288), cv2.INTER_AREA)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = img.transpose((2, 0, 1)).astype(np.float32)
img = (img - 127.5) / 127.5
img = np.expand_dims(img, axis=0)
session = onnxruntime.InferenceSession('modnet.onnx', None)
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name
result = session.run([output_name], {input_name: img})
matte = np.squeeze(result[0])
cv2.imshow('Matte', (matte * 255.).astype(np.uint8))
cv2.waitKey(0)
cv2.destroyAllWindows()