46 lines
1.9 KiB
Python
46 lines
1.9 KiB
Python
|
from DMPRUtils.DMPR_process import DMPR_process
|
|||
|
|
import tensorrt as trt
|
|||
|
|
import sys,os
|
|||
|
|
#from DMPRUtils.model.detector import DirectionalPointDetector
|
|||
|
|
from DMPRUtils.yolo_net import Model
|
|||
|
|
import torch
|
|||
|
|
|
|||
|
|
class DMPRModel(object):
|
|||
|
|
def __init__(self, weights=None,
|
|||
|
|
par={'depth_factor':32,'NUM_FEATURE_MAP_CHANNEL':6,'dmpr_thresh':0.3, 'dmprimg_size':640}
|
|||
|
|
):
|
|||
|
|
|
|||
|
|
self.par = par
|
|||
|
|
self.device = 'cuda:0'
|
|||
|
|
self.half =True
|
|||
|
|
|
|||
|
|
if weights.endswith('.engine'):
|
|||
|
|
self.infer_type ='trt'
|
|||
|
|
elif weights.endswith('.pth') or weights.endswith('.pt') :
|
|||
|
|
self.infer_type ='pth'
|
|||
|
|
else:
|
|||
|
|
print('#########ERROR:',weights,': no registered inference type, exit')
|
|||
|
|
sys.exit(0)
|
|||
|
|
|
|||
|
|
if self.infer_type=='trt':
|
|||
|
|
logger = trt.Logger(trt.Logger.ERROR)
|
|||
|
|
with open(weights, "rb") as f, trt.Runtime(logger) as runtime:
|
|||
|
|
self.model=runtime.deserialize_cuda_engine(f.read())# 输入trt本地文件，返回ICudaEngine对象
|
|||
|
|
elif self.infer_type=='pth':
|
|||
|
|
#self.model = DirectionalPointDetector(3, self.par['depth_factor'], self.par['NUM_FEATURE_MAP_CHANNEL']).to(self.device)
|
|||
|
|
confUrl = os.path.join( os.path.dirname(__file__),'DMPRUtils','config','yolov5s.yaml' )
|
|||
|
|
self.model = Model(confUrl, ch=3).to(self.device)
|
|||
|
|
self.model.load_state_dict(torch.load(weights))
|
|||
|
|
print('#######load pt model:%s success '%(weights))
|
|||
|
|
self.par['modelType']=self.infer_type
|
|||
|
|
print('#########加载模型:',weights,' 类型:',self.infer_type)
|
|||
|
|
def eval(self,image):
|
|||
|
|
det,timeInfos = DMPR_process(image, self.model, self.device, self.par)
|
|||
|
|
det = det.cpu().detach().numpy()
|
|||
|
|
return det,timeInfos
|
|||
|
|
|
|||
|
|
def get_ms(self,t1,t0):
|
|||
|
|
return (t1-t0)*1000.0
|
|||
|
|
|
|||
|
|
|