分类模型onnx推理，并生成混淆矩阵

废话不多说直接上代码

import onnxruntime
import numpy as np
import os
import cv2
import argparse
import time
import shutil
import matplotlib.pyplot as plt
from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay

labels = ["0", "1", "2", "3", "4", "5", "6", "7"]

def sigmoid(x):
    """Sigmoid function for a scalar or NumPy array."""
    return 1 / (1 + np.exp(-x))

def getFileList(dir, Filelist, ext=None):
    """
    获取文件夹及其子文件夹中文件列表
    输入 dir：文件夹根目录
    输入 ext: 扩展名
    返回： 文件路径列表
    """
    newDir = dir
    if os.path.isfile(dir):
        if ext is None:
            Filelist.append(dir)
        else:
            if ext in dir[-3:]:
                Filelist.append(dir)

    elif os.path.isdir(dir):
        for s in os.listdir(dir):
            newDir = os.path.join(dir, s)
            getFileList(newDir, Filelist, ext)

    return Filelist

def read_image(image_path, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]):
    src = cv2.imdecode(np.fromfile(image_path, dtype=np.uint8), cv2.IMREAD_COLOR)
    image = cv2.cvtColor(src, cv2.COLOR_BGR2RGB)
    image = cv2.resize(image, (224, 224))
    image = image.astype(np.float32)
    image = image / 255.0
    image = image.transpose(2, 0, 1)
    mean = np.array(mean, dtype=np.float32).reshape((3,1,1))
    std = np.array(std, dtype=np.float32).reshape((3,1,1))
    # 对图像进行归一化
    normalized_image = (image - mean) / std
    normalized_image = np.expand_dims(normalized_image, axis=0)
    return normalized_image, src

def load_onnx_model(model_path):
    providers = ['CUDAExecutionProvider']  # 使用 GPU
    # providers = ['CPUExecutionProvider']
    options = onnxruntime.SessionOptions()
    options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
    options.execution_mode = onnxruntime.ExecutionMode.ORT_SEQUENTIAL
    session = onnxruntime.InferenceSession(model_path, options, providers=providers)
    print("ONNX模型已成功加载。")
    return session

def main(image_path, session):
    image, src = read_image(image_path)
    input_name = session.get_inputs()[0].name
    output_name = session.get_outputs()[0].name
    pred = session.run([output_name], {input_name: image})[0]
    pred = np.squeeze(pred)
    pred = [sigmoid(x) for x in pred]
    return pred.index(max(pred)), max(pred), labels[pred.index(max(pred))]

def plot_confusion_matrix(y_true, y_pred, labels):
    """
    绘制混淆矩阵
    输入 y_true: 真实标签
    输入 y_pred: 预测标签
    输入 labels: 标签名称
    """
    cm = confusion_matrix(y_true, y_pred)
    disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=labels)
    disp.plot(cmap=plt.cm.Blues)
    plt.title('Confusion Matrix')
    plt.show()

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--images_path', type=str, default="/home/workspace/temp/test_result/kaideng", help='images_path')
    parser.add_argument('--model_path', type=str, default="/home/workspace/temp/47-val-Loss-0.0203-Acc-0.9942.onnx", help='model_path')
    args = parser.parse_args()
    img_list = []
    img_list = getFileList(args.images_path, img_list)
    count = 0
    session = load_onnx_model(args.model_path)
    start = time.time()
    y_true = []
    y_pred = []
    count_time = 0
    for img in img_list:
        #true_label = int(img.split('/')[-2].split('-')[0])
        true_label = img.split('/')[-3] #这一句代码是获取图像类别文件夹的名称，具体索引需要修改
        start_1 = time.time()
        predicted_index, score, label = main(img, session)
        count_time += time.time() - start_1
        y_true.append(true_label)
        #y_pred.append(predicted_index)
        y_pred.append(label)
        if label == true_label:
            count += 1
        # else:
        #     dst_path = img.replace('test', 'test_out')
        #     dst_dir = os.path.dirname(dst_path)
        #     if not os.path.exists(dst_dir):
        #         os.makedirs(dst_dir)
        #     shutil.copy(img, dst_path.replace('.jpg', "-" + label + '.jpg'))

    accuracy = count / len(img_list) * 100
    print(f"Accuracy: {accuracy:.2f}%")
    print(f"Correct predictions: {count}, Total images: {len(img_list)}")
    print(f"Time taken: {time.time() - start:.2f} seconds")
    print("推理", len(img_list), "张图像用时", count_time)
    # 绘制混淆矩阵
    plot_confusion_matrix(y_true, y_pred, labels)

1.要确保图像类别文件夹的名称和labels列表相对应，不然无法生成混淆矩阵

2.read_image函数中的预处理方式要和训练时的一致，请根据个人需要修改代码