更新 AI.py

AI.py

@@ -4,9 +4,9 @@ from segutils.segmodel import SegModel,get_largest_contours
 from segutils.trtUtils import segtrtEval,yolov5Trtforward,OcrTrtForward
 from segutils.trafficUtils import tracfficAccidentMixFunction
 
-
 from utils.torch_utils import select_device
 from utilsK.queRiver import get_labelnames,get_label_arrays,post_process_,img_pad,draw_painting_joint,detectDraw,getDetections,getDetectionsFromPreds
+from utilsK.jkmUtils import pre_process, post_process, get_return_data
 from trackUtils.sort import moving_average_wang
 
 from utils.datasets import letterbox
@@ -18,7 +18,6 @@ import torch.nn.functional as F
 from copy import deepcopy
 from scipy import interpolate
 import glob
-from loguru import logger
 
 def get_images_videos(impth, imageFixs=['.jpg','.JPG','.PNG','.png'],videoFixs=['.MP4','.mp4','.avi']):
     imgpaths=[];###获取文件里所有的图像
@@ -36,7 +35,6 @@ def get_images_videos(impth, imageFixs=['.jpg','.JPG','.PNG','.png'],videoFixs=[
     print('%s: test Images:%d , test videos:%d '%(impth, len(imgpaths), len(videopaths)))    
     return     imgpaths,videopaths
 
-
 def xywh2xyxy(box,iW=None,iH=None):
     xc,yc,w,h = box[0:4]
     x0 =max(0, xc-w/2.0)
@@ -47,7 +45,6 @@ def xywh2xyxy(box,iW=None,iH=None):
     if iH: y0,y1 = y0*iH,y1*iH
     return [x0,y0,x1,y1]
 
-
 def get_ms(t2,t1):
     return (t2-t1)*1000.0
 def get_postProcess_para(parfile):
@@ -87,6 +84,33 @@ def filter_byClass(pdetections,allowedList):
 
     return ret
 
+# 对ocr识别车牌格式化处理
+def plat_format(ocr):
+    carDct = ['黑','吉','辽','冀','晋','陕','甘','青','鲁','苏','浙','皖','闽','赣','豫','鄂',\
+          '湘','粤','琼','川','贵','云','蒙','藏','宁','新','桂','京','津','沪','渝','使','领']
+    label = ocr[0]
+    # print(label)
+    label = list(filter(lambda x: (ord(x) > 19968 and ord(x) < 63865) or (ord(x) > 96 and ord(x) < 123)
+                                  or (ord(x) > 47 and ord(x) < 58) or (ord(x) in [33, 73, 65281]), label))
+    def spt(x):
+        if x in ['I', 'i', '!', '！']:
+            return '1'
+        else:
+            return x
+
+    label = list(map(spt, label))
+    if len(label) < 7 or len(label) >8:
+        return None
+    if not label[0] in carDct:
+        return None
+
+    label.insert(2, '・')
+    label = ' '.join(label)
+    # label = label.split('I','1').split('!','1').split('i','1').split('！','1')
+    # label = label.split('I','1').split('!','1').split('i','1').split('！','1
+
+    return label.upper()
+
 def AI_process(im0s,model,segmodel,names,label_arraylist,rainbows,objectPar={ 'half':True,'device':'cuda:0' ,'conf_thres':0.25,'iou_thres':0.45,'allowedList':[0,1,2,3],'segRegionCnt':1, 'trtFlag_det':False,'trtFlag_seg':False,'score_byClass':{x:0.1 for x in range(30)} }, font={ 'line_thickness':None, 'fontSize':None,'boxLine_thickness':None,'waterLineColor':(0,255,255),'waterLineWidth':3} ,segPar={'modelSize':(640,360),'mean':(0.485, 0.456, 0.406),'std' :(0.229, 0.224, 0.225),'numpy':False, 'RGB_convert_first':True},mode='others',postPar=None):
 
     #输入参数
@@ -124,7 +148,6 @@ def AI_process(im0s,model,segmodel,names,label_arraylist,rainbows,objectPar={ 'h
     img = img[:, :, :, ::-1].transpose(0, 3, 1, 2)  # BGR to RGB, to bsx3x416x416
     img = np.ascontiguousarray(img)
 
-    
     img = torch.from_numpy(img).to(device)
     img = img.half() if half else img.float()  # uint8 to fp16/32
     img /= 255.0
@@ -132,19 +155,10 @@ def AI_process(im0s,model,segmodel,names,label_arraylist,rainbows,objectPar={ 'h
 
     if segmodel: 
         seg_pred,segstr = segmodel.eval(im0s[0] )
-        # 当不存在分割信息，无需做分类检测
+        segFlag=True
-        # segFlag = True
-        logger.info("分割信息seg_prd: {}  数据类型：{} ", seg_pred, np.count_nonzero(seg_pred))
-        if not np.any(seg_pred != 0):
-            time_info = 'No SegMentInfo'
-            return [], time_info
-
     else:
-        # seg_pred = None;
+        seg_pred = None;segFlag=False;segstr='Not implemented'
-        # segFlag = False;
+
-        # segstr = 'Not implemented'
-        time_info = 'No SegMentInfo'
-        return [], time_info
 
     time1=time.time()
     if trtFlag_det:
@@ -182,7 +196,6 @@ def AI_process(im0s,model,segmodel,names,label_arraylist,rainbows,objectPar={ 'h
     return p_result,time_info
 def default_mix(predlist,par):
     return  predlist[0],''
-    
 def AI_process_N(im0s,modelList,postProcess):
 
     #输入参数
@@ -209,7 +222,6 @@ def AI_process_N(im0s,modelList,postProcess):
     #ret就是混合处理后的结果
     ret = mixFunction( predsList, postProcess['pars'])   
     return ret[0],timeInfos+ret[1]
-    
 def getMaxScoreWords(detRets0):
     maxScore=-1;maxId=0
     for i,detRet in enumerate(detRets0):
@@ -218,7 +230,6 @@ def getMaxScoreWords(detRets0):
             maxScore = detRet[4]
     return maxId
 
-    
 def AI_process_C(im0s,modelList,postProcess):
     #函数定制的原因：
     ## 之前模型处理流是
@@ -264,7 +275,7 @@ def AI_process_C(im0s,modelList,postProcess):
         mixFunction =postProcess['function']
         crackInfos = [mixFunction(patchMask,par=parsIn) for  patchMask in detRets1]
 
-        rets = [  _detRets0[i]+ crackInfos[i] for i in range(len(imagePatches))  ]
+        rets = [detRets0[i]+ crackInfos[i] for i in range(len(imagePatches))  ]
         t3=time.time()
         outInfos='total:%.1f (det:%.1f  %d次segs:%.1f mixProcess:%.1f) '%(  (t3-t0)*1000, (t1-t0)*1000, len(detRets1),(t2-t1)*1000, (t3-t2)*1000   )
     elif   postProcess['name']=='channel2': 
@@ -289,8 +300,6 @@ def AI_process_C(im0s,modelList,postProcess):
 
     return rets,outInfos
 
-    
-       
 def AI_process_forest(im0s,model,segmodel,names,label_arraylist,rainbows,half=True,device=' cuda:0',conf_thres=0.25, iou_thres=0.45,allowedList=[0,1,2,3], font={ 'line_thickness':None, 'fontSize':None,'boxLine_thickness':None,'waterLineColor':(0,255,255),'waterLineWidth':3} ,trtFlag_det=False,SecNms=None):
     #输入参数
     #   im0s---原始图像列表
@@ -338,8 +347,6 @@ def AI_process_forest(im0s,model,segmodel,names,label_arraylist,rainbows,half=Tr
     #p_result,timeOut = post_process_(datas,conf_thres, iou_thres,names,label_arraylist,rainbows,10,object_config=allowedList,segmodel=segFlag,font=font,padInfos=padInfos)
     time_info = 'letterbox:%.1f, infer:%.1f, '%( (time1-time0)*1000,(time2-time1)*1000 )
     return p_result,time_info+timeOut
-    
-
 def AI_det_track( im0s_in,modelPar,processPar,sort_tracker,segPar=None):
     im0s,iframe=im0s_in[0],im0s_in[1]
     model = modelPar['det_Model']
@@ -664,14 +671,11 @@ def AI_det_track_batch_N(imgarray_list, iframe_list ,modelList,postProcess,sort_
             #[ifrmae,  x0 ,y0  ,x1  ,y1  ,conf,cls,trackId ]
             detResults.append( res )
 
-  
         retResults=[imgarray_list,track_det_result,detResults ]
         t2 = time.time()
         timeInfos = 'detTrack:%.1f TrackPost:%.1f, %s'%(get_ms(t1,t0),get_ms(t2,t1), timeInfos_track )
         return retResults,timeInfos
 
-                
-    
 def ocr_process(pars):
 
     img_patch,engine,context,converter,AlignCollate_normal,device=pars[0:6]
@@ -700,6 +704,63 @@ def ocr_process(pars):
 
     info_str= ('pre-process:%.2f  TRTforward:%.2f (%s)  postProcess:%2.f  decoder:%.2f, Total:%.2f , pred:%s'%(get_ms(time2,time1 ),get_ms(time3,time2 ),trtstr, get_ms(time4,time3 ), get_ms(time5,time4 ),  get_ms(time5,time1 ),  preds_str  )  )
     return  preds_str,info_str
+
+def AI_process_Ocr(im0s,modelList,device,detpar):
+    timeMixPost = ':0 ms'
+    new_device = torch.device(device)
+    time0 = time.time()
+    img, padInfos = pre_process(im0s[0], new_device)
+    ocrModel = modelList[1]
+    time1 = time.time()
+    preds,timeOut = modelList[0].eval(img)
+    time2 = time.time()
+    boxes = post_process(preds, padInfos, device, conf_thres=detpar['conf_thres'], iou_thres=detpar['iou_thres'],
+                         nc=detpar['nc'])  # 后处理
+    imagePatches = [im0s[0][int(x[1]):int(x[3]), int(x[0]):int(x[2])] for x in boxes]
+
+    detRets1 = [ocrModel.eval(patch) for patch in imagePatches]
+    time3 = time.time()
+    dets = []
+    for i, (box, ocr) in enumerate(zip(boxes, detRets1)):
+        label = plat_format(ocr)
+        if label:
+            xyxy = box[0:4]
+            dets.append([label, xyxy])
+
+    time_info = 'pre_process:%.1f, det:%.1f , ocr:%.1f  ,timeMixPost:%s ' % (
+    (time1 - time0) * 1000, (time2 - time1) * 1000, (time3 - time2) * 1000, timeMixPost)
+
+    return [im0s[0],im0s[0],dets,0],time_info
+
+
+def AI_process_Crowd(im0s,model,device,postPar):
+    timeMixPost = ':0 ms'
+    new_device = torch.device(device)
+    time0 = time.time()
+    preds = model.eval(im0s[0])
+    time1 = time.time()
+    outputs_scores = torch.nn.functional.softmax(preds['pred_logits'], -1)[:, :, 1][0]
+
+    outputs_points = preds['pred_points'][0]
+
+    points = outputs_points[outputs_scores > postPar['conf']].detach().cpu().numpy().tolist()
+    predict_cnt = int((outputs_scores > postPar['conf']).sum())
+    #img_to_draw = cv2.cvtColor(np.array(img_raw), cv2.COLOR_RGB2BGR)
+    time2 = time.time()
+    # for p in points:
+    #     img_to_draw = cv2.circle(img_to_draw, (int(p[0]), int(p[1])), line, (0, 0, 255), -1)
+    Calc_label = '当前人数: %d' % (predict_cnt)
+
+
+    dets = [[Calc_label, points]]
+    time_info = 'det:%.1f , post:%.1f  ,timeMixPost:%s ' % (
+    (time1 - time0) * 1000, (time2 - time1) * 1000, timeMixPost)
+
+
+
+    return [im0s[0],im0s[0],dets,0],time_info
+
+
 def main():
     ##预先设置的参数
     device_='1'  ##选定模型，可选 cpu,'0','1'