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 pltdef predict_class(model_path, image_path, class_index_path, top_n=1):# Load the ONNX model and create an ONNX Runtime inference sessionort_session = onnxruntime.InferenceSession(model_path, providers=['CUDAExecutionProvider'])# Load the test image and apply transformationsimg_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 inferenceort_inputs = {'input': input_tensor}pred_logits = ort_session.run(['output'], ort_inputs)[0]pred_logits = torch.tensor(pred_logits)# Apply softmax to get class probabilitiespred_softmax = F.softmax(pred_logits, dim=1)# Load the class index mappingdf = 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 scorestop_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 namespredicted_classes = [idx_to_labels[class_id] for class_id in pred_ids]# Create a list of class names and their corresponding confidence scoresresults = []for i in range(top_n):class_name = predicted_classes[i]confidence = confs[i] * 100results.append((class_name, confidence))return results,img_pilif __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 wantpredictions,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 inputplt.rcParams['font.sans-serif'] = 'SimHei'  # 黑体plt.figure(figsize=(6,6))plt.imshow(img)plt.axis('off')# add the predicted class and confidence as a titleclass_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 pltdef 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]*100plt.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=1input_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)