import cv2
import numpy as np
import onnxruntime
from scipy.special import softmax

# 加载ONNX模型
session = onnxruntime.InferenceSession("model/classify/best.onnx")

# 读取图片
img = cv2.imread("D:\\temp\\15.jpg")

# 从res.json中读取类别
with open("res1-2.json", "r") as f:
    classes = eval(f.read())

# 预处理
blob = cv2.dnn.blobFromImage(img, 1/255.0, (640, 640), swapRB=True, crop=False)

# 模型推理
outputs = session.run(None, {session.get_inputs()[0].name: blob})

# print(outputs)
# 应用softmax函数
probabilities = outputs[0]

# 找到最大概率的类别
predicted_class = np.argmax(probabilities, axis=1)[0]
max_probability = np.max(probabilities, axis=1)[0]

# 找到概率较高的前十个类别
top_ten_classes = np.argsort(probabilities, axis=1)[0][-5:]

# 输出前十个类别
print("Top 10 Classes:")
for i in top_ten_classes:
    print(f"{classes[i]}: {probabilities[0][i]}")



# 后处理
# for detection in outputs[0][0]:
#     confidence = detection[4]
#     if confidence > 0.5:
#         class_id = int(detection[5])
#         x, y, w, h = detection[:4]
#         cv2.rectangle(img, (int(x), int(y)), (int(x+w), int(y+h)), (0, 255, 0), 2)
#
# # 显示结果
# cv2.imshow("YOLOv8 Detection", img)
# cv2.waitKey(0)