ONNX runtime本地终端部署

1、class_index.csv文件：

ID,English,Chinese
0,A,你
1,B,我
2,C,他
3,D,她

2、classification.onnx
3、单张图像处理代码如下：

import onnxruntime
import torch
import torch.nn.functional as F
import pandas as pd
from PIL import Image
from torchvision import transforms
import matplotlib.pyplot as plt

def predict_class(model_path, image_path, class_index_path, top_n=1):
    # Load the ONNX model and create an ONNX Runtime inference session
    ort_session = onnxruntime.InferenceSession(model_path, providers=['CUDAExecutionProvider'])

    # Load the test image and apply transformations
    img_pil = Image.open(image_path).convert('RGB')
    test_transform = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406],
                             [0.229, 0.224, 0.225])
    ])
    input_img = test_transform(img_pil)
    input_tensor = input_img.unsqueeze(0).numpy()

    # Perform inference
    ort_inputs = {'input': input_tensor}
    pred_logits = ort_session.run(['output'], ort_inputs)[0]
    pred_logits = torch.tensor(pred_logits)

    # Apply softmax to get class probabilities
    pred_softmax = F.softmax(pred_logits, dim=1)

    # Load the class index mapping
    df = pd.read_csv(class_index_path)
    idx_to_labels = {row['ID']: row['English'] for idx, row in df.iterrows()}
    # idx_to_labels = {row['ID']: row['Chinese'] for idx, row in df.iterrows()}

    # Get the top predicted classes and their confidence scores
    top_n_results = torch.topk(pred_softmax, top_n)
    pred_ids = top_n_results.indices.numpy()[0]
    confs = top_n_results.values.numpy()[0]

    # Translate class IDs to class names
    predicted_classes = [idx_to_labels[class_id] for class_id in pred_ids]

    # Create a list of class names and their corresponding confidence scores
    results = []
    for i in range(top_n):
        class_name = predicted_classes[i]
        confidence = confs[i] * 100
        results.append((class_name, confidence))
    return results,img_pil

if __name__ == '__main__':
    model_path = 'classification.onnx'
    image_path = 'E:/Python_Project/classification/21t1Gdxsagittal0143.png'
    class_index_path = 'class_index.csv'

    top_n = 1  # Adjust the number of top predictions you want
    predictions,img = predict_class(model_path, image_path, class_index_path, top_n)

    for class_name, confidence in predictions:
        text = class_name + ': {:.3f}%'.format(confidence)
        print(text)

    # Display the image input
    plt.rcParams['font.sans-serif'] = 'SimHei'  # 黑体

    plt.figure(figsize=(6,6))
    plt.imshow(img)
    plt.axis('off')

    # add the predicted class and confidence as a title
    class_name, confidence = predictions[0]
    title_text = f'Predicted Class: {class_name}\nAccuracy: {confidence:.3f}%'
    plt.title(title_text)
    plt.show()

4、批量图像处理代码如下：

import onnxruntime
import pandas
import torch
from PIL import Image
import os
from torchvision import transforms
import torch.nn.functional as F
import matplotlib.pyplot as plt

def batch_prediction(model_path,class_index_path,top_n,input_folder,output_folder):

    # 列出输出文件夹所有图片
    input_files=os.listdir(input_folder)
    # print(input_files)

    # 针对每个文件进行处理
    for input_file in input_files:
        # 构建一个完整的路径
        input_file_path=os.path.join(input_folder,input_file)
        # 打开图像并转换成RGB格式
        img_pil=Image.open(input_file_path).convert('RGB')
        # print(image)
        # 图像预处理
        test_transform = transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406],
                                 [0.229, 0.224, 0.225])
        ])

        input_img=test_transform(img_pil)
        # 增加维度
        input_tensor=input_img.unsqueeze(0).numpy()
        # print(input_tensor.shape)

        # 加载ONNX模型并创建onnx_runtime推理
        ort_session = onnxruntime.InferenceSession(model_path, providers=['CUDAExecutionProvider'])
        # print(ort_session)

        # 平台推理
        ort_inputs={'input':input_tensor}
        pre_logits=ort_session.run(['output'],ort_inputs)[0]
        pre_logits=torch.Tensor(pre_logits)
        # print(pre_logits)

        # 应用softmax去做类别预测
        pred_softmax=F.softmax(pre_logits,dim=1)# dim=1 按行进行归一化
        # print(pred_softmax)

        # 加载类别索引
        df=pandas.read_csv(class_index_path)
        # print(df)
        idx_to_labels={row['ID'] : row['English'] for idx, row in df.iterrows()}

        # 去得到准确率和得分
        top_n_results=torch.topk(pred_softmax,top_n)
        # print(top_n_results)
        pred_ids=top_n_results.indices.numpy()[0]# 得到id
        # print(pred_ids)
        confs=top_n_results.values.numpy()[0]# 得到分数 numpy格式
        # print(type(confs))

        # 预测类别
        predicted_class=[idx_to_labels[class_id] for class_id in pred_ids]#列表格式
        # print(predicted_class)

        # 输出类别和相应的得分
        # print(predicted_class,confs)

        for i in range(top_n):
            class_name=predicted_class[i]
            confidence=confs[i]*100
            plt.rcParams['font.sans-serif'] = 'SimHei'  # 黑体中文字体

            plt.imshow(img_pil)#
            plt.axis('off')#
            title_text = f'Predicted Class: {class_name}\nAccuracy: {confidence:.3f}%'
            plt.title(title_text)
            print(title_text)

            # 确保输出文件夹存在，如果不存在则创建
            if not os.path.exists(output_folder):
                os.makedirs(output_folder)

            # 构建输出文件夹的完整路径并保存绘制的图像
            plot_output_file = os.path.join(output_folder,input_file)
            plt.savefig(plot_output_file, bbox_inches='tight', pad_inches=0.1)  # 保存绘制的图像
            plt.close()  # 关闭当前绘图以便处理下一个图像

if __name__ == '__main__':
    model_path='classification.onnx'
    class_index_path='class_index.csv'
    top_n=1
    input_folder="C:/Users/XUB/Desktop/A"
    output_folder="C:/Users/XUB/Desktop/B"
    batch_prediction(model_path,class_index_path,top_n,input_folder,output_folder)