YOLOv10導(dǎo)出onnx模型并預(yù)測(cè)

相比YOLOv8,v10不再使用NMS,預(yù)測(cè)過(guò)程也變得簡(jiǎn)單。

import onnxruntime as ort
import  matplotlib.pyplot as plt
import numpy as np
import cv2

onnx_model = ort.InferenceSession("model_path",providers=['CPUExecutionProvider'])

def ratioresize(im,  color=255):
    shape = im.shape[:2]  
    new_h, new_w = 640,640
    padded_img = np.ones((new_h, new_w, 3), dtype=np.uint8) * color

    r = min(new_h / shape[0], new_w / shape[1])
    new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))

    if shape[::-1] != new_unpad:
        im = cv2.resize(im, new_unpad, interpolation=cv2.INTER_LINEAR)
    
    padded_img[: new_unpad[1], : new_unpad[0]] = im
    padded_img = np.ascontiguousarray(padded_img)
    return padded_img, 1 / r

模型預(yù)測(cè),結(jié)果形狀為[1,300,6],每一行前4位為候選框坐標(biāo)xyxy,第五位為conf,第六位為類(lèi)別,結(jié)果已經(jīng)按照conf由大到小排列。

image = cv2.imread("image_path")
image = cv2.cvtColor(image,cv2.COLOR_BGR2RGB)
in_img,scale = ratioresize(image)
in_img = in_img/255.0 - 0.5
in_img = np.transpose(in_img,(2,0,1))
in_img = np.expand_dims(in_img,axis=0).astype(np.float32)

output = onnx_model.run(None,{'images':in_img})
result = output[0][0]
#設(shè)置閾值
boxes = result[result[:,4] > 0.25]
boxes = (boxes * scale).astype(np.int32) 
for box in boxes:
    cv2.rectangle(image ,(box[0],box[1]),(box[2],box[3]),(255,0,0),3)

plt.imshow(image)

最后編輯于
?著作權(quán)歸作者所有,轉(zhuǎn)載或內(nèi)容合作請(qǐng)聯(lián)系作者
【社區(qū)內(nèi)容提示】社區(qū)部分內(nèi)容疑似由AI輔助生成,瀏覽時(shí)請(qǐng)結(jié)合常識(shí)與多方信息審慎甄別。
平臺(tái)聲明:文章內(nèi)容(如有圖片或視頻亦包括在內(nèi))由作者上傳并發(fā)布,文章內(nèi)容僅代表作者本人觀點(diǎn),簡(jiǎn)書(shū)系信息發(fā)布平臺(tái),僅提供信息存儲(chǔ)服務(wù)。

相關(guān)閱讀更多精彩內(nèi)容

友情鏈接更多精彩內(nèi)容