ONNXRuntime模型推理

本文主要是介绍ONNXRuntime模型推理，希望对大家解决编程问题提供一定的参考价值，需要的开发者们随着小编来一起学习吧！

直接贴code

• 加载，推理

import onnxruntime as ort
import torch
import timeimport cv2
import numpy as npdef time_sync():if torch.cuda.is_available():torch.cuda.synchronize()return time.time()ort_session = ort.InferenceSession('./semseg.onnx')
onnx_input_name = ort_session.get_inputs()[0].name
onnx_outputs_names = ort_session.get_outputs()
output_names = []
for o in onnx_outputs_names:output_names.append(o.name)img = cv2.imread('./demo.png')
img = cv2.resize(img, (2048,1024)) # height = 1024, width = 2048
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = np.transpose(img, (2,0,1)) #HWC ->CHW
input_blob = np.expand_dims(img,axis=0).astype(np.float32)total_time = 0;
for i in range(0, 10): #推理10次，统计平均时间start_t = time_sync()onnx_result = ort_session.run(output_names, input_feed = {onnx_input_name:input_blob})end_t = time_sync()total_time += ((end_t - start_t) * 1E3)print(f'{i+1} Speed: %.1fms ' % ((end_t - start_t) * 1E3))print(f'Avg Speed: %.1fms ' % (total_time/10.0))

• 输出保存semseg图

//draw semseg mask image
img2 = cv2.imread('./demo.png')
img2 = cv2.resize(img2, (2048,1024)) # you can also use other way to create a temp imagemCityscapesColors = [(128, 64,128), (244, 35,232), ( 70, 70, 70), (102,102,156), (190,153,153), (153,153,153), (250,170, 30), (220,220,  0), (107,142, 35), (152,251,152), ( 70,130,180), (220, 20, 60), (255,  0,  0), (  0,  0,142), (  0,  0, 70), (  0, 60,100), (  0, 80,100), (  0,  0,230), (119, 11, 32)];for h in range(0, img2.shape[0]):for w in range(0, img2.shape[1]):img2[h,w] = mCityscapesColors[onnx_result[0][0][0][h][w]]
cv2.imwrite('./mask_semseg.png', img2) 

• 统计语义分割的类别

a = onnx_result[0][0][0]
list_b = list(np.array(a).flatten())
print set(list_b)

这篇关于ONNXRuntime模型推理的文章就介绍到这儿，希望我们推荐的文章对编程师们有所帮助！

---