65 lines
1.9 KiB
Python
65 lines
1.9 KiB
Python
|
||
import sys
|
||
#sys.path.extend(['..','../AIlib2' ])
|
||
from ocrTrt import toONNX,ONNXtoTrt
|
||
from collections import OrderedDict
|
||
import torch
|
||
import argparse
|
||
from load_obb_model import load_model_decoder_OBB
|
||
|
||
|
||
def getModel(opt):
|
||
|
||
|
||
###倾斜框（OBB）的ship目标检测
|
||
par={
|
||
'model_size':(608,608), #width,height
|
||
'K':100, #Maximum of objects'
|
||
'conf_thresh':0.18,##Confidence threshold, 0.1 for general evaluation
|
||
'device':"cuda:0",
|
||
|
||
'down_ratio':4,'num_classes':15,
|
||
'weights':opt.weights,
|
||
'dataset':'dota',
|
||
|
||
'test_dir': 'images/ship/',
|
||
'result_dir': 'images/results',
|
||
'half': False,
|
||
'mean':(0.5, 0.5, 0.5),
|
||
'std':(1, 1, 1),
|
||
'category':['0','1','2','3','4','5','6','7','8','9','10','11','12','13','boat'],
|
||
'model_size':(608,608),##width,height
|
||
'decoder':None,
|
||
'test_flag':True,
|
||
'heads': {'hm': None,'wh': 10,'reg': 2,'cls_theta': 1},
|
||
|
||
|
||
}
|
||
|
||
####加载模型
|
||
model,decoder2=load_model_decoder_OBB(par)
|
||
par['decoder']=decoder2
|
||
model = model.to(par['device'])
|
||
return model
|
||
|
||
|
||
if __name__=='__main__':
|
||
|
||
parser = argparse.ArgumentParser()
|
||
parser.add_argument('--weights', type=str, default='/mnt/thsw2/DSP2/weights/ship2/obb_608X608.pth', help='model path(s)')
|
||
parser.add_argument('--mWidth', type=int, default=608, help='segmodel mWdith')
|
||
parser.add_argument('--mHeight', type=int, default=608, help='segmodel mHeight')
|
||
opt = parser.parse_args()
|
||
|
||
pthmodel = getModel(opt)
|
||
|
||
###转换TRT模型
|
||
onnxFile=opt.weights.replace('.pth','.onnx')
|
||
trtFile=opt.weights.replace('.pth','.engine')
|
||
|
||
print('#'*20, ' begin to toONNX')
|
||
toONNX(pthmodel,onnxFile,inputShape=(1,3,opt.mHeight, opt.mWidth),device='cuda:0')
|
||
print('#'*20, ' begin to TRT')
|
||
ONNXtoTrt(onnxFile,trtFile,half=False)
|
||
|