Onnx使用预训练的 ResNet18 模型对输入图像进行分类,并将分类结果显示在图像上

本文主要是介绍Onnx使用预训练的 ResNet18 模型对输入图像进行分类,并将分类结果显示在图像上,希望对大家解决编程问题提供一定的参考价值,需要的开发者们随着小编来一起学习吧!

目录

一、整体功能概述

二、函数分析

2.1 resnet() 函数:

2.2 pre_process(img_path) 函数:

2.3 loadOnnx(img_path) 函数:

三、代码执行流程


一、整体功能概述


这段代码实现了一个图像分类系统,使用预训练的 ResNet18 模型对输入图像进行分类,并将分类结果显示在图像上。它包括以下主要步骤:
读取一个包含类别名称和对应编号的文本文件,并将其存储在字典中。
定义了几个函数,包括模型导出函数 resnet()、图像预处理函数 pre_process() 和加载 ONNX 模型进行分类的函数 loadOnnx()。
在主程序中,指定输入图像路径,调用 loadOnnx() 函数对图像进行分类并显示结果。


二、函数分析


2.1 resnet() 函数:


使用 torchvision 中的预训练 ResNet18 模型,并设置为评估模式。
生成一个随机输入张量 x,并将模型导出为 ONNX 格式,保存为 models/resnet18.onnx 文件。

def resnet():model=models.resnet18(weights=ResNet18_Weights.IMAGENET1K_V1)model.eval()x=torch.randn(1,3,224,224)torch.onnx.export(model,x,'models/resnet18.onnx',input_names=['input'],output_names=['output'])


2.2 pre_process(img_path) 函数:


读取输入图像 img_path。
调整图像大小为 224x224。
将图像颜色通道从 BGR 转换为 RGB。
对图像像素值进行归一化处理。
交换图像维度顺序,并增加一个维度。
返回预处理后的图像张量。

def pre_process(img_path):#h w c--->224,224,3#归一化#换轴#增加维度img=cv2.imread(img_path)scale_image=cv2.resize(img,dsize=(224,224))rgb_img=cv2.cvtColor(scale_image,cv2.COLOR_BGR2RGB)rgb_img=rgb_img/255rgb_img=np.transpose(rgb_img,(2,0,1))rgb_img=np.expand_dims(rgb_img,0).astype(np.float32)return rgb_img


2.3 loadOnnx(img_path) 函数:


创建一个 ONNX 推理会话,加载预导出的 ResNet18 ONNX 模型。

调用 pre_process() 函数对输入图像进行预处理。
准备输入数据并进行推理。
获取推理结果中概率最大的类别编号。
根据类别编号从字典中获取对应的类别名称,并进行翻译。
在输入图像上显示分类结果,并展示图像。

def loadOnnx(img_path):session=ort.InferenceSession(r'models\resnet18.onnx',providers=['CPUExecutionProvider'])img=pre_process(img_path)img_back=cv2.imread(img_path)intput_feed={'input':img}session_out=session.run(None,intput_feed)[0]out=np.argmax(session_out,axis=1)[0]res=str(out)# print(dict[res])ans=dict[res].split(',')[1].split(']')[0].strip()ans = translator.translate(ans)cv2.putText(img_back,ans,(100,100),fontFace=1,fontScale=2.0,color=(0,0,255),thickness=3,lineType=cv2.LINE_AA)cv2.imshow('win',img_back)cv2.waitKey(0)cv2.destroyAllWindows()print(ans)

完整代码如下

import cv2
import numpy as np
import torch
from torchvision import models
from torchvision.models import ResNet18_Weights
import onnxruntime as ort
from translate import Translator
translator=Translator(to_lang='Chinese')#翻译成中文
dict={}
with open('类别.txt','r',encoding='utf-8') as f:lines=f.readlines()for line in lines:name=line.split('\t')[0]value=line.split('\t')[1]dict[name]=value
# print(dict)
def resnet():model=models.resnet18(weights=ResNet18_Weights.IMAGENET1K_V1)model.eval()x=torch.randn(1,3,224,224)torch.onnx.export(model,x,'models/resnet18.onnx',input_names=['input'],output_names=['output'])
def pre_process(img_path):#h w c--->224,224,3#归一化#换轴#增加维度img=cv2.imread(img_path)scale_image=cv2.resize(img,dsize=(224,224))rgb_img=cv2.cvtColor(scale_image,cv2.COLOR_BGR2RGB)rgb_img=rgb_img/255rgb_img=np.transpose(rgb_img,(2,0,1))rgb_img=np.expand_dims(rgb_img,0).astype(np.float32)return rgb_img#RGB
def loadOnnx(img_path):session=ort.InferenceSession(r'models\resnet18.onnx',providers=['CPUExecutionProvider'])img=pre_process(img_path)img_back=cv2.imread(img_path)intput_feed={'input':img}session_out=session.run(None,intput_feed)[0]out=np.argmax(session_out,axis=1)[0]res=str(out)# print(dict[res])ans=dict[res].split(',')[1].split(']')[0].strip()ans = translator.translate(ans)cv2.putText(img_back,ans,(100,100),fontFace=1,fontScale=2.0,color=(0,0,255),thickness=3,lineType=cv2.LINE_AA)cv2.imshow('win',img_back)cv2.waitKey(0)cv2.destroyAllWindows()print(ans)pass
if __name__ == '__main__':img_path='dog.png'# resnet()#导出模型loadOnnx(img_path)


三、代码执行流程


在 if __name__ == '__main__': 部分:
定义输入图像路径 img_path。
可以选择调用 resnet() 函数导出模型(注释状态,通常只在第一次运行或模型更新时使用)。
调用 loadOnnx(img_path) 函数对输入图像进行分类和显示结果。

 

 

这篇关于Onnx使用预训练的 ResNet18 模型对输入图像进行分类,并将分类结果显示在图像上的文章就介绍到这儿,希望我们推荐的文章对编程师们有所帮助!



http://www.chinasem.cn/article/1101594

相关文章

Java 中的 @SneakyThrows 注解使用方法(简化异常处理的利与弊)

《Java中的@SneakyThrows注解使用方法(简化异常处理的利与弊)》为了简化异常处理,Lombok提供了一个强大的注解@SneakyThrows,本文将详细介绍@SneakyThro... 目录1. @SneakyThrows 简介 1.1 什么是 Lombok?2. @SneakyThrows

使用Python和Pyecharts创建交互式地图

《使用Python和Pyecharts创建交互式地图》在数据可视化领域,创建交互式地图是一种强大的方式,可以使受众能够以引人入胜且信息丰富的方式探索地理数据,下面我们看看如何使用Python和Pyec... 目录简介Pyecharts 简介创建上海地图代码说明运行结果总结简介在数据可视化领域,创建交互式地

Java Stream流使用案例深入详解

《JavaStream流使用案例深入详解》:本文主要介绍JavaStream流使用案例详解,本文通过实例代码给大家介绍的非常详细,对大家的学习或工作具有一定的参考借鉴价值,需要的朋友参考下吧... 目录前言1. Lambda1.1 语法1.2 没参数只有一条语句或者多条语句1.3 一个参数只有一条语句或者多

利用python实现对excel文件进行加密

《利用python实现对excel文件进行加密》由于文件内容的私密性,需要对Excel文件进行加密,保护文件以免给第三方看到,本文将以Python语言为例,和大家讲讲如何对Excel文件进行加密,感兴... 目录前言方法一:使用pywin32库(仅限Windows)方法二:使用msoffcrypto-too

Java Spring 中 @PostConstruct 注解使用原理及常见场景

《JavaSpring中@PostConstruct注解使用原理及常见场景》在JavaSpring中,@PostConstruct注解是一个非常实用的功能,它允许开发者在Spring容器完全初... 目录一、@PostConstruct 注解概述二、@PostConstruct 注解的基本使用2.1 基本代

C#使用StackExchange.Redis实现分布式锁的两种方式介绍

《C#使用StackExchange.Redis实现分布式锁的两种方式介绍》分布式锁在集群的架构中发挥着重要的作用,:本文主要介绍C#使用StackExchange.Redis实现分布式锁的... 目录自定义分布式锁获取锁释放锁自动续期StackExchange.Redis分布式锁获取锁释放锁自动续期分布式

springboot使用Scheduling实现动态增删启停定时任务教程

《springboot使用Scheduling实现动态增删启停定时任务教程》:本文主要介绍springboot使用Scheduling实现动态增删启停定时任务教程,具有很好的参考价值,希望对大家有... 目录1、配置定时任务需要的线程池2、创建ScheduledFuture的包装类3、注册定时任务,增加、删

使用Python实现矢量路径的压缩、解压与可视化

《使用Python实现矢量路径的压缩、解压与可视化》在图形设计和Web开发中,矢量路径数据的高效存储与传输至关重要,本文将通过一个Python示例,展示如何将复杂的矢量路径命令序列压缩为JSON格式,... 目录引言核心功能概述1. 路径命令解析2. 路径数据压缩3. 路径数据解压4. 可视化代码实现详解1

Pandas透视表(Pivot Table)的具体使用

《Pandas透视表(PivotTable)的具体使用》透视表用于在数据分析和处理过程中进行数据重塑和汇总,本文就来介绍一下Pandas透视表(PivotTable)的具体使用,感兴趣的可以了解一下... 目录前言什么是透视表?使用步骤1. 引入必要的库2. 读取数据3. 创建透视表4. 查看透视表总结前言

Python 交互式可视化的利器Bokeh的使用

《Python交互式可视化的利器Bokeh的使用》Bokeh是一个专注于Web端交互式数据可视化的Python库,本文主要介绍了Python交互式可视化的利器Bokeh的使用,具有一定的参考价值,感... 目录1. Bokeh 简介1.1 为什么选择 Bokeh1.2 安装与环境配置2. Bokeh 基础2