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)