134 lines
5.1 KiB
Python
134 lines
5.1 KiB
Python
from loguru import logger
|
||
from models.experimental import attempt_load
|
||
import tensorrt as trt
|
||
import torch
|
||
import sys
|
||
from DrGraph.util.segutils.trtUtils import segPreProcess_image,segTrtForward,segPreProcess_image_torch
|
||
from DrGraph.util.segutils.model_stages import BiSeNet_STDC
|
||
import time,cv2
|
||
import numpy as np
|
||
|
||
from DrGraph.util.drHelper import *
|
||
|
||
class stdcModel(object):
|
||
def __init__(self, weights=None,
|
||
par={'modelSize':(640,360),'dynamic':False,'nclass':2,'predResize':True,'mean':(0.485, 0.456, 0.406),'std' :(0.229, 0.224, 0.225),'numpy':False, 'RGB_convert_first':True}
|
||
):
|
||
|
||
self.par = par
|
||
self.device = 'cuda:0'
|
||
self.half =True
|
||
if 'dynamic' not in par.keys():
|
||
self.dynamic=False
|
||
else: self.dynamic=par['dynamic']
|
||
|
||
if weights.endswith('.engine'):
|
||
self. infer_type ='trt'
|
||
elif weights.endswith('.pth') or weights.endswith('.pt') :
|
||
self. infer_type ='pth'
|
||
else:
|
||
logger.error(f'{weights}: no registered inference type, exit')
|
||
sys.exit(0)
|
||
|
||
if self.infer_type=='trt':
|
||
if self.dynamic :
|
||
logger.error('STDC动态模型不能采用trt格式')
|
||
|
||
trt_logger = trt.Logger(trt.Logger.ERROR)
|
||
with open(weights, "rb") as f, trt.Runtime(trt_logger) as runtime:
|
||
self.model=runtime.deserialize_cuda_engine(f.read())# 输入trt本地文件，返回ICudaEngine对象
|
||
elif self.infer_type=='pth':
|
||
if self.dynamic: modelSize=None
|
||
else: modelSize=( self.par['modelSize'][1], self.par['modelSize'][0] )
|
||
self.model = BiSeNet_STDC(backbone='STDCNet813', n_classes=par['seg_nclass'],
|
||
use_boundary_2=False, use_boundary_4=False,
|
||
use_boundary_8=True, use_boundary_16=False,
|
||
use_conv_last=False,
|
||
modelSize = modelSize
|
||
)
|
||
|
||
self.model.load_state_dict(torch.load(weights, map_location=torch.device(self.device) ))
|
||
self.model= self.model.to(self.device)
|
||
logger.info('加载 stdcModel 模型:',weights,' 类型:',self.infer_type)
|
||
def preprocess_image(self,image):
|
||
image = self.RB_convert(image)
|
||
re_size = (640, 360)
|
||
if 'modelSize' in self.par:
|
||
if self.dynamic:
|
||
H,W=image.shape[0:2];
|
||
yscale = self.par['modelSize'][1]/H
|
||
xscale = self.par['modelSize'][0]/W
|
||
dscale = min(yscale,xscale)
|
||
re_size = ( int((dscale*W)//4*4), int( (dscale*H)//4*4 ) )
|
||
else: re_size = self.par['modelSize']
|
||
else:
|
||
logger.warning('modelSize not in par, use default size(640, 360)')
|
||
#print('####line 58:,', re_size,image.shape)
|
||
image = cv2.resize(image, re_size, interpolation=cv2.INTER_LINEAR)
|
||
|
||
image = image.astype(np.float32)
|
||
image /= 255.0
|
||
|
||
if 'mean' not in self.par:
|
||
self.par['mean'] = (0.485, 0.456, 0.406)
|
||
logger.warning('mean not in par, use default mean(0.485, 0.456, 0.406)')
|
||
if 'std' not in self.par:
|
||
self.par['std'] = (0.229, 0.224, 0.225)
|
||
logger.warning('std not in par, use default std(0.229, 0.224, 0.225)')
|
||
|
||
image[:, :, 0] -= self.par['mean'][0]
|
||
image[:, :, 1] -= self.par['mean'][1]
|
||
image[:, :, 2] -= self.par['mean'][2]
|
||
|
||
image[:, :, 0] /= self.par['std'][0]
|
||
image[:, :, 1] /= self.par['std'][1]
|
||
image[:, :, 2] /= self.par['std'][2]
|
||
|
||
image = np.transpose(image, (2, 0, 1))
|
||
image = torch.from_numpy(image).float()
|
||
image = image.unsqueeze(0)
|
||
if self.device != 'cpu':
|
||
image = image.to(self.device)
|
||
|
||
return image
|
||
|
||
def RB_convert(self,image):
|
||
image_c = image.copy()
|
||
image_c[:,:,0] = image[:,:,2]
|
||
image_c[:,:,2] = image[:,:,0]
|
||
return image_c
|
||
|
||
def eval(self,image):
|
||
time0 = time.time()
|
||
imageH, imageW, _ = image.shape
|
||
|
||
img = self.preprocess_image(image)
|
||
time1 = time.time()
|
||
|
||
if self.infer_type=='trt':
|
||
pred=segTrtForward(self.model,[img])
|
||
|
||
elif self.infer_type=='pth':
|
||
self.model.eval()
|
||
with torch.no_grad():
|
||
pred = self.model(img)
|
||
|
||
time2 = time.time()
|
||
pred=torch.argmax(pred,dim=1).cpu().numpy()[0]
|
||
time3 = time.time()
|
||
pred = cv2.resize(pred.astype(np.uint8),(imageW,imageH))
|
||
time4 = time.time()
|
||
outstr= 'pre-precess:%.1f ,infer:%.1f ,post-cpu-argmax:%.1f ,post-resize:%.1f, total:%.1f \n '%( \
|
||
timeHelper.deltaTime_MS(time1,time0),\
|
||
timeHelper.deltaTime_MS(time2,time1),\
|
||
timeHelper.deltaTime_MS(time3,time2),\
|
||
timeHelper.deltaTime_MS(time4,time3),\
|
||
timeHelper.deltaTime_MS(time4,time0) )
|
||
|
||
return pred,outstr
|
||
|
||
|
||
|
||
|
||
|
||
|