65 lines
2.5 KiB
Python
65 lines
2.5 KiB
Python
import torch
|
||
import numpy as np
|
||
import cv2
|
||
import time
|
||
import os
|
||
import sys
|
||
sys.path.extend(['../AIlib2/obbUtils'])
|
||
import matplotlib.pyplot as plt
|
||
import func_utils
|
||
import time
|
||
import torchvision.transforms as transforms
|
||
from obbmodels import ctrbox_net
|
||
import decoder
|
||
import tensorrt as trt
|
||
import onnx
|
||
import onnxruntime as ort
|
||
|
||
def load_model_decoder_OBB(par={'down_ratio':4,'num_classes':15,'weights':'weights_dota/obb.pth'}):
|
||
weights=par['weights']
|
||
heads = par['heads']
|
||
heads['hm']=par['num_classes']
|
||
par['heads']=heads
|
||
if weights.endswith('.pth') or weights.endswith('.pt'):
|
||
resume=par['weights']
|
||
down_ratio = par['down_ratio']
|
||
model = ctrbox_net.CTRBOX(heads=heads,
|
||
pretrained=True,
|
||
down_ratio=down_ratio,
|
||
final_kernel=1,
|
||
head_conv=256)
|
||
|
||
|
||
checkpoint = torch.load(resume, map_location=lambda storage, loc: storage)
|
||
print('loaded weights from {}, epoch {}'.format(resume, checkpoint['epoch']))
|
||
state_dict_ = checkpoint['model_state_dict']
|
||
model.load_state_dict(state_dict_, strict=True)
|
||
model.eval()
|
||
model = model.to(par['device'])
|
||
model = model.half() if par['half'] else model
|
||
par['saveType']='pth'
|
||
elif weights.endswith('.onnx'):
|
||
onnx_model = onnx.load(weights)
|
||
onnx.checker.check_model(onnx_model)
|
||
# 设置模型session以及输入信息
|
||
sess = ort.InferenceSession(str(weights),providers= ort.get_available_providers())
|
||
print('len():',len( sess.get_inputs() ))
|
||
input_name = sess.get_inputs()[0].name
|
||
model = {'sess':sess,'input_name':input_name}
|
||
par['saveType']='onnx'
|
||
elif weights.endswith('.engine'):
|
||
logger = trt.Logger(trt.Logger.ERROR)
|
||
with open(weights, "rb") as f, trt.Runtime(logger) as runtime:
|
||
model = runtime.deserialize_cuda_engine(f.read())# 输入trt本地文件，返回ICudaEngine对象
|
||
print('#####load TRT file:',weights,'success #####')
|
||
par['saveType']='trt'
|
||
|
||
|
||
decoder2 = decoder.DecDecoder(K=par['K'],
|
||
conf_thresh=par['conf_thresh'],
|
||
num_classes=par['num_classes'])
|
||
|
||
|
||
|
||
return model, decoder2
|