熱影像物件偵測 (Object Detection on Thermal Image)- 2 使用TensorRT C++ Win10 - 測試onnx model
上篇文章我們得到了onnx格式的model，先做個測試確認這個model是可以使用的，需要在onnx framework下執行這個model，這邊使用到python onnx runtime。
一、下載onnx runtime
可以在 https://www.onnxruntime.ai/ 下載你需要的版本。
我下載的是GPU版，在python環境下輸入：
pip install onnxruntime-gpu
二、寫執行程式
1.import
import cv2> import numpy as np import onnxruntime as ort import onnxruntime.backend
2.讀入 onnx model
modelName = 'weights/export.onnx' sess = ort.InferenceSession(modelName)
3.顯示讀入model的相關資訊
print("device", ort.get_device())
print("in count", len(sess.get_inputs()))
print("out count", len(sess.get_outputs()))
input_name = sess.get_inputs()[0].name
print("input name", input_name)
input_shape = sess.get_inputs()[0].shape
print("input shape", input_shape)
input_type = sess.get_inputs()[0].type
print("input type", input_type)
output_name0 = sess.get_outputs()[0].name
print("output0 name", output_name0)
output_shape0 = sess.get_outputs()[0].shape
print("output0 shape", output_shape0)
output_type0 = sess.get_outputs()[0].type
print("output0 type", output_type0)
output_name1 = sess.get_outputs()[1].name
print("output1 name", output_name1)
output_shape1 = sess.get_outputs()[1].shape
print("output1 shape", output_shape1)
output_type1 = sess.get_outputs()[1].type
print("output1 type", output_type1)
執行結果：
device GPU in count 1 out count 2 input name input0 input shape [1, 3, 320, 416] input type tensor(float) output0 name output0 output0 shape [8190, 80] output0 type tensor(float) output1 name output1 output1 shape [8190, 4] output1 type tensor(float)
4.讀入測試用圖片
imgs = []
img_orig = cv2.imread("image.jpeg")
img = cv2.resize(img_orig, (416,320))
imgs.append(img)
imgs = np.stack(imgs, 0)
imgs = imgs[:, :, :, ::-1].transpose(0, 3, 1, 2) # BGR to RGB
imgs = np.ascontiguousarray(imgs, dtype=np.float32) # uint8 to fp16/fp32
imgs /= 255.0 # 0 - 255 to 0.0 - 1.0
print("imgs.shape", imgs.shape)
執行結果：
imgs.shape (1, 3, 320, 416)
5.將圖片輸入至model做inference
res2 = sess.run([output_name0, output_name1], {input_name: imgs})
6.整理inference result
class_scores = res[0][:]
bbox = res[1][:]
# 80個種類中哪一個種類分數最高，紀錄種類的index及分數
class_IDs = np.argmax(class_scores, axis=1)
objects_class_score = []
for i, idx in enumerate(class_IDs):
objects_class_score.append(class_scores[i,idx])
objects_class_score = np.array(objects_class_score)
# 將class_ID、class_score、bbox集合起來
preds = np.stack([class_IDs,objects_class_score], axis=1)
preds = np.concatenate((preds, bbox), axis=1)
# 篩選出 class_score > 0.3 的物體
i = (objects_class_score>0.3) & np.isfinite(preds).all()
preds = preds[i]
print("result shape:", preds.shape)
執行結果：
result shape: (15, 6)
7.show inference result
drawimg = img.copy()
imgH, imgW = drawimg.shape[:2]
# 取得 bbox 的部分，bbox的結構是[center point x, center point y, object width, object height]
for x,y,w,h in preds[:,2:]:
L=int((x-w/2)*imgW)
T=int((y-h/2)*imgH)
R=int(L+w*imgW)
B=int(T+h*imgH)
cv2.rectangle(drawimg, (L,T), (R,B), (255,0,0), 1)
cv2.imshow("img", drawimg)
key = cv2.waitKey(0)
cv2.destroyAllWindows()
執行結果：
三、結論
有抓出物體大致位置代表這個model運作正常，下篇開始講解tensorRT的部分
四、參考資料來源：
https://medium.com/@joehoeller/object-detection-on-thermal-images-f9526237686a
https://github.com/joehoeller/Object-Detection-on-Thermal-Images
留言
張貼留言