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)
|
