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