30 lines
797 B
Python
30 lines
797 B
Python
|
"""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()
|