AIlib2/crowd.py

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)