From 9cba2fda3a933871b63bf62571d1e4b38f0b71a8 Mon Sep 17 00:00:00 2001 From: zhoushuliang Date: Thu, 10 Jul 2025 17:49:37 +0800 Subject: [PATCH] =?UTF-8?q?=E6=9B=B4=E6=96=B0=20AI.py?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- AI.py | 1525 ++++++++++++++++++++++++++++++--------------------------- 1 file changed, 793 insertions(+), 732 deletions(-) diff --git a/AI.py b/AI.py index 9fb0da2..f8d860c 100644 --- a/AI.py +++ b/AI.py @@ -1,80 +1,77 @@ -import cv2,os,time,json -from models.experimental import attempt_load -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 trackUtils.sort import moving_average_wang - -from utils.datasets import letterbox -import numpy as np -import torch -import math -from PIL import Image -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=[];###获取文件里所有的图像 - videopaths=[]###获取文件里所有的视频 - if os.path.isdir(impth): - for postfix in imageFixs: - imgpaths.extend(glob.glob('%s/*%s'%(impth,postfix )) ) - for postfix in videoFixs: - videopaths.extend(glob.glob('%s/*%s'%(impth,postfix )) ) - else: - postfix = os.path.splitext(impth)[-1] - if postfix in imageFixs: imgpaths=[ impth ] - if postfix in videoFixs: videopaths = [impth ] - - 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) - x1 =min(1, xc+w/2.0) - y0=max(0, yc-h/2.0) - y1=min(1,yc+h/2.0) - if iW: x0,x1 = x0*iW,x1*iW - 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): - with open(parfile) as fp: - par = json.load(fp) - assert 'post_process' in par.keys(), ' parfile has not key word:post_process' - parPost=par['post_process'] - - return parPost["conf_thres"],parPost["iou_thres"],parPost["classes"],parPost["rainbows"] -def get_postProcess_para_dic(parfile): - with open(parfile) as fp: - par = json.load(fp) - parPost=par['post_process'] - return parPost -def score_filter_byClass(pdetections,score_para_2nd): - ret=[] - for det in pdetections: - score,cls = det[4],det[5] - if int(cls) in score_para_2nd.keys(): - score_th = score_para_2nd[int(cls)] - elif str(int(cls)) in score_para_2nd.keys(): - score_th = score_para_2nd[str(int(cls))] - else: - score_th = 0.7 - if score > score_th: - ret.append(det) - return ret +import cv2,os,time,json +from models.experimental import attempt_load +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 +import numpy as np +import torch +import math +from PIL import Image +import torch.nn.functional as F +from copy import deepcopy +from scipy import interpolate +import glob + +def get_images_videos(impth, imageFixs=['.jpg','.JPG','.PNG','.png'],videoFixs=['.MP4','.mp4','.avi']): + imgpaths=[];###获取文件里所有的图像 + videopaths=[]###获取文件里所有的视频 + if os.path.isdir(impth): + for postfix in imageFixs: + imgpaths.extend(glob.glob('%s/*%s'%(impth,postfix )) ) + for postfix in videoFixs: + videopaths.extend(glob.glob('%s/*%s'%(impth,postfix )) ) + else: + postfix = os.path.splitext(impth)[-1] + if postfix in imageFixs: imgpaths=[ impth ] + if postfix in videoFixs: videopaths = [impth ] + + 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) + x1 =min(1, xc+w/2.0) + y0=max(0, yc-h/2.0) + y1=min(1,yc+h/2.0) + if iW: x0,x1 = x0*iW,x1*iW + 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): + with open(parfile) as fp: + par = json.load(fp) + assert 'post_process' in par.keys(), ' parfile has not key word:post_process' + parPost=par['post_process'] + + return parPost["conf_thres"],parPost["iou_thres"],parPost["classes"],parPost["rainbows"] +def get_postProcess_para_dic(parfile): + with open(parfile) as fp: + par = json.load(fp) + parPost=par['post_process'] + return parPost +def score_filter_byClass(pdetections,score_para_2nd): + ret=[] + for det in pdetections: + score,cls = det[4],det[5] + if int(cls) in score_para_2nd.keys(): + score_th = score_para_2nd[int(cls)] + elif str(int(cls)) in score_para_2nd.keys(): + score_th = score_para_2nd[str(int(cls))] + else: + score_th = 0.7 + if score > score_th: + ret.append(det) + return ret # 按类过滤 def filter_byClass(pdetections,allowedList): ret=[] @@ -86,662 +83,726 @@ def filter_byClass(pdetections,allowedList): ret.append(det) return ret - -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): - - #输入参数 - # im0s---原始图像列表 - # model---检测模型，segmodel---分割模型（如若没有用到，则为None） - # - #输出：两个元素（列表，字符）构成的元组，[im0s[0],im0,det_xywh,iframe],strout - # [im0s[0],im0,det_xywh,iframe]中， - # im0s[0]--原始图像，im0--AI处理后的图像,iframe--帧号/暂时不需用到。 - # det_xywh--检测结果，是一个列表。 - # 其中每一个元素表示一个目标构成如：[ xc,yc,w,h, float(conf_c),float(cls_c) ] ,2023.08.03修改输出格式 - # #cls_c--类别，如0,1,2,3; xc,yc,w,h--中心点坐标及宽；conf_c--得分, 取值范围在0-1之间 - # #strout---统计AI处理个环节的时间 - # Letterbox - - half,device,conf_thres,iou_thres,allowedList = objectPar['half'],objectPar['device'],objectPar['conf_thres'],objectPar['iou_thres'],objectPar['allowedList'] - - trtFlag_det,trtFlag_seg,segRegionCnt = objectPar['trtFlag_det'],objectPar['trtFlag_seg'],objectPar['segRegionCnt'] - if 'ovlap_thres_crossCategory' in objectPar.keys(): ovlap_thres = objectPar['ovlap_thres_crossCategory'] - else: ovlap_thres = None - - if 'score_byClass' in objectPar.keys(): score_byClass = objectPar['score_byClass'] - else: score_byClass = None - - time0=time.time() - if trtFlag_det: - img, padInfos = img_pad(im0s[0], size=(640,640,3)) ;img = [img] - else: - #print('####line72:',im0s[0][10:12,10:12,2]) - img = [letterbox(x, 640, auto=True, stride=32)[0] for x in im0s];padInfos=None - #print('####line74:',img[0][10:12,10:12,2]) - # Stack - img = np.stack(img, 0) - # Convert - 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 - time01=time.time() - - if segmodel: - seg_pred, segstr = segmodel.eval(im0s[0]) - # 当不存在分割信息，无需做分类检测 - # 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 + +# 对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): + + #输入参数 + # im0s---原始图像列表 + # model---检测模型，segmodel---分割模型（如若没有用到，则为None） + # + #输出：两个元素（列表，字符）构成的元组，[im0s[0],im0,det_xywh,iframe],strout + # [im0s[0],im0,det_xywh,iframe]中， + # im0s[0]--原始图像，im0--AI处理后的图像,iframe--帧号/暂时不需用到。 + # det_xywh--检测结果，是一个列表。 + # 其中每一个元素表示一个目标构成如：[ xc,yc,w,h, float(conf_c),float(cls_c) ] ,2023.08.03修改输出格式 + # #cls_c--类别，如0,1,2,3; xc,yc,w,h--中心点坐标及宽；conf_c--得分, 取值范围在0-1之间 + # #strout---统计AI处理个环节的时间 + # Letterbox + + half,device,conf_thres,iou_thres,allowedList = objectPar['half'],objectPar['device'],objectPar['conf_thres'],objectPar['iou_thres'],objectPar['allowedList'] + + trtFlag_det,trtFlag_seg,segRegionCnt = objectPar['trtFlag_det'],objectPar['trtFlag_seg'],objectPar['segRegionCnt'] + if 'ovlap_thres_crossCategory' in objectPar.keys(): ovlap_thres = objectPar['ovlap_thres_crossCategory'] + else: ovlap_thres = None + + if 'score_byClass' in objectPar.keys(): score_byClass = objectPar['score_byClass'] + else: score_byClass = None + + time0=time.time() + if trtFlag_det: + img, padInfos = img_pad(im0s[0], size=(640,640,3)) ;img = [img] + else: + #print('####line72:',im0s[0][10:12,10:12,2]) + img = [letterbox(x, 640, auto=True, stride=32)[0] for x in im0s];padInfos=None + #print('####line74:',img[0][10:12,10:12,2]) + # Stack + img = np.stack(img, 0) + # Convert + 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 + time01=time.time() + + if segmodel: + seg_pred,segstr = segmodel.eval(im0s[0] ) + segFlag=True + else: + seg_pred = None;segFlag=False;segstr='Not implemented' + + + time1=time.time() + if trtFlag_det: + pred = yolov5Trtforward(model,img) + else: + #print('####line96:',img[0,0,10:12,10:12]) + pred = model(img,augment=False)[0] + + time2=time.time() + + + p_result, timeOut = getDetectionsFromPreds(pred,img,im0s[0],conf_thres=conf_thres,iou_thres=iou_thres,ovlap_thres=ovlap_thres,padInfos=padInfos) + if score_byClass: + p_result[2] = score_filter_byClass(p_result[2],score_byClass) + #if mode=='highWay3.0': + #if segmodel: + if segPar and segPar['mixFunction']['function']: + + mixFunction = segPar['mixFunction']['function'];H,W = im0s[0].shape[0:2] + parMix = segPar['mixFunction']['pars'];#print('###line117:',parMix,p_result[2]) + parMix['imgSize'] = (W,H) + #print(' -----------line110: ',p_result[2] ,'\n', seg_pred) + p_result[2] , timeMixPost= mixFunction(p_result[2], seg_pred, pars=parMix ) + #print(' -----------line112: ',p_result[2] ) + p_result.append(seg_pred) else: - # seg_pred = None; - # segFlag = False; - # segstr = 'Not implemented' - time_info = 'No SegMentInfo' - return [], time_info - - time1=time.time() - if trtFlag_det: - pred = yolov5Trtforward(model,img) - else: - #print('####line96:',img[0,0,10:12,10:12]) - pred = model(img,augment=False)[0] - - time2=time.time() - - - p_result, timeOut = getDetectionsFromPreds(pred,img,im0s[0],conf_thres=conf_thres,iou_thres=iou_thres,ovlap_thres=ovlap_thres,padInfos=padInfos) - if score_byClass: - p_result[2] = score_filter_byClass(p_result[2],score_byClass) - #if mode=='highWay3.0': - #if segmodel: - if segPar and segPar['mixFunction']['function']: - - mixFunction = segPar['mixFunction']['function'];H,W = im0s[0].shape[0:2] - parMix = segPar['mixFunction']['pars'];#print('###line117:',parMix,p_result[2]) - parMix['imgSize'] = (W,H) - #print(' -----------line110: ',p_result[2] ,'\n', seg_pred) - p_result[2] , timeMixPost= mixFunction(p_result[2], seg_pred, pars=parMix ) - #print(' -----------line112: ',p_result[2] ) - p_result.append(seg_pred) - - else: - timeMixPost=':0 ms' - #print('#### line121: segstr:%s timeMixPost:%s timeOut:%s'%( segstr.strip(), timeMixPost,timeOut )) - time_info = 'letterbox:%.1f, seg:%.1f , infer:%.1f,%s, seginfo:%s ,timeMixPost:%s '%( (time01-time0)*1000, (time1-time01)*1000 ,(time2-time1)*1000,timeOut , segstr.strip(),timeMixPost ) + timeMixPost=':0 ms' + #print('#### line121: segstr:%s timeMixPost:%s timeOut:%s'%( segstr.strip(), timeMixPost,timeOut )) + time_info = 'letterbox:%.1f, seg:%.1f , infer:%.1f,%s, seginfo:%s ,timeMixPost:%s '%( (time01-time0)*1000, (time1-time01)*1000 ,(time2-time1)*1000,timeOut , segstr.strip(),timeMixPost ) if allowedList: p_result[2] = filter_byClass(p_result[2],allowedList) - - print('-'*10,p_result[2]) - return p_result,time_info -def default_mix(predlist,par): - return predlist[0],'' - -def AI_process_N(im0s,modelList,postProcess): - - #输入参数 - ## im0s---原始图像列表 - ## modelList--所有的模型 - # postProcess--字典{},包括后处理函数，及其参数 - #输出参数 - ##ret[0]--检测结果； - ##ret[1]--时间信息 - - #modelList包括模型，每个模型是一个类，里面的eval函数可以输出该模型的推理结果 - modelRets=[ model.eval(im0s[0]) for model in modelList] - - timeInfos = [ x[1] for x in modelRets] - timeInfos=''.join(timeInfos) - timeInfos=timeInfos - - #postProcess['function']--后处理函数，输入的就是所有模型输出结果 - mixFunction =postProcess['function'] - predsList = [ modelRet[0] for modelRet in modelRets ] - H,W = im0s[0].shape[0:2] - postProcess['pars']['imgSize'] = (W,H) - - #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): - if detRet[4]>maxScore: - maxId=i - maxScore = detRet[4] - return maxId - - -def AI_process_C(im0s,modelList,postProcess): - #函数定制的原因： - ## 之前模型处理流是 - ## 图片---> 模型1-->result1；图片---> 模型2->result2；[result1,result2]--->后处理函数 - ## 本函数的处理流程是 - ## 图片---> 模型1-->result1；[图片,result1]---> 模型2->result2；[result1,result2]--->后处理函数 - ## 模型2的输入，是有模型1的输出决定的。如模型2是ocr模型，需要将模型1检测出来的船名抠图出来输入到模型2. - ## 之前的模型流都是模型2是分割模型，输入就是原始图片，与模型1的输出无关。 - #输入参数 - ## im0s---原始图像列表 - ## modelList--所有的模型 - # postProcess--字典{},包括后处理函数，及其参数 - #输出参数 - ##ret[0]--检测结果； - ##ret[1]--时间信息 - - #modelList包括模型，每个模型是一个类，里面的eval函数可以输出该模型的推理结果 - - t0=time.time() - detRets0 = modelList[0].eval(im0s[0]) - - #detRets0=[[12, 46, 1127, 1544, 0.2340087890625, 2.0], [1884, 1248, 2992, 1485, 0.64208984375, 1.0]] - detRets0 = detRets0[0] - parsIn=postProcess['pars'] - - _detRets0_obj = list(filter(lambda x: x[5] in parsIn['objs'], detRets0 )) - _detRets0_others = list(filter(lambda x: x[5] not in parsIn['objs'], detRets0 )) - _detRets0 = [] - if postProcess['name']=='channel2': - if len(_detRets0_obj)>0: - maxId=getMaxScoreWords(_detRets0_obj) - _detRets0 = _detRets0_obj[maxId:maxId+1] - else: _detRets0 = detRets0 - - - t1=time.time() - imagePatches = [ im0s[0][int(x[1]):int(x[3] ) ,int(x[0]):int(x[2])] for x in _detRets0 ] - detRets1 = [modelList[1].eval(patch) for patch in imagePatches] - print('###line240:',detRets1) - if postProcess['name']=='crackMeasurement': - detRets1 = [x[0]*255 for x in detRets1] - t2=time.time() - mixFunction =postProcess['function'] - crackInfos = [mixFunction(patchMask,par=parsIn) for patchMask in detRets1] - - 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': - H,W = im0s[0].shape[0:2];parsIn['imgSize'] = (W,H) - mixFunction =postProcess['function'] - _detRets0_others = mixFunction([_detRets0_others], parsIn) - ocrInfo='no ocr' - if len(_detRets0_obj)>0: - res_real = detRets1[0][0] - res_real="".join( list(filter(lambda x:(ord(x) >19968 and ord(x)<63865 ) or (ord(x) >47 and ord(x)<58 ),res_real))) - - #detRets1[0][0]="".join( list(filter(lambda x:(ord(x) >19968 and ord(x)<63865 ) or (ord(x) >47 and ord(x)<58 ),detRets1[0][0]))) - _detRets0_obj[maxId].append(res_real ) - _detRets0_obj = [_detRets0_obj[maxId]]##只输出有OCR的那个船名结果 - ocrInfo=detRets1[0][1] - print( ' _detRets0_obj:{} _detRets0_others:{} '.format( _detRets0_obj, _detRets0_others ) ) - rets=_detRets0_obj+_detRets0_others - t3=time.time() - outInfos='total:%.1f ,where det:%.1f, ocr:%s'%( (t3-t0)*1000, (t1-t0)*1000, ocrInfo) - - #print('###line233:',detRets1,detRets0 ) - - 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---原始图像列表 - # model---检测模型，segmodel---分割模型（如若没有用到，则为None） - #输出：两个元素（列表，字符）构成的元组，[im0s[0],im0,det_xywh,iframe],strout - # [im0s[0],im0,det_xywh,iframe]中， - # im0s[0]--原始图像，im0--AI处理后的图像,iframe--帧号/暂时不需用到。 - # det_xywh--检测结果，是一个列表。 - # 其中每一个元素表示一个目标构成如：[ xc,yc,w,h, float(conf_c),float(cls_c)],#2023.08.03,修改输出格式 - # #cls_c--类别，如0,1,2,3; xc,yc,w,h--中心点坐标及宽；conf_c--得分, 取值范围在0-1之间 - # #strout---统计AI处理个环节的时间 - - # Letterbox - time0=time.time() - if trtFlag_det: - img, padInfos = img_pad(im0s[0], size=(640,640,3)) ;img = [img] - else: - img = [letterbox(x, 640, auto=True, stride=32)[0] for x in im0s];padInfos=None - #img = [letterbox(x, 640, auto=True, stride=32)[0] for x in im0s] - # Stack - img = np.stack(img, 0) - # Convert - 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 # 0 - 255 to 0.0 - 1.0 - if segmodel: - seg_pred,segstr = segmodel.eval(im0s[0] ) - segFlag=True - else: - seg_pred = None;segFlag=False - time1=time.time() - pred = yolov5Trtforward(model,img) if trtFlag_det else model(img,augment=False)[0] - - - time2=time.time() - datas = [[''], img, im0s, None,pred,seg_pred,10] - - ObjectPar={ 'object_config':allowedList, 'slopeIndex':[] ,'segmodel':segFlag,'segRegionCnt':0 } - p_result,timeOut = post_process_(datas,conf_thres, iou_thres,names,label_arraylist,rainbows,10,ObjectPar=ObjectPar,font=font,padInfos=padInfos,ovlap_thres=SecNms) - #print('###line274:',p_result[2]) - #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'] - segmodel = modelPar['seg_Model'] - half,device,conf_thres, iou_thres,trtFlag_det = processPar['half'], processPar['device'], processPar['conf_thres'], processPar['iou_thres'],processPar['trtFlag_det'] - if 'score_byClass' in processPar.keys(): score_byClass = processPar['score_byClass'] - else: score_byClass = None - - iou2nd = processPar['iou2nd'] - time0=time.time() - - if trtFlag_det: - img, padInfos = img_pad(im0s[0], size=(640,640,3)) ;img = [img] - else: - img = [letterbox(x, 640, auto=True, stride=32)[0] for x in im0s];padInfos=None - img = np.stack(img, 0) - # Convert - 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 # 0 - 255 to 0.0 - 1.0 - - seg_pred = None;segFlag=False - time1=time.time() - pred = yolov5Trtforward(model,img) if trtFlag_det else model(img,augment=False)[0] - - time2=time.time() - - #p_result,timeOut = getDetections(datas,conf_thres, iou_thres,names,label_arraylist,rainbows,10,ObjectPar=ObjectPar,font=font,padInfos=padInfos) - p_result, timeOut = getDetectionsFromPreds(pred,img,im0s[0],conf_thres=conf_thres,iou_thres=iou_thres,ovlap_thres=iou2nd,padInfos=padInfos) - if score_byClass: - p_result[2] = score_filter_byClass(p_result[2],score_byClass) - if segmodel: - seg_pred,segstr = segmodel.eval(im0s[0] ) - segFlag=True - else: - seg_pred = None;segFlag=False;segstr='No segmodel' - - - if segPar and segPar['mixFunction']['function']: - mixFunction = segPar['mixFunction']['function'] - - H,W = im0s[0].shape[0:2] - parMix = segPar['mixFunction']['pars'];#print('###line117:',parMix,p_result[2]) - parMix['imgSize'] = (W,H) - - - p_result[2],timeInfos_post = mixFunction(p_result[2], seg_pred, pars=parMix ) - timeInfos_seg_post = 'segInfer:%s ,postMixProcess:%s'%( segstr, timeInfos_post ) - else: - timeInfos_seg_post = ' ' - ''' - if segmodel: - timeS1=time.time() - #seg_pred,segstr = segtrtEval(segmodel,im0s[0],par=segPar) if segPar['trtFlag_seg'] else segmodel.eval(im0s[0] ) - seg_pred,segstr = segmodel.eval(im0s[0] ) - timeS2=time.time() - mixFunction = segPar['mixFunction']['function'] - - p_result[2],timeInfos_post = mixFunction(p_result[2], seg_pred, pars=segPar['mixFunction']['pars'] ) - - timeInfos_seg_post = 'segInfer:%.1f ,postProcess:%s'%( (timeS2-timeS1)*1000, timeInfos_post ) - - else: - timeInfos_seg_post = ' ' - #print('######line341:',seg_pred.shape,np.max(seg_pred),np.min(seg_pred) , len(p_result[2]) ) - ''' - time_info = 'letterbox:%.1f, detinfer:%.1f, '%( (time1-time0)*1000,(time2-time1)*1000 ) - - if sort_tracker: - #在这里增加设置调用追踪器的频率 - #..................USE TRACK FUNCTION.................... - #pass an empty array to sort - dets_to_sort = np.empty((0,7), dtype=np.float32) - - # NOTE: We send in detected object class too - #for detclass,x1,y1,x2,y2,conf in p_result[2]: - for x1,y1,x2,y2,conf, detclass in p_result[2]: - #print('#######line342:',x1,y1,x2,y2,img.shape,[x1, y1, x2, y2, conf, detclass,iframe]) - dets_to_sort = np.vstack((dets_to_sort, - np.array([x1, y1, x2, y2, conf, detclass,iframe],dtype=np.float32) )) - - # Run SORT - tracked_dets = deepcopy(sort_tracker.update(dets_to_sort) ) - tracks =sort_tracker.getTrackers() - p_result.append(tracked_dets) ###index=4 - p_result.append(tracks) ###index=5 - - return p_result,time_info+timeOut+timeInfos_seg_post -def AI_det_track_batch(imgarray_list, iframe_list ,modelPar,processPar,sort_tracker,trackPar,segPar=None): - ''' - 输入： - imgarray_list--图像列表 - iframe_list -- 帧号列表 - modelPar--模型参数,字典,modelPar={'det_Model':,'seg_Model':} - processPar--字典，存放检测相关参数，'half', 'device', 'conf_thres', 'iou_thres','trtFlag_det' - sort_tracker--对象，初始化的跟踪对象。为了保持一致，即使是单帧也要有。 - trackPar--跟踪参数，关键字包括：det_cnt，windowsize - segPar--None,分割模型相关参数。如果用不到，则为None - 输入：retResults,timeInfos - retResults：list - retResults[0]--imgarray_list - retResults[1]--所有结果用numpy格式，所有的检测结果，包括8类，每列分别是x1, y1, x2, y2, conf, detclass,iframe,trackId - retResults[2]--所有结果用list表示,其中每一个元素为一个list，表示每一帧的检测结果，每一个结果是由多个list构成，每个list表示一个框，格式为[ x0 ,y0 ,x1 ,y1 ,conf, cls ,ifrmae,trackId ]，如 retResults[2][j][k]表示第j帧的第k个框。2023.08.03,修改输出格式 - ''' - - det_cnt,windowsize = trackPar['det_cnt'] ,trackPar['windowsize'] - trackers_dic={} - index_list = list(range( 0, len(iframe_list) ,det_cnt )); - if len(index_list)>1 and index_list[-1]!= iframe_list[-1]: - index_list.append( len(iframe_list) - 1 ) - - if len(imgarray_list)==1: #如果是单帧图片，则不用跟踪 - retResults = [] - p_result,timeOut = AI_det_track( [ [imgarray_list[0]] ,iframe_list[0] ],modelPar,processPar,None,segPar ) - ##下面4行内容只是为了保持格式一致 - detArray = np.array(p_result[2]) - #print('##line371:',detArray) - if len(p_result[2])==0:res=[] - else: - cnt = detArray.shape[0];trackIds=np.zeros((cnt,1));iframes = np.zeros((cnt,1)) + iframe_list[0] - - #detArray = np.hstack( (detArray[:,1:5], detArray[:,5:6] ,detArray[:,0:1],iframes, trackIds ) ) - detArray = np.hstack( (detArray[:,0:4], detArray[:,4:6] ,iframes, trackIds ) ) ##2023.08.03 修改输入格式 - res = [[ b[0],b[1],b[2],b[3],b[4],b[5],b[6],b[7] ] for b in detArray ] - retResults=[imgarray_list,detArray,res ] - #print('##line380:',retResults[2]) - return retResults,timeOut - - else: - t0 = time.time() - timeInfos_track='' - for iframe_index, index_frame in enumerate(index_list): - p_result,timeOut = AI_det_track( [ [imgarray_list[index_frame]] ,iframe_list[index_frame] ],modelPar,processPar,sort_tracker,segPar ) - timeInfos_track='%s:%s'%(timeInfos_track,timeOut) - - for tracker in p_result[5]: - trackers_dic[tracker.id]=deepcopy(tracker) - t1 = time.time() - - track_det_result = np.empty((0,8)) - for trackId in trackers_dic.keys(): - tracker = trackers_dic[trackId] - bbox_history = np.array(tracker.bbox_history) - if len(bbox_history)<2: continue - ###把(x0,y0,x1,y1)转换成(xc,yc,w,h) - xcs_ycs = (bbox_history[:,0:2] + bbox_history[:,2:4] )/2 - whs = bbox_history[:,2:4] - bbox_history[:,0:2] - bbox_history[:,0:2] = xcs_ycs;bbox_history[:,2:4] = whs; - - arrays_box = bbox_history[:,0:7].transpose();frames=bbox_history[:,6] - #frame_min--表示该批次图片的起始帧，如该批次是[1,100],则frame_min=1，[101,200]--frame_min=101 - #frames[0]--表示该目标出现的起始帧，如[1,11,21,31,41],则frames[0]=1，frames[0]可能会在frame_min之前出现，即一个横跨了多个批次。 - - ##如果要最好化插值范围，则取内区间[frame_min,则frame_max ]和[frames[0],frames[-1] ]的交集 - #inter_frame_min = int(max(frame_min, frames[0])); inter_frame_max = int(min( frame_max, frames[-1] )) ## - - ##如果要求得到完整的目标轨迹，则插值区间要以目标出现的起始点为准 - inter_frame_min=int(frames[0]);inter_frame_max=int(frames[-1]) - new_frames= np.linspace(inter_frame_min,inter_frame_max,inter_frame_max-inter_frame_min+1 ) - f_linear = interpolate.interp1d(frames,arrays_box); interpolation_x0s = (f_linear(new_frames)).transpose() - move_cnt_use =(len(interpolation_x0s)+1)//2*2-1 if len(interpolation_x0s)1 and index_list[-1]!= iframe_list[-1]: - index_list.append( len(iframe_list) - 1 ) - - if len(imgarray_list)==1: #如果是单帧图片，则不用跟踪 - retResults = [] - p_result,timeOut = AI_det_track_N( [ [imgarray_list[0]] ,iframe_list[0] ],modelList,postProcess,None ) - ##下面4行内容只是为了保持格式一致 - detArray = np.array(p_result[2]) - if len(p_result[2])==0:res=[] - else: - cnt = detArray.shape[0];trackIds=np.zeros((cnt,1));iframes = np.zeros((cnt,1)) + iframe_list[0] - - #detArray = np.hstack( (detArray[:,1:5], detArray[:,5:6] ,detArray[:,0:1],iframes, trackIds ) ) - detArray = np.hstack( (detArray[:,0:4], detArray[:,4:6] ,iframes, trackIds ) ) ##2023.08.03 修改输入格式 - res = [[ b[0],b[1],b[2],b[3],b[4],b[5],b[6],b[7] ] for b in detArray ] - retResults=[imgarray_list,detArray,res ] - #print('##line380:',retResults[2]) - return retResults,timeOut - - else: - t0 = time.time() - timeInfos_track='' - for iframe_index, index_frame in enumerate(index_list): - p_result,timeOut = AI_det_track_N( [ [imgarray_list[index_frame]] ,iframe_list[index_frame] ],modelList,postProcess,sort_tracker ) - timeInfos_track='%s:%s'%(timeInfos_track,timeOut) - - for tracker in p_result[5]: - trackers_dic[tracker.id]=deepcopy(tracker) - t1 = time.time() - - track_det_result = np.empty((0,8)) - for trackId in trackers_dic.keys(): - tracker = trackers_dic[trackId] - bbox_history = np.array(tracker.bbox_history).copy() - if len(bbox_history)<2: continue - ###把(x0,y0,x1,y1)转换成(xc,yc,w,h) - xcs_ycs = (bbox_history[:,0:2] + bbox_history[:,2:4] )/2 - whs = bbox_history[:,2:4] - bbox_history[:,0:2] - bbox_history[:,0:2] = xcs_ycs;bbox_history[:,2:4] = whs; - - #2023.11.17添加的。目的是修正跟踪链上所有的框的类别一样 - chainClsId = get_tracker_cls(bbox_history,scId=4,clsId=5) - bbox_history[:,5] = chainClsId - - arrays_box = bbox_history[:,0:7].transpose();frames=bbox_history[:,6] - #frame_min--表示该批次图片的起始帧，如该批次是[1,100],则frame_min=1，[101,200]--frame_min=101 - #frames[0]--表示该目标出现的起始帧，如[1,11,21,31,41],则frames[0]=1，frames[0]可能会在frame_min之前出现，即一个横跨了多个批次。 - - ##如果要最好化插值范围，则取内区间[frame_min,则frame_max ]和[frames[0],frames[-1] ]的交集 - #inter_frame_min = int(max(frame_min, frames[0])); inter_frame_max = int(min( frame_max, frames[-1] )) ## - - ##如果要求得到完整的目标轨迹，则插值区间要以目标出现的起始点为准 - inter_frame_min=int(frames[0]);inter_frame_max=int(frames[-1]) - new_frames= np.linspace(inter_frame_min,inter_frame_max,inter_frame_max-inter_frame_min+1 ) - f_linear = interpolate.interp1d(frames,arrays_box); interpolation_x0s = (f_linear(new_frames)).transpose() - move_cnt_use =(len(interpolation_x0s)+1)//2*2-1 if len(interpolation_x0s)maxScore: + maxId=i + maxScore = detRet[4] + return maxId + +def AI_process_C(im0s,modelList,postProcess): + #函数定制的原因： + ## 之前模型处理流是 + ## 图片---> 模型1-->result1；图片---> 模型2->result2；[result1,result2]--->后处理函数 + ## 本函数的处理流程是 + ## 图片---> 模型1-->result1；[图片,result1]---> 模型2->result2；[result1,result2]--->后处理函数 + ## 模型2的输入，是有模型1的输出决定的。如模型2是ocr模型，需要将模型1检测出来的船名抠图出来输入到模型2. + ## 之前的模型流都是模型2是分割模型，输入就是原始图片，与模型1的输出无关。 + #输入参数 + ## im0s---原始图像列表 + ## modelList--所有的模型 + # postProcess--字典{},包括后处理函数，及其参数 + #输出参数 + ##ret[0]--检测结果； + ##ret[1]--时间信息 + + #modelList包括模型，每个模型是一个类，里面的eval函数可以输出该模型的推理结果 + + t0=time.time() + detRets0 = modelList[0].eval(im0s[0]) + + #detRets0=[[12, 46, 1127, 1544, 0.2340087890625, 2.0], [1884, 1248, 2992, 1485, 0.64208984375, 1.0]] + detRets0 = detRets0[0] + parsIn=postProcess['pars'] + + _detRets0_obj = list(filter(lambda x: x[5] in parsIn['objs'], detRets0 )) + _detRets0_others = list(filter(lambda x: x[5] not in parsIn['objs'], detRets0 )) + _detRets0 = [] + if postProcess['name']=='channel2': + if len(_detRets0_obj)>0: + maxId=getMaxScoreWords(_detRets0_obj) + _detRets0 = _detRets0_obj[maxId:maxId+1] + else: _detRets0 = detRets0 + + + t1=time.time() + imagePatches = [ im0s[0][int(x[1]):int(x[3] ) ,int(x[0]):int(x[2])] for x in _detRets0 ] + detRets1 = [modelList[1].eval(patch) for patch in imagePatches] + print('###line240:',detRets1) + if postProcess['name']=='crackMeasurement': + detRets1 = [x[0]*255 for x in detRets1] + t2=time.time() + mixFunction =postProcess['function'] + crackInfos = [mixFunction(patchMask,par=parsIn) for patchMask in detRets1] + + 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': + H,W = im0s[0].shape[0:2];parsIn['imgSize'] = (W,H) + mixFunction =postProcess['function'] + _detRets0_others = mixFunction([_detRets0_others], parsIn) + ocrInfo='no ocr' + if len(_detRets0_obj)>0: + res_real = detRets1[0][0] + res_real="".join( list(filter(lambda x:(ord(x) >19968 and ord(x)<63865 ) or (ord(x) >47 and ord(x)<58 ),res_real))) + + #detRets1[0][0]="".join( list(filter(lambda x:(ord(x) >19968 and ord(x)<63865 ) or (ord(x) >47 and ord(x)<58 ),detRets1[0][0]))) + _detRets0_obj[maxId].append(res_real ) + _detRets0_obj = [_detRets0_obj[maxId]]##只输出有OCR的那个船名结果 + ocrInfo=detRets1[0][1] + print( ' _detRets0_obj:{} _detRets0_others:{} '.format( _detRets0_obj, _detRets0_others ) ) + rets=_detRets0_obj+_detRets0_others + t3=time.time() + outInfos='total:%.1f ,where det:%.1f, ocr:%s'%( (t3-t0)*1000, (t1-t0)*1000, ocrInfo) + + #print('###line233:',detRets1,detRets0 ) + + 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---原始图像列表 + # model---检测模型，segmodel---分割模型（如若没有用到，则为None） + #输出：两个元素（列表，字符）构成的元组，[im0s[0],im0,det_xywh,iframe],strout + # [im0s[0],im0,det_xywh,iframe]中， + # im0s[0]--原始图像，im0--AI处理后的图像,iframe--帧号/暂时不需用到。 + # det_xywh--检测结果，是一个列表。 + # 其中每一个元素表示一个目标构成如：[ xc,yc,w,h, float(conf_c),float(cls_c)],#2023.08.03,修改输出格式 + # #cls_c--类别，如0,1,2,3; xc,yc,w,h--中心点坐标及宽；conf_c--得分, 取值范围在0-1之间 + # #strout---统计AI处理个环节的时间 + + # Letterbox + time0=time.time() + if trtFlag_det: + img, padInfos = img_pad(im0s[0], size=(640,640,3)) ;img = [img] + else: + img = [letterbox(x, 640, auto=True, stride=32)[0] for x in im0s];padInfos=None + #img = [letterbox(x, 640, auto=True, stride=32)[0] for x in im0s] + # Stack + img = np.stack(img, 0) + # Convert + 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 # 0 - 255 to 0.0 - 1.0 + if segmodel: + seg_pred,segstr = segmodel.eval(im0s[0] ) + segFlag=True + else: + seg_pred = None;segFlag=False + time1=time.time() + pred = yolov5Trtforward(model,img) if trtFlag_det else model(img,augment=False)[0] + + + time2=time.time() + datas = [[''], img, im0s, None,pred,seg_pred,10] + + ObjectPar={ 'object_config':allowedList, 'slopeIndex':[] ,'segmodel':segFlag,'segRegionCnt':0 } + p_result,timeOut = post_process_(datas,conf_thres, iou_thres,names,label_arraylist,rainbows,10,ObjectPar=ObjectPar,font=font,padInfos=padInfos,ovlap_thres=SecNms) + #print('###line274:',p_result[2]) + #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'] + segmodel = modelPar['seg_Model'] + half,device,conf_thres, iou_thres,trtFlag_det = processPar['half'], processPar['device'], processPar['conf_thres'], processPar['iou_thres'],processPar['trtFlag_det'] + if 'score_byClass' in processPar.keys(): score_byClass = processPar['score_byClass'] + else: score_byClass = None + + iou2nd = processPar['iou2nd'] + time0=time.time() + + if trtFlag_det: + img, padInfos = img_pad(im0s[0], size=(640,640,3)) ;img = [img] + else: + img = [letterbox(x, 640, auto=True, stride=32)[0] for x in im0s];padInfos=None + img = np.stack(img, 0) + # Convert + 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 # 0 - 255 to 0.0 - 1.0 + + seg_pred = None;segFlag=False + time1=time.time() + pred = yolov5Trtforward(model,img) if trtFlag_det else model(img,augment=False)[0] + + time2=time.time() + + #p_result,timeOut = getDetections(datas,conf_thres, iou_thres,names,label_arraylist,rainbows,10,ObjectPar=ObjectPar,font=font,padInfos=padInfos) + p_result, timeOut = getDetectionsFromPreds(pred,img,im0s[0],conf_thres=conf_thres,iou_thres=iou_thres,ovlap_thres=iou2nd,padInfos=padInfos) + if score_byClass: + p_result[2] = score_filter_byClass(p_result[2],score_byClass) + if segmodel: + seg_pred,segstr = segmodel.eval(im0s[0] ) + segFlag=True + else: + seg_pred = None;segFlag=False;segstr='No segmodel' + + + if segPar and segPar['mixFunction']['function']: + mixFunction = segPar['mixFunction']['function'] + + H,W = im0s[0].shape[0:2] + parMix = segPar['mixFunction']['pars'];#print('###line117:',parMix,p_result[2]) + parMix['imgSize'] = (W,H) + + + p_result[2],timeInfos_post = mixFunction(p_result[2], seg_pred, pars=parMix ) + timeInfos_seg_post = 'segInfer:%s ,postMixProcess:%s'%( segstr, timeInfos_post ) + else: + timeInfos_seg_post = ' ' + ''' + if segmodel: + timeS1=time.time() + #seg_pred,segstr = segtrtEval(segmodel,im0s[0],par=segPar) if segPar['trtFlag_seg'] else segmodel.eval(im0s[0] ) + seg_pred,segstr = segmodel.eval(im0s[0] ) + timeS2=time.time() + mixFunction = segPar['mixFunction']['function'] + + p_result[2],timeInfos_post = mixFunction(p_result[2], seg_pred, pars=segPar['mixFunction']['pars'] ) + + timeInfos_seg_post = 'segInfer:%.1f ,postProcess:%s'%( (timeS2-timeS1)*1000, timeInfos_post ) + + else: + timeInfos_seg_post = ' ' + #print('######line341:',seg_pred.shape,np.max(seg_pred),np.min(seg_pred) , len(p_result[2]) ) + ''' + time_info = 'letterbox:%.1f, detinfer:%.1f, '%( (time1-time0)*1000,(time2-time1)*1000 ) + + if sort_tracker: + #在这里增加设置调用追踪器的频率 + #..................USE TRACK FUNCTION.................... + #pass an empty array to sort + dets_to_sort = np.empty((0,7), dtype=np.float32) + + # NOTE: We send in detected object class too + #for detclass,x1,y1,x2,y2,conf in p_result[2]: + for x1,y1,x2,y2,conf, detclass in p_result[2]: + #print('#######line342:',x1,y1,x2,y2,img.shape,[x1, y1, x2, y2, conf, detclass,iframe]) + dets_to_sort = np.vstack((dets_to_sort, + np.array([x1, y1, x2, y2, conf, detclass,iframe],dtype=np.float32) )) + + # Run SORT + tracked_dets = deepcopy(sort_tracker.update(dets_to_sort) ) + tracks =sort_tracker.getTrackers() + p_result.append(tracked_dets) ###index=4 + p_result.append(tracks) ###index=5 + + return p_result,time_info+timeOut+timeInfos_seg_post +def AI_det_track_batch(imgarray_list, iframe_list ,modelPar,processPar,sort_tracker,trackPar,segPar=None): + ''' + 输入： + imgarray_list--图像列表 + iframe_list -- 帧号列表 + modelPar--模型参数,字典,modelPar={'det_Model':,'seg_Model':} + processPar--字典，存放检测相关参数，'half', 'device', 'conf_thres', 'iou_thres','trtFlag_det' + sort_tracker--对象，初始化的跟踪对象。为了保持一致，即使是单帧也要有。 + trackPar--跟踪参数，关键字包括：det_cnt，windowsize + segPar--None,分割模型相关参数。如果用不到，则为None + 输入：retResults,timeInfos + retResults：list + retResults[0]--imgarray_list + retResults[1]--所有结果用numpy格式，所有的检测结果，包括8类，每列分别是x1, y1, x2, y2, conf, detclass,iframe,trackId + retResults[2]--所有结果用list表示,其中每一个元素为一个list，表示每一帧的检测结果，每一个结果是由多个list构成，每个list表示一个框，格式为[ x0 ,y0 ,x1 ,y1 ,conf, cls ,ifrmae,trackId ]，如 retResults[2][j][k]表示第j帧的第k个框。2023.08.03,修改输出格式 + ''' + + det_cnt,windowsize = trackPar['det_cnt'] ,trackPar['windowsize'] + trackers_dic={} + index_list = list(range( 0, len(iframe_list) ,det_cnt )); + if len(index_list)>1 and index_list[-1]!= iframe_list[-1]: + index_list.append( len(iframe_list) - 1 ) + + if len(imgarray_list)==1: #如果是单帧图片，则不用跟踪 + retResults = [] + p_result,timeOut = AI_det_track( [ [imgarray_list[0]] ,iframe_list[0] ],modelPar,processPar,None,segPar ) + ##下面4行内容只是为了保持格式一致 + detArray = np.array(p_result[2]) + #print('##line371:',detArray) + if len(p_result[2])==0:res=[] + else: + cnt = detArray.shape[0];trackIds=np.zeros((cnt,1));iframes = np.zeros((cnt,1)) + iframe_list[0] + + #detArray = np.hstack( (detArray[:,1:5], detArray[:,5:6] ,detArray[:,0:1],iframes, trackIds ) ) + detArray = np.hstack( (detArray[:,0:4], detArray[:,4:6] ,iframes, trackIds ) ) ##2023.08.03 修改输入格式 + res = [[ b[0],b[1],b[2],b[3],b[4],b[5],b[6],b[7] ] for b in detArray ] + retResults=[imgarray_list,detArray,res ] + #print('##line380:',retResults[2]) + return retResults,timeOut + + else: + t0 = time.time() + timeInfos_track='' + for iframe_index, index_frame in enumerate(index_list): + p_result,timeOut = AI_det_track( [ [imgarray_list[index_frame]] ,iframe_list[index_frame] ],modelPar,processPar,sort_tracker,segPar ) + timeInfos_track='%s:%s'%(timeInfos_track,timeOut) + + for tracker in p_result[5]: + trackers_dic[tracker.id]=deepcopy(tracker) + t1 = time.time() + + track_det_result = np.empty((0,8)) + for trackId in trackers_dic.keys(): + tracker = trackers_dic[trackId] + bbox_history = np.array(tracker.bbox_history) + if len(bbox_history)<2: continue + ###把(x0,y0,x1,y1)转换成(xc,yc,w,h) + xcs_ycs = (bbox_history[:,0:2] + bbox_history[:,2:4] )/2 + whs = bbox_history[:,2:4] - bbox_history[:,0:2] + bbox_history[:,0:2] = xcs_ycs;bbox_history[:,2:4] = whs; + + arrays_box = bbox_history[:,0:7].transpose();frames=bbox_history[:,6] + #frame_min--表示该批次图片的起始帧，如该批次是[1,100],则frame_min=1，[101,200]--frame_min=101 + #frames[0]--表示该目标出现的起始帧，如[1,11,21,31,41],则frames[0]=1，frames[0]可能会在frame_min之前出现，即一个横跨了多个批次。 + + ##如果要最好化插值范围，则取内区间[frame_min,则frame_max ]和[frames[0],frames[-1] ]的交集 + #inter_frame_min = int(max(frame_min, frames[0])); inter_frame_max = int(min( frame_max, frames[-1] )) ## + + ##如果要求得到完整的目标轨迹，则插值区间要以目标出现的起始点为准 + inter_frame_min=int(frames[0]);inter_frame_max=int(frames[-1]) + new_frames= np.linspace(inter_frame_min,inter_frame_max,inter_frame_max-inter_frame_min+1 ) + f_linear = interpolate.interp1d(frames,arrays_box); interpolation_x0s = (f_linear(new_frames)).transpose() + move_cnt_use =(len(interpolation_x0s)+1)//2*2-1 if len(interpolation_x0s)1 and index_list[-1]!= iframe_list[-1]: + index_list.append( len(iframe_list) - 1 ) + + if len(imgarray_list)==1: #如果是单帧图片，则不用跟踪 + retResults = [] + p_result,timeOut = AI_det_track_N( [ [imgarray_list[0]] ,iframe_list[0] ],modelList,postProcess,None ) + ##下面4行内容只是为了保持格式一致 + detArray = np.array(p_result[2]) + if len(p_result[2])==0:res=[] + else: + cnt = detArray.shape[0];trackIds=np.zeros((cnt,1));iframes = np.zeros((cnt,1)) + iframe_list[0] + + #detArray = np.hstack( (detArray[:,1:5], detArray[:,5:6] ,detArray[:,0:1],iframes, trackIds ) ) + detArray = np.hstack( (detArray[:,0:4], detArray[:,4:6] ,iframes, trackIds ) ) ##2023.08.03 修改输入格式 + res = [[ b[0],b[1],b[2],b[3],b[4],b[5],b[6],b[7] ] for b in detArray ] + retResults=[imgarray_list,detArray,res ] + #print('##line380:',retResults[2]) + return retResults,timeOut + + else: + t0 = time.time() + timeInfos_track='' + for iframe_index, index_frame in enumerate(index_list): + p_result,timeOut = AI_det_track_N( [ [imgarray_list[index_frame]] ,iframe_list[index_frame] ],modelList,postProcess,sort_tracker ) + timeInfos_track='%s:%s'%(timeInfos_track,timeOut) + + for tracker in p_result[5]: + trackers_dic[tracker.id]=deepcopy(tracker) + t1 = time.time() + + track_det_result = np.empty((0,8)) + for trackId in trackers_dic.keys(): + tracker = trackers_dic[trackId] + bbox_history = np.array(tracker.bbox_history).copy() + if len(bbox_history)<2: continue + ###把(x0,y0,x1,y1)转换成(xc,yc,w,h) + xcs_ycs = (bbox_history[:,0:2] + bbox_history[:,2:4] )/2 + whs = bbox_history[:,2:4] - bbox_history[:,0:2] + bbox_history[:,0:2] = xcs_ycs;bbox_history[:,2:4] = whs; + + #2023.11.17添加的。目的是修正跟踪链上所有的框的类别一样 + chainClsId = get_tracker_cls(bbox_history,scId=4,clsId=5) + bbox_history[:,5] = chainClsId + + arrays_box = bbox_history[:,0:7].transpose();frames=bbox_history[:,6] + #frame_min--表示该批次图片的起始帧，如该批次是[1,100],则frame_min=1，[101,200]--frame_min=101 + #frames[0]--表示该目标出现的起始帧，如[1,11,21,31,41],则frames[0]=1，frames[0]可能会在frame_min之前出现，即一个横跨了多个批次。 + + ##如果要最好化插值范围，则取内区间[frame_min,则frame_max ]和[frames[0],frames[-1] ]的交集 + #inter_frame_min = int(max(frame_min, frames[0])); inter_frame_max = int(min( frame_max, frames[-1] )) ## + + ##如果要求得到完整的目标轨迹，则插值区间要以目标出现的起始点为准 + inter_frame_min=int(frames[0]);inter_frame_max=int(frames[-1]) + new_frames= np.linspace(inter_frame_min,inter_frame_max,inter_frame_max-inter_frame_min+1 ) + f_linear = interpolate.interp1d(frames,arrays_box); interpolation_x0s = (f_linear(new_frames)).transpose() + move_cnt_use =(len(interpolation_x0s)+1)//2*2-1 if len(interpolation_x0s) 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' + + ##以下参数目前不可改 + Detweights = "weights/yolov5/class5/best_5classes.pt" + seg_nclass = 2 + Segweights = "weights/BiSeNet/checkpoint.pth" + conf_thres,iou_thres,classes= 0.25,0.45,5 + labelnames = "weights/yolov5/class5/labelnames.json" + rainbows = [ [0,0,255],[0,255,0],[255,0,0],[255,0,255],[255,255,0],[255,129,0],[255,0,127],[127,255,0],[0,255,127],[0,127,255],[127,0,255],[255,127,255],[255,255,127],[127,255,255],[0,255,255],[255,127,255],[127,255,255], [0,127,0],[0,0,127],[0,255,255]] + allowedList=[0,1,2,3] + + + ##加载模型，准备好显示字符 + device = select_device(device_) + names=get_labelnames(labelnames) + label_arraylist = get_label_arrays(names,rainbows,outfontsize=40,fontpath="conf/platech.ttf") + half = device.type != 'cpu' # half precision only supported on CUDA + model = attempt_load(Detweights, map_location=device) # load FP32 model + if half: model.half() + segmodel = SegModel(nclass=seg_nclass,weights=Segweights,device=device) + + + ##图像测试 + #url='images/examples/20220624_响水河_12300_1621.jpg' + impth = 'images/examples/' + outpth = 'images/results/' + folders = os.listdir(impth) + for i in range(len(folders)): + imgpath = os.path.join(impth, folders[i]) + im0s=[cv2.imread(imgpath)] + time00 = time.time() + p_result,timeOut = AI_process(im0s,model,segmodel,names,label_arraylist,rainbows,half,device,conf_thres, iou_thres,allowedList,fontSize=1.0) + time11 = time.time() + image_array = p_result[1] + cv2.imwrite( os.path.join( outpth,folders[i] ) ,image_array ) + #print('----process:%s'%(folders[i]), (time.time() - time11) * 1000) + + + + + +if __name__=="__main__": + main()