191 lines
8.5 KiB
Python
191 lines
8.5 KiB
Python
import argparse
|
||
from PIL import Image
|
||
from crowdUtils.engine import standard_transforms,preprocess,postprocess,DictToObject,AnchorPointsf
|
||
from crowdUtils.models import build_model
|
||
from segutils import trtUtils2
|
||
import os,torch,cv2,time
|
||
import numpy as np
|
||
import warnings
|
||
import tensorrt as trt
|
||
from copy import deepcopy
|
||
warnings.filterwarnings('ignore')
|
||
|
||
class crowdModel(object):
|
||
def __init__(self, weights=None,
|
||
par={'mean':[0.485, 0.456, 0.406], 'std':[0.229, 0.224, 0.225],'threshold':0.5,
|
||
'modelPar':{'backbone':'vgg16_bn', 'gpu_id':0,'anchorFlag':False,'line':2,'width':None,'height':None , 'output_dir':'./output', 'row':2}
|
||
}
|
||
):
|
||
print('-'*20,par['modelPar'] )
|
||
self.mean = par['mean']
|
||
self.std =par['std']
|
||
|
||
self.width = par['modelPar']['width']
|
||
self.height = par['modelPar']['height']
|
||
self.minShape = par['input_profile_shapes'][0]
|
||
self.maxShape = par['input_profile_shapes'][2]
|
||
self.IOShapes0,self.IOShapes1 = [ None,None,None ],[ None,None,None ]
|
||
self.Oshapes0,self.Oshapes1 = [ None,None,None ], [ None,None,None ]
|
||
|
||
|
||
self.modelPar = DictToObject(par['modelPar'])
|
||
self.threshold = par['threshold']
|
||
self.device = 'cuda:0'
|
||
|
||
|
||
|
||
if weights.endswith('.engine') or weights.endswith('.trt'):
|
||
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.engine=runtime.deserialize_cuda_engine(f.read())# 输入trt本地文件，返回ICudaEngine对象
|
||
#self.stream=cuda.Stream()
|
||
self.bindingNames=[ self.engine.get_binding_name(ib) for ib in range(len(self.engine)) ]
|
||
print('############load seg model trt success: ',weights,self.bindingNames)
|
||
self.inputs,self.outputs,self.bindings,self.stream=None,None,None,None
|
||
self.context = self.engine.create_execution_context()
|
||
|
||
elif self.infer_type=='pth':
|
||
#self.model = DirectionalPointDetector(3, self.par['depth_factor'], self.par['NUM_FEATURE_MAP_CHANNEL']).to(self.device)
|
||
|
||
self.model = build_model(self.modelPar)
|
||
checkpoint = torch.load(args.weight_path, map_location='cpu')
|
||
self.model.load_state_dict(checkpoint['model'])
|
||
self.model=self.model.to(self.device)
|
||
if not self.modelPar.anchorFlag:
|
||
if self.infer_type=='trt':
|
||
self.anchors = AnchorPointsf(pyramid_levels=[3,], strides=None, row=self.modelPar.row, line=self.modelPar.line,device='cpu' )
|
||
elif self.infer_type=='pth':
|
||
self.anchors = AnchorPointsf(pyramid_levels=[3,], strides=None, row=self.modelPar.row, line=self.modelPar.line ,device='cuda:0')
|
||
|
||
print('#########加载模型:',weights,' 类型:',self.infer_type)
|
||
|
||
def preprocess(self,img):
|
||
tmpImg = preprocess(img,mean=self.mean, std=self.std,minShape=self.minShape,maxShape=self.maxShape)
|
||
if self.infer_type=='pth':
|
||
tmpImg = torch.from_numpy(tmpImg)
|
||
tmpImg = torch.Tensor(tmpImg).unsqueeze(0)
|
||
|
||
elif self.infer_type=='trt':
|
||
#if not self.height:
|
||
chs, height, width= tmpImg.shape[0:3]
|
||
self.width, self.height = width,height
|
||
self.IOShapes1 = [ (1, chs, height, width ),(1, height//4*width//4,2),(1, height//4*width//4,2) ]
|
||
self.Oshapes1 = [ (1, height//4*width//4,2),(1, height//4*width//4,2) ]
|
||
tmpImg = tmpImg[np.newaxis,:,:,:]#CHW->NCHW
|
||
|
||
return tmpImg
|
||
def ms(self,t1,t0):
|
||
return '%.1f'%( (t1-t0)*1000 )
|
||
def eval(self,img):
|
||
time0 = time.time()
|
||
img_b = img.copy()
|
||
#print('-----line54:',img.shape)
|
||
samples = self.preprocess(img)
|
||
|
||
time1 = time.time()
|
||
if self.infer_type=='pth':
|
||
samples = samples.to(self.device)
|
||
elif self.infer_type=='trt' :
|
||
|
||
#print('##### line83: 决定是否申请 内存 ',self.IOShapes1, self.IOShapes0,self.IOShapes1==self.IOShapes0)
|
||
#if self.IOShapes1 != self.IOShapes0:
|
||
self.inputs,self.outputs,self.bindings,self.stream = trtUtils2.allocate_buffers(self.engine,self.IOShapes1)
|
||
#print('##### line96: 开辟新内存成功 ' ,self.height,self.width)
|
||
self.IOShapes0=deepcopy(self.IOShapes1)
|
||
|
||
|
||
|
||
time2 = time.time()
|
||
if not self.modelPar.anchorFlag:
|
||
self.anchor_points = self.anchors.eval(samples)
|
||
if self.infer_type=='pth':
|
||
# run inference
|
||
self.model.eval()
|
||
with torch.no_grad():
|
||
outputs = self.model(samples)
|
||
outputs['pred_points'] = outputs['pred_points'] + self.anchor_points
|
||
#print('###line64:',outputs.keys(), outputs['pred_points'].shape, outputs['pred_logits'].shape)
|
||
elif self.infer_type=='trt':
|
||
outputs = trtUtils2.trt_inference( samples,self.height,self.width,self.context,self.inputs,self.outputs,self.bindings,self.stream,input_name = self.bindingNames[0])
|
||
for i in range(len(self.Oshapes1)):
|
||
outputs[i] = torch.from_numpy( np.reshape(outputs[i],self.Oshapes1[i]))
|
||
|
||
outputs={'pred_points':outputs[0], 'pred_logits':outputs[1]}
|
||
#print('###line117:',outputs.keys(), outputs['pred_points'].shape, outputs['pred_logits'].shape)
|
||
outputs['pred_points'] = outputs['pred_points'] + self.anchor_points
|
||
|
||
|
||
time3 = time.time()
|
||
points,scores = self.postprocess(outputs)
|
||
time4 = time.time()
|
||
infos = 'precess:%s datacopy:%s infer:%s post:%s'%( self.ms(time1,time0) , self.ms(time2,time1), self.ms(time3,time2), self.ms(time4,time3) )
|
||
|
||
p2 = self.toOBBformat(points,scores,cls=0 )
|
||
|
||
presults=[ img_b, points,p2 ]
|
||
|
||
|
||
|
||
|
||
return presults, infos
|
||
|
||
def postprocess(self,outputs):
|
||
return postprocess(outputs,threshold=self.threshold)
|
||
def toOBBformat(self,points,scores,cls=0):
|
||
outs = []
|
||
for i in range(len(points)):
|
||
pt,score = points[i],scores[i]
|
||
pts4=[pt]*4
|
||
ret = [ pts4,score,cls]
|
||
outs.append(ret)
|
||
return outs
|
||
|
||
def main():
|
||
|
||
par={'mean':[0.485, 0.456, 0.406], 'std':[0.229, 0.224, 0.225],'threshold':0.5, 'output_dir':'./output','input_profile_shapes':[(1,3,256,256),(1,3,1024,1024),(1,3,2048,2048)],'modelPar':{'backbone':'vgg16_bn', 'gpu_id':0,'anchorFlag':False, 'width':None,'height':None ,'line':2, 'row':2}
|
||
}
|
||
weights='weights/best_mae_dynamic.engine'
|
||
#weights='weights/best_mae.pth'
|
||
cmodel = crowdModel(weights,par)
|
||
|
||
img_path = "./testImages"
|
||
File = os.listdir(img_path)
|
||
targetList = []
|
||
for file in File[0:]:
|
||
COORlist = []
|
||
imgPath = img_path + os.sep + file
|
||
img_raw = np.array(Image.open(imgPath).convert('RGB') )
|
||
points, infos = cmodel.eval(img_raw)
|
||
print(file,infos,img_raw.shape)
|
||
img_to_draw = cv2.cvtColor(np.array(img_raw), cv2.COLOR_RGB2BGR)
|
||
# 打印预测图像中人头的个数
|
||
for p in points:
|
||
img_to_draw = cv2.circle(img_to_draw, (int(p[0]), int(p[1])), 2, (0, 255, 0), -1)
|
||
COORlist.append((int(p[0]), int(p[1])))
|
||
# 将各测试图像中的人头坐标存储在targetList中, 格式：[[(x1, y1),(x2, y2),...], [(X1, Y1),(X2, Y2),..], ...]
|
||
targetList.append(COORlist)
|
||
time.sleep(2)
|
||
# 保存预测图片
|
||
cv2.imwrite(os.path.join(par['output_dir'], file), img_to_draw)
|
||
#print(targetList )
|
||
|
||
|
||
|
||
|
||
|
||
|
||
if __name__ == '__main__':
|
||
|
||
par = {'backbone':'vgg16_bn', 'gpu_id':0, 'line':2, 'output_dir':'./output', 'row':2,'anchorFlag':False, 'weight_path':'./weights/best_mae.pth'}
|
||
args = DictToObject(par)
|
||
|
||
targetList = main()
|
||
print("line81", targetList)
|