|
-
- 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)
-
|
