584 lines
22 KiB
Python
584 lines
22 KiB
Python
import torch
|
||
import argparse
|
||
import sys,os
|
||
sys.path.extend(['segutils'])
|
||
from core.models.bisenet import BiSeNet
|
||
from model_stages import BiSeNet_STDC
|
||
from torchvision import transforms
|
||
import cv2,glob
|
||
import numpy as np
|
||
import matplotlib.pyplot as plt
|
||
import time
|
||
from pathlib import Path
|
||
from trtUtils import TRTModule,segTrtForward,segtrtEval,segPreProcess_image,get_ms
|
||
from concurrent.futures import ThreadPoolExecutor
|
||
import tensorrt as trt
|
||
from copy import deepcopy
|
||
import onnx
|
||
import numpy as np
|
||
import onnxruntime as ort
|
||
import cv2
|
||
#import pycuda.driver as cuda
|
||
class SegModel_BiSeNet(object):
|
||
def __init__(self, nclass=2,weights=None,modelsize=512,device='cuda:0'):
|
||
#self.args = args
|
||
self.model = BiSeNet(nclass)
|
||
checkpoint = torch.load(weights)
|
||
if isinstance(modelsize,list) or isinstance(modelsize,tuple):
|
||
self.modelsize = modelsize
|
||
else: self.modelsize = (modelsize,modelsize)
|
||
self.model.load_state_dict(checkpoint['model'])
|
||
self.device = device
|
||
self.model= self.model.to(self.device)
|
||
'''self.composed_transforms = transforms.Compose([
|
||
|
||
transforms.Normalize(mean=(0.335, 0.358, 0.332), std=(0.141, 0.138, 0.143)),
|
||
transforms.ToTensor()]) '''
|
||
self.mean = (0.335, 0.358, 0.332)
|
||
self.std = (0.141, 0.138, 0.143)
|
||
def eval(self,image):
|
||
time0 = time.time()
|
||
imageH,imageW,imageC = image.shape
|
||
image = self.preprocess_image(image)
|
||
time1 = time.time()
|
||
self.model.eval()
|
||
image = image.to(self.device)
|
||
with torch.no_grad():
|
||
output = self.model(image)
|
||
|
||
time2 = time.time()
|
||
pred = output.data.cpu().numpy()
|
||
pred = np.argmax(pred, axis=1)[0]#得到每行
|
||
time3 = time.time()
|
||
pred = cv2.resize(pred.astype(np.uint8),(imageW,imageH))
|
||
time4 = time.time()
|
||
outstr= 'pre-precess:%.1f ,infer:%.1f ,post-precess:%.1f ,post-resize:%.1f, total:%.1f \n '%( self.get_ms(time1,time0),self.get_ms(time2,time1),self.get_ms(time3,time2),self.get_ms(time4,time3),self.get_ms(time4,time0) )
|
||
|
||
#print('pre-precess:%.1f ,infer:%.1f ,post-precess:%.1f ,post-resize:%.1f, total:%.1f '%( self.get_ms(time1,time0),self.get_ms(time2,time1),self.get_ms(time3,time2),self.get_ms(time4,time3),self.get_ms(time4,time0) ))
|
||
return pred,outstr
|
||
def get_ms(self,t1,t0):
|
||
return (t1-t0)*1000.0
|
||
def preprocess_image(self,image):
|
||
|
||
time0 = time.time()
|
||
image = cv2.resize(image,self.modelsize)
|
||
time1 = time.time()
|
||
image = image.astype(np.float32)
|
||
image /= 255.0
|
||
time2 = time.time()
|
||
image[:,:,0] -=self.mean[0]
|
||
image[:,:,1] -=self.mean[1]
|
||
image[:,:,2] -=self.mean[2]
|
||
time3 = time.time()
|
||
image[:,:,0] /= self.std[0]
|
||
image[:,:,1] /= self.std[1]
|
||
image[:,:,2] /= self.std[2]
|
||
time4 = time.time()
|
||
image = cv2.cvtColor( image,cv2.COLOR_RGB2BGR)
|
||
#image -= self.mean
|
||
#image /= self.std
|
||
image = np.transpose(image, ( 2, 0, 1))
|
||
|
||
image = torch.from_numpy(image).float()
|
||
image = image.unsqueeze(0)
|
||
time5 = time.time()
|
||
print('resize:%1f ,normalize:%.1f ,Demean:%.1f ,DeVar:%.1f ,other:%.1f'%( self.get_ms(time1,time0 ), self.get_ms(time2,time1 ), self.get_ms(time3,time2 ), self.get_ms(time4,time3 ), self.get_ms(time5,time4 ) ))
|
||
|
||
return image
|
||
class SegModel_STDC(object):
|
||
def __init__(self, nclass=2,weights=None,modelsize=512,device='cuda:0',modelSize=(360,640)):
|
||
#self.args = args
|
||
self.model = BiSeNet_STDC(backbone='STDCNet813', n_classes=nclass,
|
||
use_boundary_2=False, use_boundary_4=False,
|
||
use_boundary_8=True, use_boundary_16=False,
|
||
use_conv_last=False,modelSize=modelSize)
|
||
self.device = device
|
||
self.model.load_state_dict(torch.load(weights, map_location=torch.device(self.device) ))
|
||
self.model= self.model.to(self.device)
|
||
self.mean = (0.485, 0.456, 0.406)
|
||
self.std = (0.229, 0.224, 0.225)
|
||
self.modelSize = modelSize
|
||
|
||
def eval(self,image):
|
||
time0 = time.time()
|
||
imageH, imageW, _ = image.shape
|
||
|
||
image = self.RB_convert(image)
|
||
|
||
img = self.preprocess_image(image)
|
||
|
||
if self.device != 'cpu':
|
||
imgs = img.to(self.device)
|
||
else:imgs=img
|
||
time1 = time.time()
|
||
self.model.eval()
|
||
|
||
with torch.no_grad():
|
||
output = self.model(imgs)
|
||
|
||
time2 = time.time()
|
||
pred = output.data.cpu().numpy()
|
||
pred = np.argmax(pred, axis=1)[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 '%( self.get_ms(time1,time0),self.get_ms(time2,time1),self.get_ms(time3,time2),self.get_ms(time4,time3),self.get_ms(time4,time0) )
|
||
|
||
return pred,outstr
|
||
def get_ms(self,t1,t0):
|
||
return (t1-t0)*1000.0
|
||
def preprocess_image(self,image):
|
||
|
||
image = cv2.resize(image, (self.modelSize[1],self.modelSize[0] ), interpolation=cv2.INTER_LINEAR)
|
||
image = image.astype(np.float32)
|
||
image /= 255.0
|
||
|
||
image[:, :, 0] -= self.mean[0]
|
||
image[:, :, 1] -= self.mean[1]
|
||
image[:, :, 2] -= self.mean[2]
|
||
|
||
image[:, :, 0] /= self.std[0]
|
||
image[:, :, 1] /= self.std[1]
|
||
image[:, :, 2] /= self.std[2]
|
||
|
||
image = np.transpose(image, (2, 0, 1))
|
||
image = torch.from_numpy(image).float()
|
||
image = image.unsqueeze(0)
|
||
|
||
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 get_largest_contours(contours):
|
||
areas = [cv2.contourArea(x) for x in contours]
|
||
max_area = max(areas)
|
||
max_id = areas.index(max_area)
|
||
|
||
return max_id
|
||
|
||
def infer_usage(par):
|
||
#par={'modelSize':(inputShape[3],inputShape[2]),'mean':(0.485, 0.456, 0.406),'std':(0.229, 0.224, 0.225),'RGB_convert_first':True,
|
||
# 'weights':trtFile,'device':device,'max_threads':1,
|
||
# 'image_dir':'../../AIdemo2/images/trafficAccident/','out_dir' :'results'}
|
||
|
||
|
||
segmodel = par['segmodel']
|
||
|
||
image_urls=glob.glob('%s/*'%(par['image_dir']))
|
||
out_dir =par['out_dir']
|
||
os.makedirs(out_dir,exist_ok=True)
|
||
for im,image_url in enumerate(image_urls[0:1]):
|
||
#image_url = '/home/thsw2/WJ/data/THexit/val/images/54(199).JPG'
|
||
image_array0 = cv2.imread(image_url)
|
||
H,W,C = image_array0.shape
|
||
time_1=time.time()
|
||
pred,outstr = segmodel.eval(image_array0 )
|
||
|
||
|
||
binary0 = pred.copy()
|
||
|
||
|
||
time0 = time.time()
|
||
contours, hierarchy = cv2.findContours(binary0,cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE)
|
||
max_id = -1
|
||
time1 = time.time()
|
||
time2 = time.time()
|
||
|
||
cv2.drawContours(image_array0,contours,max_id,(0,255,255),3)
|
||
time3 = time.time()
|
||
out_url='%s/%s'%(out_dir,os.path.basename(image_url))
|
||
ret = cv2.imwrite(out_url,image_array0)
|
||
cv2.imwrite(out_url.replace('.','_mask.'),(pred*50).astype(np.uint8))
|
||
time4 = time.time()
|
||
|
||
print('image:%d,%s ,%d*%d,eval:%.1f ms, %s,findcontours:%.1f ms,draw:%.1f total:%.1f'%(im,os.path.basename(image_url),H,W,get_ms(time0,time_1),outstr,get_ms(time1,time0), get_ms(time3,time2),get_ms(time3,time_1)) )
|
||
|
||
def colorstr(*input):
|
||
# Colors a string https://en.wikipedia.org/wiki/ANSI_escape_code, i.e. colorstr('blue', 'hello world')
|
||
*args, string = input if len(input) > 1 else ('blue', 'bold', input[0]) # color arguments, string
|
||
colors = {'black': '\033[30m', # basic colors
|
||
'red': '\033[31m',
|
||
'green': '\033[32m',
|
||
'yellow': '\033[33m',
|
||
'blue': '\033[34m',
|
||
'magenta': '\033[35m',
|
||
'cyan': '\033[36m',
|
||
'white': '\033[37m',
|
||
'bright_black': '\033[90m', # bright colors
|
||
'bright_red': '\033[91m',
|
||
'bright_green': '\033[92m',
|
||
'bright_yellow': '\033[93m',
|
||
'bright_blue': '\033[94m',
|
||
'bright_magenta': '\033[95m',
|
||
'bright_cyan': '\033[96m',
|
||
'bright_white': '\033[97m',
|
||
'end': '\033[0m', # misc
|
||
'bold': '\033[1m',
|
||
'underline': '\033[4m'}
|
||
return ''.join(colors[x] for x in args) + f'{string}' + colors['end']
|
||
def file_size(path):
|
||
# Return file/dir size (MB)
|
||
path = Path(path)
|
||
if path.is_file():
|
||
return path.stat().st_size / 1E6
|
||
elif path.is_dir():
|
||
return sum(f.stat().st_size for f in path.glob('**/*') if f.is_file()) / 1E6
|
||
else:
|
||
return 0.0
|
||
|
||
|
||
def toONNX(seg_model,onnxFile,inputShape=(1,3,360,640),device=torch.device('cuda:0'),dynamic=False ):
|
||
|
||
import onnx
|
||
|
||
im = torch.rand(inputShape).to(device)
|
||
seg_model.eval()
|
||
out=seg_model(im)
|
||
print('###test model infer example over ####')
|
||
train=False
|
||
dynamic = False
|
||
opset=11
|
||
print('####begin to export to onnx')
|
||
|
||
torch.onnx.export(seg_model, im,onnxFile, opset_version=opset,
|
||
training=torch.onnx.TrainingMode.TRAINING if train else torch.onnx.TrainingMode.EVAL,
|
||
do_constant_folding=not train,
|
||
input_names=['images'],
|
||
output_names=['output'],
|
||
#dynamic_axes={'images': {0: 'batch', 2: 'height', 3: 'width'}, # shape(1,3,640,640)
|
||
# 'output': {0: 'batch', 1: 'anchors'} # shape(1,25200,85)
|
||
# } if dynamic else None
|
||
dynamic_axes={
|
||
'images': {0: 'batch_size', 2: 'in_width', 3: 'int_height'},
|
||
'output': {0: 'batch_size', 2: 'out_width', 3: 'out_height'}} if dynamic else None
|
||
|
||
)
|
||
'''
|
||
input_name='images'
|
||
output_name='output'
|
||
torch.onnx.export(seg_model,
|
||
im,
|
||
onnxFile,
|
||
opset_version=11,
|
||
input_names=[input_name],
|
||
output_names=[output_name],
|
||
dynamic_axes={
|
||
input_name: {0: 'batch_size', 2: 'in_width', 3: 'int_height'},
|
||
output_name: {0: 'batch_size', 2: 'out_width', 3: 'out_height'}}
|
||
)
|
||
|
||
'''
|
||
|
||
|
||
print('output onnx file:',onnxFile)
|
||
def ONNXtoTrt(onnxFile,trtFile):
|
||
import tensorrt as trt
|
||
#onnx = Path('../weights/BiSeNet/checkpoint.onnx')
|
||
#onnxFile = Path('../weights/STDC/model_maxmIOU75_1720_0.946_360640.onnx')
|
||
time0=time.time()
|
||
half=True;verbose=True;workspace=4;prefix=colorstr('TensorRT:')
|
||
#f = onnx.with_suffix('.engine') # TensorRT engine file
|
||
f=trtFile
|
||
logger = trt.Logger(trt.Logger.INFO)
|
||
if verbose:
|
||
logger.min_severity = trt.Logger.Severity.VERBOSE
|
||
|
||
builder = trt.Builder(logger)
|
||
config = builder.create_builder_config()
|
||
config.max_workspace_size = workspace * 1 << 30
|
||
|
||
flag = (1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
|
||
network = builder.create_network(flag)
|
||
parser = trt.OnnxParser(network, logger)
|
||
if not parser.parse_from_file(str(onnxFile)):
|
||
raise RuntimeError(f'failed to load ONNX file: {onnx}')
|
||
|
||
inputs = [network.get_input(i) for i in range(network.num_inputs)]
|
||
outputs = [network.get_output(i) for i in range(network.num_outputs)]
|
||
print(f'{prefix} Network Description:')
|
||
for inp in inputs:
|
||
print(f'{prefix}\tinput "{inp.name}" with shape {inp.shape} and dtype {inp.dtype}')
|
||
for out in outputs:
|
||
print(f'{prefix}\toutput "{out.name}" with shape {out.shape} and dtype {out.dtype}')
|
||
|
||
half &= builder.platform_has_fast_fp16
|
||
print(f'{prefix} building FP{16 if half else 32} engine in {f}')
|
||
if half:
|
||
config.set_flag(trt.BuilderFlag.FP16)
|
||
with builder.build_engine(network, config) as engine, open(f, 'wb') as t:
|
||
t.write(engine.serialize())
|
||
print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
|
||
time1=time.time()
|
||
print('output trtfile from ONNX, time:%.4f s ,'%(time1-time0),trtFile)
|
||
def ONNX_eval(par):
|
||
|
||
|
||
model_path = par['weights'];
|
||
modelSize=par['modelSize']
|
||
mean = par['mean']
|
||
std = par['std']
|
||
image_urls=glob.glob('%s/*'%(par['image_dir'] ))
|
||
out_dir = par['out_dir']
|
||
|
||
# 验证模型合法性
|
||
onnx_model = onnx.load(model_path)
|
||
onnx.checker.check_model(onnx_model)
|
||
|
||
# 设置模型session以及输入信息
|
||
sess = ort.InferenceSession(str(model_path),providers= ort.get_available_providers())
|
||
print('len():',len( sess.get_inputs() ))
|
||
input_name1 = sess.get_inputs()[0].name
|
||
|
||
|
||
half = False;device = 'cuda:0'
|
||
os.makedirs(out_dir,exist_ok=True)
|
||
|
||
for im,image_url in enumerate(image_urls[0:1]):
|
||
image_array0 = cv2.imread(image_url)
|
||
#img=segPreProcess_image(image_array0).to(device)
|
||
img=segPreProcess_image(image_array0,modelSize=modelSize,mean=mean,std=std,numpy=True,RGB_convert_first=par['RGB_convert_first'])
|
||
#img = cv2.resize(img,(512,512)).transpose(2,0,1)
|
||
img = np.array(img)[np.newaxis, :, :, :].astype(np.float32)
|
||
H,W,C = image_array0.shape
|
||
time_1=time.time()
|
||
#pred,outstr = segmodel.eval(image_array0 )
|
||
print('###line343:',img.shape, os.path.basename(image_url))
|
||
print('###line343:img[0,0,10:12,10:12] ',img[0,0,10:12,10:12])
|
||
output = sess.run(None, {input_name1: img})
|
||
pred =output[0]
|
||
|
||
#pred = pred.data.cpu().numpy()
|
||
pred = np.argmax(pred, axis=1)[0]#得到每行
|
||
pred = cv2.resize(pred.astype(np.uint8),(W,H))
|
||
print('###line362:',np.max(pred))
|
||
outstr='###---###'
|
||
|
||
binary0 = pred.copy()
|
||
|
||
|
||
time0 = time.time()
|
||
contours, hierarchy = cv2.findContours(binary0,cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE)
|
||
max_id = -1
|
||
|
||
time1 = time.time()
|
||
|
||
|
||
time2 = time.time()
|
||
|
||
#cv2.drawContours(image_array0,contours,max_id,(0,255,255),3)
|
||
cv2.drawContours(image_array0,contours,-1,(0,255,255),3)
|
||
time3 = time.time()
|
||
out_url='%s/%s'%(out_dir,os.path.basename(image_url))
|
||
ret = cv2.imwrite(out_url,image_array0)
|
||
ret = cv2.imwrite(out_url.replace('.jpg','_mask.jpg').replace('.png','_mask.png' ),(pred*50).astype(np.uint8))
|
||
time4 = time.time()
|
||
|
||
print('image:%d,%s ,%d*%d,eval:%.1f ms, %s,findcontours:%.1f ms,draw:%.1f total:%.1f'%(im,os.path.basename(image_url),H,W,get_ms(time0,time_1),outstr,get_ms(time1,time0), get_ms(time3,time2),get_ms(time3,time_1)) )
|
||
print('outimage:',out_url)
|
||
|
||
|
||
#print(output)
|
||
|
||
class SegModel_STDC_trt(object):
|
||
def __init__(self,weights=None,modelsize=512,std=(0.229, 0.224, 0.225),mean=(0.485, 0.456, 0.406),device='cuda:0'):
|
||
|
||
logger = trt.Logger(trt.Logger.INFO)
|
||
with open(weights, "rb") as f, trt.Runtime(logger) as runtime:
|
||
engine=runtime.deserialize_cuda_engine(f.read())# 输入trt本地文件，返回ICudaEngine对象
|
||
self.model = TRTModule(engine, ["images"], ["output"])
|
||
self.mean = mean
|
||
self.std = std
|
||
self.device = device
|
||
self.modelsize = modelsize
|
||
|
||
|
||
def eval(self,image):
|
||
time0=time.time()
|
||
H,W,C=image.shape
|
||
img_input=self.segPreProcess_image(image)
|
||
time1=time.time()
|
||
pred=self.model(img_input)
|
||
time2=time.time()
|
||
pred=torch.argmax(pred,dim=1).cpu().numpy()[0]
|
||
#pred = np.argmax(pred.cpu().numpy(), axis=1)[0]#得到每行
|
||
time3 = time.time()
|
||
pred = cv2.resize(pred.astype(np.uint8),(W,H))
|
||
time4 = time.time()
|
||
outstr= 'pre-precess:%.1f ,infer:%.1f ,post-cpu-argmax:%.1f ,post-resize:%.1f, total:%.1f \n '%( self.get_ms(time1,time0),self.get_ms(time2,time1),self.get_ms(time3,time2),self.get_ms(time4,time3),self.get_ms(time4,time0) )
|
||
|
||
return pred,outstr
|
||
def segPreProcess_image(self,image):
|
||
|
||
image = cv2.resize(image,self.modelsize)
|
||
image = cv2.cvtColor( image,cv2.COLOR_RGB2BGR)
|
||
|
||
image = image.astype(np.float32)
|
||
image /= 255.0
|
||
image[:,:,0] -=self.mean[0]
|
||
image[:,:,1] -=self.mean[1]
|
||
image[:,:,2] -=self.mean[2]
|
||
|
||
image[:,:,0] /= self.std[0]
|
||
image[:,:,1] /= self.std[1]
|
||
image[:,:,2] /= self.std[2]
|
||
image = np.transpose(image, ( 2, 0, 1))
|
||
|
||
image = torch.from_numpy(image).float()
|
||
image = image.unsqueeze(0)
|
||
return image.to(self.device)
|
||
def get_ms(self,t1,t0):
|
||
return (t1-t0)*1000.0
|
||
|
||
|
||
|
||
def EngineInfer_onePic_thread(pars_thread):
|
||
|
||
engine,image_array0,out_dir,image_url,im ,par= pars_thread[0:6]
|
||
out_url='%s/%s'%(out_dir,os.path.basename(image_url))
|
||
|
||
H,W,C = image_array0.shape
|
||
time0=time.time()
|
||
|
||
time1=time.time()
|
||
# 运行模型
|
||
|
||
|
||
#pred,segInfoStr=segtrtEval(engine,image_array0,par={'modelSize':(640,360),'mean':(0.485, 0.456, 0.406),'std' :(0.229, 0.224, 0.225),'numpy':False, 'RGB_convert_first':True})
|
||
pred,segInfoStr=segtrtEval(engine,image_array0,par=par)
|
||
cv2.imwrite(out_url.replace('.','_mask.'),(pred*50).astype(np.uint8))
|
||
|
||
pred = 1 - pred
|
||
time2=time.time()
|
||
|
||
outstr='###---###'
|
||
binary0 = pred.copy()
|
||
time3 = time.time()
|
||
|
||
contours, hierarchy = cv2.findContours(binary0,cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE)
|
||
max_id = -1
|
||
#if len(contours)>0:
|
||
# max_id = get_largest_contours(contours)
|
||
# binary0[:,:] = 0
|
||
# cv2.fillPoly(binary0, [contours[max_id][:,0,:]], 1)
|
||
time4 = time.time()
|
||
|
||
cv2.drawContours(image_array0,contours,max_id,(0,255,255),3)
|
||
time5 = time.time()
|
||
|
||
ret = cv2.imwrite(out_url,image_array0)
|
||
time6 = time.time()
|
||
|
||
print('image:%d,%s ,%d*%d, %s,,findcontours:%.1f ms,draw:%.1f total:%.1f'%(im,os.path.basename(image_url),H,W,segInfoStr, get_ms(time4,time3),get_ms(time5,time4),get_ms(time5,time0) ))
|
||
|
||
|
||
return 'success'
|
||
|
||
|
||
def EngineInfer(par):
|
||
|
||
modelSize=par['modelSize'];mean = par['mean'] ;std = par['std'] ;RGB_convert_first=par['RGB_convert_first'];device=par['device']
|
||
weights=par['weights']; image_dir=par['image_dir']
|
||
max_threads=par['max_threads'];par['numpy']=False
|
||
image_urls=glob.glob('%s/*'%(image_dir))
|
||
out_dir =par['out_dir']
|
||
|
||
os.makedirs(out_dir,exist_ok=True)
|
||
|
||
#trt_model = SegModel_STDC_trt(weights=weights,modelsize=modelSize,std=std,mean=mean,device=device)
|
||
logger = trt.Logger(trt.Logger.ERROR)
|
||
with open(weights, "rb") as f, trt.Runtime(logger) as runtime:
|
||
engine=runtime.deserialize_cuda_engine(f.read())# 输入trt本地文件，返回ICudaEngine对象
|
||
print('#####load TRT file:',weights,'success #####')
|
||
|
||
pars_thread=[]
|
||
pars_threads=[]
|
||
for im,image_url in enumerate(image_urls[0:]):
|
||
image_array0 = cv2.imread(image_url)
|
||
pars_thread=[engine,image_array0,out_dir,image_url,im,par]
|
||
pars_threads.append(pars_thread)
|
||
#EngineInfer_onePic_thread(pars_thread)
|
||
t1=time.time()
|
||
if max_threads==1:
|
||
for i in range(len(pars_threads[0:])):
|
||
EngineInfer_onePic_thread(pars_threads[i])
|
||
|
||
'''
|
||
pred,segInfoStr=segtrtEval(pars_threads[i][0],pars_threads[i][1],par)
|
||
bname=os.path.basename( pars_threads[i][3] )
|
||
outurl= os.path.join( out_dir , bname.replace( '.png','_mask.png').replace('.jpg','._mask.jpg') )
|
||
ret=cv2.imwrite( outurl,(pred*50).astype(np.uint8))
|
||
print(ret,outurl)'''
|
||
|
||
else:
|
||
with ThreadPoolExecutor(max_workers=max_threads) as t:
|
||
for result in t.map(EngineInfer_onePic_thread, pars_threads):
|
||
tt=result
|
||
|
||
t2=time.time()
|
||
print('All %d images time:%.1f ms, each:%.1f ms , with %d threads'%(len(image_urls),(t2-t1)*1000, (t2-t1)*1000.0/len(image_urls), max_threads) )
|
||
|
||
|
||
|
||
if __name__=='__main__':
|
||
|
||
parser = argparse.ArgumentParser()
|
||
parser.add_argument('--weights', type=str, default='stdc_360X640.pth', help='model path(s)')
|
||
parser.add_argument('--nclass', type=int, default=2, help='segmodel nclass')
|
||
parser.add_argument('--mWidth', type=int, default=640, help='segmodel mWdith')
|
||
parser.add_argument('--mHeight', type=int, default=360, help='segmodel mHeight')
|
||
opt = parser.parse_args()
|
||
print( opt.weights )
|
||
#pthFile = Path('../../../yolov5TRT/weights/river/stdc_360X640.pth')
|
||
pthFile = Path(opt.weights)
|
||
onnxFile = str(pthFile.with_suffix('.onnx')).replace('360X640', '%dX%d'%( opt.mWidth,opt.mHeight ))
|
||
trtFile = onnxFile.replace('.onnx','.engine' )
|
||
|
||
nclass = opt.nclass; device=torch.device('cuda:0');
|
||
|
||
'''###BiSeNet
|
||
weights = '../weights/BiSeNet/checkpoint.pth';;inputShape =(1, 3, 512,512)
|
||
segmodel = SegModel_BiSeNet(nclass=nclass,weights=weights)
|
||
seg_model=segmodel.model
|
||
'''
|
||
|
||
##STDC net
|
||
weights = pthFile
|
||
inputShape =(1, 3, opt.mHeight,opt.mWidth)#(bs,channels,height,width)
|
||
#inputShape =(1, 3, 360,640)#(bs,channels,height,width)
|
||
segmodel = SegModel_STDC(nclass=nclass,weights=weights,modelSize=(inputShape[2],inputShape[3]));
|
||
|
||
|
||
seg_model=segmodel.model
|
||
|
||
par={'modelSize':(inputShape[3],inputShape[2]),'mean':(0.485, 0.456, 0.406),'std':(0.229, 0.224, 0.225),'RGB_convert_first':True,
|
||
'weights':trtFile,'device':device,'max_threads':1,'predResize':True,
|
||
'image_dir':'../../AIdemo2/images/trafficAccident/','out_dir' :'results'}
|
||
|
||
par_onnx =deepcopy( par)
|
||
par_onnx['weights']=onnxFile
|
||
par_pth =deepcopy( par);par_pth['segmodel']=segmodel;
|
||
#infer_usage(par_pth)
|
||
|
||
toONNX(seg_model,onnxFile,inputShape=inputShape,device=device,dynamic=True)
|
||
print('####trt to onnx over###')
|
||
ONNXtoTrt(onnxFile,trtFile)
|
||
|
||
#EngineInfer(par)
|
||
|
||
|
||
#ONNX_eval(par_onnx)
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|