1. Torchvison模型推理
import torch
import torchvision
import cv2
import onnx
import numpy as np
import timm
import os
from PIL import Image
from torchvision import transforms
import onnxruntime
from onnxsim import simplify
print(torch.__version__)
print(cv2.__version__)
print(np.__version__)
print(onnx.__version__)
classes = None
class_file = r'F:\Pycharm\PyCharm_Study\Others\ONNX\Opencv-Onnx\classification_classes_ILSVRC2012.txt'
with open(class_file, 'rt') as f:
classes = f.read().rstrip('\n').split('\n')
def init_model(model_name):
if model_name == 'alexnet':
model = torchvision.models.alexnet(pretrained=True)
if model_name == 'densnet':
model = torchvision.models.densenet121(pretrained=True)
if model_name == 'resnet':
model = torchvision.models.resnet50(pretrained=True)
if model_name == 'mobilenet':
model = torchvision.models.mobilenet_v2(pretrained=True)
if model_name == 'squeezenet':
model = torchvision.models.squeezenet1_1(pretrained=True)
if model_name == 'inception':
model = torchvision.models.inception_v3(pretrained=False)
if model_name == 'googlenet':
model = torchvision.models.googlenet(pretrained=True)
if model_name == 'vgg16':
model = torchvision.models.vgg16(pretrained=True)
if model_name == 'vgg19':
model = torchvision.models.vgg19(pretrained=True)
if model_name == 'shufflenet':
model = torchvision.models.shufflenet_v2_x1_0(pretrained=True)
if model_name == 'cspdarknet53':
model = timm.create_model('cspdarknet53', pretrained=True)
if model_name == 'seresnet18':
model = timm.create_model('seresnet18', pretrained=True)
if model_name