直接贴code
加载，推理
import onnxruntime as ortimport torchimport timeimport cv2import 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].nameonnx_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 = 2048img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)img = np.transpose(img, (2,0,1)) #HWC ->CHWinput_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 imageimg2 = 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)