ONNXRuntime模型推理

直接贴code

• 加载，推理

import onnxruntime as ort
import torch
import time

import cv2
import numpy as np


def 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 image

mCityscapesColors = [
    (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)