|
import os
|
import cv2
|
import numpy as np
|
import onnxruntime
|
|
if __name__ == '__main__':
|
img = cv2.imread("data/images/2/20250114200010.jpg")
|
identify_session = onnxruntime.InferenceSession("model/herb_identify.onnx")
|
herbs = eval(identify_session.get_modelmeta().custom_metadata_map['names'])
|
# 预处理
|
herb_blob = cv2.dnn.blobFromImage(img, 1 / 255.0, (640, 640), swapRB=True, crop=False)
|
# 模型推理
|
herb_outputs = identify_session.run(None, {identify_session.get_inputs()[0].name: herb_blob})
|
herb_probabilities = herb_outputs[0]
|
|
top_five_classes = np.argsort(herb_probabilities, axis=1)[0][-5:][::-1]
|
name = ""
|
for i, class_id in enumerate(top_five_classes):
|
# 保留两位小数
|
probability = round(herb_probabilities[0][class_id], 2)
|
|
herb_class = herbs[class_id]
|
name = name + " " + herb_class + ":" + str(probability)
|
# 将识别结果打印到图片上
|
cv2.putText(img, herb_class + ":" + str(probability), (10, (i+1)*40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
|
|
# 显示图片
|
cv2.imshow("Output", img)
|
print(name)
|
# 等待
|
cv2.waitKey(0)
|
