import sys, yaml from easydict import EasyDict as edict from concurrent.futures import ThreadPoolExecutor sys.path.extend(['..','../AIlib2' ]) from AI import AI_process,AI_process_forest,get_postProcess_para,get_postProcess_para_dic,ocr_process,AI_det_track,AI_det_track_batch import cv2,os,time from segutils.segmodel import SegModel from stdc import stdcModel from segutils.trafficUtils import tracfficAccidentMixFunction from models.experimental import attempt_load from utils.torch_utils import select_device from utilsK.queRiver import get_labelnames,get_label_arrays,save_problem_images,riverDetSegMixProcess from ocrUtils.ocrUtils import CTCLabelConverter,AlignCollate from trackUtils.sort import Sort,track_draw_boxAndTrace,track_draw_trace_boxes,moving_average_wang,drawBoxTraceSimplied from trackUtils.sort_obb import OBB_Sort,obbTohbb,track_draw_all_boxes,track_draw_trace from obbUtils.shipUtils import OBB_infer,OBB_tracker,draw_obb,OBB_tracker_batch from utilsK.noParkingUtils import mixNoParking_road_postprocess from obbUtils.load_obb_model import load_model_decoder_OBB import numpy as np import torch,glob import tensorrt as trt from utilsK.masterUtils import get_needed_objectsIndex from copy import deepcopy from scipy import interpolate from utilsK.drownUtils import mixDrowing_water_postprocess #import warnings #warnings.filterwarnings("error") def view_bar(num, total,time1,prefix='prefix'): rate = num / total time_n=time.time() rate_num = int(rate * 30) rate_nums = np.round(rate * 100) r = '\r %s %d / %d [%s%s] %.2f s'%(prefix,num,total, ">" * rate_num, " " * (30 - rate_num), time_n-time1 ) sys.stdout.write(r) sys.stdout.flush() ''' 多线程 ''' def process_v1(frame): #try: print('demo.py beging to :',frame[8]) time00 = time.time() H,W,C = frame[0][0].shape p_result,timeOut = AI_process(frame[0],frame[1],frame[2],frame[3],frame[4],frame[5],objectPar=frame[6],font=frame[7],segPar=frame[9],mode=frame[10],postPar=frame[11]) time11 = time.time() image_array = p_result[1] cv2.imwrite(os.path.join('images/results/',frame[8] ) ,image_array) bname = frame[8].split('.')[0] if len(p_result)==5: image_mask = p_result[4] cv2.imwrite(os.path.join('images/results/',bname+'_mask.png' ) , (image_mask).astype(np.uint8)) boxes=p_result[2] with open( os.path.join('images/results/',bname+'.txt' ),'w' ) as fp: for box in boxes: box_str=[str(x) for x in box] out_str=','.join(box_str)+'\n' fp.write(out_str) time22 = time.time() print('%s,%d*%d,AI-process: %.1f,image save:%.1f , %s'%(frame[8],H,W, (time11 - time00) * 1000.0, (time22-time11)*1000.0,timeOut), boxes) return 'success' #except Exception as e: # return 'failed:'+str(e) def process_video(video,par0,mode='detSeg'): cap=cv2.VideoCapture(video) if not cap.isOpened(): print('#####error url:',video) return False bname=os.path.basename(video).split('.')[0] fps = int(cap.get(cv2.CAP_PROP_FPS)+0.5) width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH )+0.5) height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)+0.5) framecnt=int(cap.get(7)+0.5) save_path_AI = os.path.join(par0['outpth'],os.path.basename(video)) problem_image_dir= os.path.join( par0['outpth'], 'probleImages' ) os.makedirs(problem_image_dir,exist_ok=True) vid_writer_AI = cv2.VideoWriter(save_path_AI, cv2.VideoWriter_fourcc(*'mp4v'), fps, (width,height)) num=0 iframe=0;post_results=[];fpsample=30*10 imgarray_list = []; iframe_list = [] patch_cnt = par0['trackPar']['patchCnt'] ##windowsize 对逐帧插值后的结果做平滑，windowsize为平滑的长度,没隔det_cnt帧做一次跟踪。 trackPar={'det_cnt':10,'windowsize':29 } ##track_det_result_update= np.empty((0,8)) ###每100帧跑出来的结果，放在track_det_result_update，只保留当前100帧里有的tracker Id. while cap.isOpened(): ret, imgarray = cap.read() #读取摄像头画面 iframe +=1 if not ret:break if mode=='detSeg': p_result,timeOut = AI_process([imgarray],par0['model'],par0['segmodel'],par0['names'],par0['label_arraylist'],par0['rainbows'],objectPar=par0['objectPar'],font=par0['digitFont'],segPar=par0['segPar']) elif mode == 'track': #sampleCount=10 imgarray_list.append( imgarray ) iframe_list.append(iframe ) if iframe%patch_cnt==0: time_patch0 = time.time() retResults,timeInfos = AI_det_track_batch(imgarray_list, iframe_list ,par0['modelPar'],par0['processPar'],par0['sort_tracker'] ,par0['trackPar'],segPar=par0['segPar']) #print('###line111:',retResults[2]) ###需要保存成一个二维list，每一个list是一帧检测结果。 ###track_det_result 内容格式：x1, y1, x2, y2, conf, cls,iframe,trackId time_patch2 = time.time() frame_min = iframe_list[0];frame_max=iframe_list[-1] for iiframe in range(frame_min,frame_max+1): img_draw = imgarray_list[ iiframe- frame_min ] img_draw = drawBoxTraceSimplied(retResults[1] ,iiframe, img_draw,rainbows=par0['drawPar']['rainbows'],boxFlag=True,traceFlag=True,names=par0['drawPar']['names'] ) ret = vid_writer_AI.write(img_draw) view_bar(iiframe, framecnt,time.time(),prefix=os.path.basename(video)) imgarray_list=[];iframe_list=[] elif mode =='obbTrack': imgarray_list.append( imgarray ) iframe_list.append(iframe ) if iframe%patch_cnt==0: time_patch0 = time.time() track_det_results, timeInfos = OBB_tracker_batch(imgarray_list,iframe_list,par0['modelPar'],par0['obbModelPar'],par0['sort_tracker'],par0['trackPar'],segPar=None) print( timeInfos ) #对结果画图 track_det_np = track_det_results[1] frame_min = iframe_list[0];frame_max=iframe_list[-1] for iiframe in range(frame_min,frame_max+1): img_draw = imgarray_list[ iiframe- frame_min ] if len( track_det_results[2][ iiframe- frame_min]) > 0: img_draw = draw_obb( track_det_results[2][iiframe- frame_min ] ,img_draw,par0['drawPar']) if True: frameIdex=12;trackIdex=13; boxes_oneFrame = track_det_np[ track_det_np[:,frameIdex]==iiframe ] ###在某一帧上，画上轨迹 track_ids = boxes_oneFrame[:,trackIdex].tolist() boxes_before_oneFrame = track_det_np[ track_det_np[:,frameIdex]<=iiframe ] for trackId in track_ids: boxes_before_oneFrame_oneId = boxes_before_oneFrame[boxes_before_oneFrame[:,trackIdex]==trackId] xcs = boxes_before_oneFrame_oneId[:,8] ycs = boxes_before_oneFrame_oneId[:,9] [cv2.line(img_draw, ( int(xcs[i]) , int(ycs[i]) ), ( int(xcs[i+1]),int(ycs[i+1]) ),(255,0,0), thickness=2) for i,_ in enumerate(xcs) if i < len(xcs)-1 ] ret = vid_writer_AI.write(img_draw) #sys.exit(0) #print('vide writer ret:',ret) imgarray_list=[];iframe_list=[] view_bar(iframe, framecnt,time.time(),prefix=os.path.basename(video)) else: p_result,timeOut = AI_process_forest([imgarray],par0['model'],par0['segmodel'],par0['names'],par0['label_arraylist'],par0['rainbows'],par0['half'],par0['device'],par0['conf_thres'], par0['iou_thres'],par0['allowedList'],font=par0['digitFont'],trtFlag_det=par0['trtFlag_det']) if mode not in [ 'track','obbTrack']: image_array = p_result[1];num+=1 ret = vid_writer_AI.write(image_array) view_bar(num, framecnt,time.time(),prefix=os.path.basename(video)) ##每隔 fpsample帧处理一次，如果有问题就保存图片 if (iframe % fpsample == 0) and (len(post_results)>0) : parImage=save_problem_images(post_results,iframe,par0['names'],streamName=bname,outImaDir=problem_image_dir,imageTxtFile=False) post_results=[] if len(p_result[2] )>0: post_results.append(p_result) vid_writer_AI.release(); def det_track_demo(business ): ''' 跟踪参数说明： 'trackPar':{'sort_max_age':2,'sort_min_hits':3,'sort_iou_thresh':0.2,'det_cnt':10,'windowsize':29,'patchCnt':100} sort_max_age--跟踪链断裂时允许目标消失最大的次数。超过之后，会认为是新的目标。 sort_min_hits--每隔目标连续出现的次数，超过这个次数才认为是一个目标。 sort_iou_thresh--检测最小的置信度。 det_cnt--每隔几次做一个跟踪和检测，默认10。 windowsize--轨迹平滑长度，一定是奇数，表示每隔几帧做一平滑，默认29。 patchCnt--每次送入图像的数量，不宜少于100帧。 ''' ''' 以下是基于检测和分割的跟踪模型，分割用来修正检测的结果''' ####河道巡检的跟踪模型参数 if opt['business'] == 'river' or opt['business'] == 'river2' : par={ 'device':'0', ###显卡号,如果用TRT模型，只支持0（单显卡） 'labelnames':"../AIlib2/weights/conf/%s/labelnames.json"%(opt['business']), ###检测类别对照表 'gpuname':'2080Ti',###显卡名称 'max_workers':1, ###并行线程数 'half':True, 'Detweights':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###检测模型路径 'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,4,5,6] ],###控制哪些检测类别显示、输出 'seg_nclass':2,###分割模型类别数目，默认2类 'segRegionCnt':0,###分割模型结果需要保留的等值线数目 'segPar':{'modelSize':(640,360),'mean':(0.485, 0.456, 0.406),'std' :(0.229, 0.224, 0.225),'numpy':False, 'RGB_convert_first':True,#分割模型预处理参数 'mixFunction':{'function':riverDetSegMixProcess,'pars':{'slopeIndex':[1,3,4,7], 'riverIou':0.1}} #分割和检测混合处理的函数 }, 'trackPar':{'sort_max_age':2,'sort_min_hits':3,'sort_iou_thresh':0.2,'det_cnt':10,'windowsize':29,'patchCnt':100}, 'Segweights' : "../weights/%s/AIlib2/%s/stdc_360X640_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###分割模型权重位置 'postFile': '../AIlib2/weights/conf/%s/para.json'%( opt['business'] ),###后处理参数文件 'txtFontSize':80,###文本字符的大小 'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'waterLineWidth':3},###显示框、线设置 #'testImgPath':'images/videos/river',###测试图像的位置 'testImgPath':'images/tt',###测试图像的位置 'testOutPath':'images/results/',###输出测试图像位置 } if opt['business'] == 'highWay2': par={ 'device':'0', ###显卡号,如果用TRT模型，只支持0（单显卡） 'labelnames':"../AIlib2/weights/conf/%s/labelnames.json"%( opt['business'] ), ###检测类别对照表 'half':True, 'gpuname':'3090',###显卡名称 'max_workers':1, ###并行线程数 'Detweights':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###检测模型路径 #'Detweights':"../AIlib2/weights/conf/highWay2/yolov5.pt", 'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,5,6,7,8,9] ],###控制哪些检测类别显示、输出 'seg_nclass':3,###分割模型类别数目，默认2类 'segRegionCnt':2,###分割模型结果需要保留的等值线数目 'segPar':{'modelSize':(640,360),'mean':(0.485, 0.456, 0.406),'std' :(0.229, 0.224, 0.225),'predResize':True,'numpy':False, 'RGB_convert_first':True,###分割模型预处理参数 'mixFunction':{'function':tracfficAccidentMixFunction, 'pars':{ 'RoadArea': 16000, 'vehicleArea': 10, 'roadVehicleAngle': 15, 'speedRoadVehicleAngleMax': 75,'radius': 50 , 'roundness': 1.0, 'cls': 9, 'vehicleFactor': 0.1,'cls':9, 'confThres':0.25,'roadIou':0.6,'vehicleFlag':False,'distanceFlag': False } } }, 'trackPar':{'sort_max_age':2,'sort_min_hits':3,'sort_iou_thresh':0.2,'det_cnt':5,'windowsize':29,'patchCnt':100}, 'mode':'highWay3.0', 'Segweights' : "../weights/%s/AIlib2/%s/stdc_360X640_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###分割模型权重位置 'postFile': '../AIlib2/weights/conf/%s/para.json'%(opt['business'] ),###后处理参数文件 'txtFontSize':20,###文本字符的大小 'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':0.5,'waterLineColor':(0,255,255),'segLineShow':True,'waterLineWidth':2},###显示框、线设置 #'testImgPath':'images/trafficAccident/8.png',###测试图像的位置 'testImgPath':'/home/chenyukun/777-7-42.mp4',###测试图像的位置 'testOutPath':'images/results/',###输出测试图像位置 } par['segPar']['mixFunction']['pars']['modelSize'] = par['segPar']['modelSize'] if opt['business'] == 'noParking': par={ 'device':'0', ###显卡号,如果用TRT模型，只支持0（单显卡） 'labelnames':"../AIlib2/weights/conf/%s/labelnames.json"%( opt['business'] ), ###检测类别对照表 'half':True, 'gpuname':'3090',###显卡名称 'max_workers':1, ###并行线程数 'Detweights':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###检测模型路径 #'Detweights':"../AIlib2/weights/conf/highWay2/yolov5.pt", 'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,5,6,7,8,9] ],###控制哪些检测类别显示、输出 'seg_nclass':4,###分割模型类别数目，默认2类 'segRegionCnt':2,###分割模型结果需要保留的等值线数目 'segPar':{'modelSize':(640,360),'mean':(0.485, 0.456, 0.406),'std' :(0.229, 0.224, 0.225),'predResize':True,'numpy':False, 'RGB_convert_first':True,###分割模型预处理参数 'mixFunction':{'function':mixNoParking_road_postprocess, 'pars': { 'roundness': 0.3, 'cls': 9, 'laneArea': 10, 'laneAngleCha': 5 ,'RoadArea': 16000,'fitOrder':2} } }, 'trackPar':{'sort_max_age':2,'sort_min_hits':3,'sort_iou_thresh':0.2,'det_cnt':10,'windowsize':29,'patchCnt':100}, 'mode':'highWay3.0', 'Segweights' : "../weights/%s/AIlib2/%s/stdc_360X640_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###分割模型权重位置 'postFile': '../AIlib2/weights/conf/%s/para.json'%('highWay2' ),###后处理参数文件 'txtFontSize':20,###文本字符的大小 'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'segLineShow':True,'waterLineWidth':2},###显示框、线设置 'testImgPath':'images/noParking/',###测试图像的位置 'testOutPath':'images/results/',###输出测试图像位置 } par['segPar']['mixFunction']['pars']['modelSize'] = par['segPar']['modelSize'] if opt['business'] == 'cityMangement2': from DMPR import DMPRModel from DMPRUtils.jointUtil import dmpr_yolo par={ 'device':'0', ###显卡号,如果用TRT模型，只支持0（单显卡） 'labelnames':"../AIlib2/weights/conf/%s/labelnames.json"%(opt['business']), ###检测类别对照表 'max_workers':1, ###并行线程数 'half':True, 'Detweights':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###检测模型路径 #'Detweights':"/mnt/thsw2/DSP2/weights/cityMangement2/weights/urbanManagement/yolo/best.pt", 'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,5,6,7,8,9] ],###控制哪些检测类别显示、输出 'seg_nclass':4,###分割模型类别数目，默认2类 'segRegionCnt':2,###分割模型结果需要保留的等值线数目 'segPar':{ 'depth_factor':32,'NUM_FEATURE_MAP_CHANNEL':6,'dmpr_thresh':0.3, 'dmprimg_size':640, 'mixFunction':{'function':dmpr_yolo, 'pars':{'carCls':0 ,'illCls':3,'scaleRatio':0.5,'border':80} } }, 'trackPar':{'sort_max_age':2,'sort_min_hits':3,'sort_iou_thresh':0.2,'det_cnt':10,'windowsize':29,'patchCnt':100}, #'Segweights' : '/mnt/thsw2/DSP2/weights/cityMangement2/weights/urbanManagement/DMPR/dp_detector_499.engine',###分割模型权重位置 'Segweights':"../weights/%s/AIlib2/%s/dmpr_%s.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###检测模型路径 'postFile': '../AIlib2/weights/conf/%s/para.json'%(opt['business'] ),###后处理参数文件 'txtFontSize':20,###文本字符的大小 'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'segLineShow':True,'waterLineWidth':2},###显示框、线设置 #'testImgPath':'/mnt/thsw2/DSP2/demoImages/illParking',###测试图像的位置 'testImgPath':'/mnt/thsw2/DSP2/weights/cityMangement2_0916/images/input', #'testImgPath':'images/cityMangement/', 'testOutPath':'images/results/',###输出测试图像位置 } if opt['business'] == 'drowning': par={ 'device':'0', ###显卡号,如果用TRT模型，只支持0（单显卡） 'labelnames':"../AIlib2/weights/conf/%s/labelnames.json"%( opt['business'] ), ###检测类别对照表 'half':True, 'gpuname':'3090',###显卡名称 'max_workers':1, ###并行线程数 'Detweights':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###检测模型路径 #'Detweights':"../AIlib2/weights/conf/highWay2/yolov5.pt", 'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,5,6,7,8,9] ],###控制哪些检测类别显示、输出 'seg_nclass':2,###分割模型类别数目，默认2类 'segRegionCnt':2,###分割模型结果需要保留的等值线数目 'segPar':{'modelSize':(640,360),'mean':(0.485, 0.456, 0.406),'std' :(0.229, 0.224, 0.225),'predResize':True,'numpy':False, 'RGB_convert_first':True,###分割模型预处理参数 'mixFunction':{'function':mixDrowing_water_postprocess, 'pars':{ } } }, 'trackPar':{'sort_max_age':2,'sort_min_hits':3,'sort_iou_thresh':0.2,'det_cnt':10,'windowsize':29,'patchCnt':100}, 'Segweights' : "../weights/%s/AIlib2/%s/stdc_360X640_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###分割模型权重位置 'postFile': '../AIlib2/weights/conf/%s/para.json'%('highWay2' ),###后处理参数文件 'txtFontSize':20,###文本字符的大小 'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'segLineShow':True,'waterLineWidth':2},###显示框、线设置 'testImgPath':'images/drowning/',###测试图像的位置 'testOutPath':'images/results/',###输出测试图像位置 } par['segPar']['mixFunction']['pars']['modelSize'] = par['segPar']['modelSize'] ''' 以下是基于检测的跟踪模型，只有检测没有分割 ''' if opt['business'] == 'forest2': par={ 'device':'0', ###显卡号,如果用TRT模型，只支持0（单显卡） 'labelnames':"../AIlib2/weights/conf/forest2/labelnames.json", ###检测类别对照表 'gpuname':opt['gpu'],###显卡名称 'max_workers':1, ###并行线程数 'half':True, 'trtFlag_det':True,###检测模型是否采用TRT 'trtFlag_seg':False,###分割模型是否采用TRT 'Detweights':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###检测模型路径 #'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,4,5,6] ],###控制哪些检测类别显示、输出 'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [] ], 'trackPar':{'sort_max_age':2,'sort_min_hits':3,'sort_iou_thresh':0.2,'det_cnt':10,'windowsize':29,'patchCnt':100}, 'seg_nclass':2,###分割模型类别数目，默认2类 'segRegionCnt':0,###分割模型结果需要保留的等值线数目 'segPar':None,###分割模型预处理参数 'Segweights' : None,###分割模型权重位置 'postFile': '../AIlib2/weights/conf/forest/para.json',###后处理参数文件 'txtFontSize':80,###文本字符的大小 'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'waterLineWidth':3},###显示框、线设置 'testImgPath':'../AIdemo2/images/forest2/',###测试图像的位置 'testOutPath':'images/results/',###输出测试图像位置 } ###车辆巡检参数 if opt['business'] == 'vehicle': par={ 'device':'0', ###显卡号,如果用TRT模型，只支持0（单显卡） 'labelnames':"../AIlib2/weights/conf/vehicle/labelnames.json", ###检测类别对照表 'gpuname':'2080T',###显卡名称 'half':True, 'max_workers':1, ###并行线程数 'trtFlag_det':True,###检测模型是否采用TRT 'trtFlag_seg':False,###分割模型是否采用TRT 'Detweights':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###检测模型路径 'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,4,5,6] ],###控制哪些检测类别显示、输出 'trackPar':{'sort_max_age':2,'sort_min_hits':3,'sort_iou_thresh':0.2,'det_cnt':10,'windowsize':29,'patchCnt':100}, 'seg_nclass':2,###分割模型类别数目，默认2类 'segRegionCnt':0,###分割模型结果需要保留的等值线数目 'segPar':None,###分割模型预处理参数 'Segweights' : None,###分割模型权重位置 'postFile': '../AIlib2/weights/conf/vehicle/para.json',###后处理参数文件 'txtFontSize':40,###文本字符的大小 'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'waterLineWidth':3},###显示框、线设置 'testImgPath':'images/videos/vehicle/',###测试图像的位置 'testOutPath':'images/results/',###输出测试图像位置 } ###行人检测模型 if opt['business'] == 'pedestrian': par={ 'device':'0', ###显卡号,如果用TRT模型，只支持0（单显卡） 'labelnames':"../AIlib2/weights/conf/pedestrian/labelnames.json", ###检测类别对照表 'gpuname':'2080T',###显卡名称 'half':True, 'max_workers':1, ###并行线程数 'trtFlag_det':True,###检测模型是否采用TRT 'trtFlag_seg':False,###分割模型是否采用TRT 'Detweights':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###检测模型路径 'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,4,5,6] ],###控制哪些检测类别显示、输出 'trackPar':{'sort_max_age':2,'sort_min_hits':3,'sort_iou_thresh':0.2,'det_cnt':10,'windowsize':29,'patchCnt':100}, 'seg_nclass':2,###分割模型类别数目，默认2类 'segRegionCnt':0,###分割模型结果需要保留的等值线数目 'segPar':None,###分割模型预处理参数 'Segweights' : None,###分割模型权重位置 'postFile': '../AIlib2/weights/conf/pedestrian/para.json',###后处理参数文件 'txtFontSize':40,###文本字符的大小 'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'waterLineWidth':3},###显示框、线设置 'testImgPath':'../AIdemo2/images/pedestrian/',###测试图像的位置 'testOutPath':'images/results/',###输出测试图像位置 } if opt['business'] == 'smogfire': par={ 'device':'0', ###显卡号,如果用TRT模型，只支持0（单显卡） 'labelnames':"../AIlib2/weights/conf/smogfire/labelnames.json", ###检测类别对照表 'gpuname':'2080T',###显卡名称 'half':True, 'max_workers':1, ###并行线程数 'trtFlag_det':True,###检测模型是否采用TRT 'trtFlag_seg':False,###分割模型是否采用TRT 'Detweights':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###检测模型路径 'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,4,5,6] ],###控制哪些检测类别显示、输出 'trackPar':{'sort_max_age':2,'sort_min_hits':3,'sort_iou_thresh':0.2,'det_cnt':10,'windowsize':29,'patchCnt':100}, 'seg_nclass':2,###没有分割模型，此处不用 'segRegionCnt':0,###没有分割模型，此处不用 'segPar':None,###分割模型预处理参数 'Segweights' : None,###分割模型权重位置 'postFile': '../AIlib2/weights/conf/smogfire/para.json',###后处理参数文件 'txtFontSize':40,###文本字符的大小 'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'waterLineWidth':3},###显示框、线设置 'testImgPath':'../AIdemo2/images/smogfire/',###测试图像的位置 'testOutPath':'images/results/',###输出测试图像位置 } ###钓鱼游泳检测 if opt['business'] == 'AnglerSwimmer': par={ 'device':'0', ###显卡号,如果用TRT模型，只支持0（单显卡） 'labelnames':"../AIlib2/weights/conf/AnglerSwimmer/labelnames.json", ###检测类别对照表 'gpuname':'2080T',###显卡名称 'half':True, 'max_workers':1, ###并行线程数 'trtFlag_det':True,###检测模型是否采用TRT 'trtFlag_seg':False,###分割模型是否采用TRT 'Detweights':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###检测模型路径 'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,4,5,6] ],###控制哪些检测类别显示、输出 'trackPar':{'sort_max_age':2,'sort_min_hits':3,'sort_iou_thresh':0.2,'det_cnt':10,'windowsize':29,'patchCnt':100}, 'seg_nclass':2,###没有分割模型，此处不用 'segRegionCnt':0,###没有分割模型，此处不用 'segPar':None,###分割模型预处理参数 'Segweights' : None,###分割模型权重位置 'postFile': '../AIlib2/weights/conf/AnglerSwimmer/para.json',###后处理参数文件 'txtFontSize':40,###文本字符的大小 'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'waterLineWidth':3},###显示框、线设置 'testImgPath':'../AIdemo2/images/AnglerSwimmer/',###测试图像的位置 'testOutPath':'images/results/',###输出测试图像位置 } ###航道应急，做落水人员检测, channelEmergency if opt['business'] == 'channelEmergency': par={ 'device':'0', ###显卡号,如果用TRT模型，只支持0（单显卡） 'labelnames':"../AIlib2/weights/conf/channelEmergency/labelnames.json", ###检测类别对照表 'gpuname':'2080T',###显卡名称 'half':True, 'max_workers':1, ###并行线程数 'trtFlag_det':True,###检测模型是否采用TRT 'trtFlag_seg':False,###分割模型是否采用TRT 'Detweights':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###检测模型路径 #'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,4,5,6] ],###控制哪些检测类别显示、输出 'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [] ],###控制哪些检测类别显示、输出 'trackPar':{'sort_max_age':2,'sort_min_hits':3,'sort_iou_thresh':0.2,'det_cnt':10,'windowsize':29,'patchCnt':100}, 'seg_nclass':2,###没有分割模型，此处不用 'segRegionCnt':0,###没有分割模型，此处不用 'segPar':None,###分割模型预处理参数 'Segweights' : None,###分割模型权重位置 'postFile': '../AIlib2/weights/conf/channelEmergency/para.json',###后处理参数文件 'txtFontSize':40,###文本字符的大小 'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'waterLineWidth':3},###显示框、线设置 'testImgPath':'../AIdemo2/images/channelEmergency/',###测试图像的位置 'testOutPath':'images/results/',###输出测试图像位置 } ###乡村路违法种植 if opt['business'] == 'countryRoad': par={ 'device':'0', ###显卡号,如果用TRT模型，只支持0（单显卡） 'labelnames':"../AIlib2/weights/conf/countryRoad/labelnames.json", ###检测类别对照表 'gpuname':'2080T',###显卡名称 'half':True, 'max_workers':1, ###并行线程数 'trtFlag_det':True,###检测模型是否采用TRT 'trtFlag_seg':False,###分割模型是否采用TRT 'Detweights':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###检测模型路径 'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,4,5,6] ],###控制哪些检测类别显示、输出 'trackPar':{'sort_max_age':2,'sort_min_hits':3,'sort_iou_thresh':0.2,'det_cnt':10,'windowsize':29,'patchCnt':100}, 'seg_nclass':2,###没有分割模型，此处不用 'segRegionCnt':0,###没有分割模型，此处不用 'segPar':None,###分割模型预处理参数 'Segweights' : None,###分割模型权重位置 'postFile': '../AIlib2/weights/conf/countryRoad/para.json',###后处理参数文件 'txtFontSize':40,###文本字符的大小 'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'waterLineWidth':3},###显示框、线设置 'testImgPath':'../AIdemo2/images/countryRoad/',###测试图像的位置 'testOutPath':'images/results/',###输出测试图像位置 } ###城管项目，检测城市垃圾和车辆 if opt['business'] == 'cityMangement': par={ 'device':'0', ###显卡号,如果用TRT模型，只支持0（单显卡） 'labelnames':"../AIlib2/weights/conf/%s/labelnames.json"%(opt['business']), ###检测类别对照表 'gpuname':'2080Ti',###显卡名称 'half':True, 'max_workers':1, ###并行线程数 'trtFlag_det':True,###检测模型是否采用TRT 'trtFlag_seg':False,###分割模型是否采用TRT 'Detweights':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###检测模型路径 'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,4,5,6] ],###控制哪些检测类别显示、输出 'trackPar':{'sort_max_age':2,'sort_min_hits':3,'sort_iou_thresh':0.2,'det_cnt':10,'windowsize':29,'patchCnt':100}, 'seg_nclass':2,###没有分割模型，此处不用 'segRegionCnt':0,###没有分割模型，此处不用 'segPar':None,###分割模型预处理参数 'Segweights' : None,###分割模型权重位置 'postFile': '../AIlib2/weights/conf/%s/para.json'%(opt['business']),###后处理参数文件 'txtFontSize':40,###文本字符的大小 'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'waterLineWidth':3},###显示框、线设置 'testImgPath':'images/cityMangement',###测试图像的位置 'testOutPath':'images/results/',###输出测试图像位置 } ###城管项目，检测道路情况，输入类别为五个："护栏","交通标志","非交通标志","施工","施工“(第4，第5类别合并，名称相同) ###实际模型检测输出的类别为："护栏","交通标志","非交通标志","锥桶","水马" if opt['business'] == 'cityRoad': par={ 'device':'0', ###显卡号,如果用TRT模型，只支持0（单显卡） 'labelnames':"../AIlib2/weights/conf/%s/labelnames.json"%(opt['business']), ###检测类别对照表 'gpuname':'2080Ti',###显卡名称 'half':True, 'max_workers':1, ###并行线程数 'trtFlag_det':True,###检测模型是否采用TRT 'trtFlag_seg':False,###分割模型是否采用TRT 'Detweights':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###检测模型路径 'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,4,5,6] ],###控制哪些检测类别显示、输出 'trackPar':{'sort_max_age':2,'sort_min_hits':3,'sort_iou_thresh':0.2,'det_cnt':10,'windowsize':29,'patchCnt':100}, 'seg_nclass':2,###没有分割模型，此处不用 'segRegionCnt':0,###没有分割模型，此处不用 'segPar':None,###分割模型预处理参数 'Segweights' : None,###分割模型权重位置 'postFile': '../AIlib2/weights/conf/%s/para.json'%(opt['business']),###后处理参数文件 'txtFontSize':40,###文本字符的大小 'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'waterLineWidth':3},###显示框、线设置 'testImgPath':'images/%s'%(opt['business']),###测试图像的位置 'testOutPath':'images/results/',###输出测试图像位置 } if opt['business'] == 'illParking': from utilsK.illParkingUtils import illParking_postprocess par={ 'device':'0', ###显卡号,如果用TRT模型，只支持0（单显卡） 'half':True, 'labelnames':"../AIlib2/weights/conf/%s/labelnames.json"%(opt['business']), ###检测类别对照表 'max_workers':1, ###并行线程数 'Detweights':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###检测模型路径 'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,5,6,7,8,9] ],###控制哪些检测类别显示、输出 'trackPar':{'sort_max_age':2,'sort_min_hits':3,'sort_iou_thresh':0.2,'det_cnt':10,'windowsize':29,'patchCnt':100}, 'seg_nclass':4,###没有分割模型，此处不用 'segRegionCnt':2,###没有分割模型，此处不用 'segPar':{ 'mixFunction':{'function':illParking_postprocess, 'pars':{ } } }, 'Segweights' : None,###分割模型权重位置 'postFile': '../AIlib2/weights/conf/%s/para.json'%(opt['business'] ),###后处理参数文件 'txtFontSize':20,###文本字符的大小 'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'waterLineWidth':2},###显示框、线设置 'testImgPath':'images/cityMangement',###测试图像的位置 'testOutPath':'images/results/',###输出测试图像位置 } par['trtFlag_det']=True if par['Detweights'].endswith('.engine') else False if par['Segweights']: par['segPar']['trtFlag_seg']=True if par['Segweights'].endswith('.engine') else False ##使用森林，道路模型,business 控制['forest','road'] ##预先设置的参数 #gpuname=par['gpuname']#如果用trt就需要此参数，只能是"3090" "2080Ti" device_=par['device'] ##选定模型，可选 cpu,'0','1' device = select_device(device_) half = device.type != 'cpu' # half precision only supported on CUDA trtFlag_det=par['trtFlag_det'] ###是否采用TRT模型加速 ##以下参数目前不可改 imageW=1080 ####道路模型 digitFont= par['digitFont'] ####加载检测模型 if trtFlag_det: Detweights=par['Detweights'] logger = trt.Logger(trt.Logger.ERROR) with open(Detweights, "rb") as f, trt.Runtime(logger) as runtime: model=runtime.deserialize_cuda_engine(f.read())# 输入trt本地文件，返回ICudaEngine对象 print('####load TRT model :%s'%(Detweights)) else: Detweights=par['Detweights'] model = attempt_load(Detweights, map_location=device) # load FP32 model if half: model.half() ####加载分割模型 seg_nclass = par['seg_nclass'] segPar=par['segPar'] if par['Segweights']: if opt['business'] == 'cityMangement2': segmodel = DMPRModel(weights=par['Segweights'], par = par['segPar']) else: segmodel = stdcModel(weights=par['Segweights'], par = par['segPar']) ''' if par['segPar']['trtFlag_seg']: Segweights = par['Segweights'] logger = trt.Logger(trt.Logger.ERROR) with open(Segweights, "rb") as f, trt.Runtime(logger) as runtime: segmodel=runtime.deserialize_cuda_engine(f.read())# 输入trt本地文件，返回ICudaEngine对象 print('############locad seg model trt success: ',Segweights) else: Segweights = par['Segweights'] segmodel = SegModel(nclass=seg_nclass,weights=Segweights,device=device) print('############locad seg model pth success:',Segweights) ''' else: segmodel=None trackPar=par['trackPar'] sort_tracker = Sort(max_age=trackPar['sort_max_age'], min_hits=trackPar['sort_min_hits'], iou_threshold=trackPar['sort_iou_thresh']) labelnames = par['labelnames'] postFile= par['postFile'] print( Detweights,labelnames ) conf_thres,iou_thres,classes,rainbows=get_postProcess_para(postFile) detPostPar = get_postProcess_para_dic(postFile) conf_thres,iou_thres,classes,rainbows = detPostPar["conf_thres"],detPostPar["iou_thres"],detPostPar["classes"],detPostPar["rainbows"] if 'ovlap_thres_crossCategory' in detPostPar.keys(): iou2nd=detPostPar['ovlap_thres_crossCategory'] else:iou2nd = None if 'score_byClass' in detPostPar.keys(): score_byClass=detPostPar['score_byClass'] else: score_byClass = None ####模型选择参数用如下： mode_paras=par['detModelpara'] allowedList,allowedList_string=get_needed_objectsIndex(mode_paras) #slopeIndex = par['slopeIndex'] ##只加载检测模型，准备好显示字符 names=get_labelnames(labelnames) #imageW=4915;###默认是1920，在森林巡检的高清图像中是4920 outfontsize=int(imageW/1920*40);### label_arraylist = get_label_arrays(names,rainbows,outfontsize=par['txtFontSize'],fontpath="../AIlib2/conf/platech.ttf") ##图像测试和视频 outpth = par['testOutPath'] impth = par['testImgPath'] imgpaths=[]###获取文件里所有的图像 videopaths=[]###获取文件里所有的视频 img_postfixs = ['.jpg','.JPG','.PNG','.png']; vides_postfixs= ['.MP4','.mp4','.avi'] if os.path.isdir(impth): for postfix in img_postfixs: imgpaths.extend(glob.glob('%s/*%s'%(impth,postfix )) ) for postfix in ['.MP4','.mp4','.avi']: videopaths.extend(glob.glob('%s/*%s'%(impth,postfix )) ) else: postfix = os.path.splitext(impth)[-1] if postfix in img_postfixs: imgpaths=[ impth ] if postfix in vides_postfixs: videopaths = [impth ] imgpaths.sort() modelPar={ 'det_Model': model,'seg_Model':segmodel } processPar={'half':par['half'],'device':device,'conf_thres':conf_thres,'iou_thres':iou_thres,'trtFlag_det':trtFlag_det,'iou2nd':iou2nd,'score_byClass':score_byClass} drawPar={'names':names,'label_arraylist':label_arraylist,'rainbows':rainbows,'font': par['digitFont'],'allowedList':allowedList} for i in range(len(imgpaths)): #for i in range(2): #imgpath = os.path.join(impth, folders[i]) imgpath = imgpaths[i] bname = os.path.basename(imgpath ) im0s=[cv2.imread(imgpath)] time00 = time.time() retResults,timeOut = AI_det_track_batch(im0s, [i] ,modelPar,processPar,sort_tracker ,trackPar,segPar) #print('###line627:',retResults[2]) #retResults,timeInfos = AI_det_track_batch(imgarray_list, iframe_list ,par0['modelPar'],par0['processPar'],par0['sort_tracker'] ,par0['trackPar'],segPar=par0['segPar']) if len(retResults[1])>0: retResults[0][0] = drawBoxTraceSimplied(retResults[1],i, retResults[0][0],rainbows=rainbows,boxFlag=True,traceFlag=False,names=drawPar['names']) time11 = time.time() image_array = retResults[0][0] ''' 返回值retResults[2] --list,其中每一个元素为一个list，表示每一帧的检测结果，每一个结果是由多个list构成，每个list表示一个框，格式为[ cls , x0 ,y0 ,x1 ,y1 ,conf,ifrmae,trackId ] --etc. retResults[2][j][k]表示第j帧的第k个框。 ''' cv2.imwrite( os.path.join( outpth,bname ) ,image_array ) print('----image:%s, process:%s ( %s ),save:%s'%(bname,(time11-time00) * 1000, timeOut,(time.time() - time11) * 1000) ) ##process video print('##begin to process videos, total %d videos'%( len(videopaths))) for i,video in enumerate(videopaths): print('process video%d :%s '%(i,video)) par0={'modelPar':modelPar,'processPar':processPar,'drawPar':drawPar,'outpth':par['testOutPath'], 'sort_tracker':sort_tracker,'trackPar':trackPar,'segPar':segPar} process_video(video,par0,mode='track') def OCR_demo2(opt): from ocrUtils2 import crnn_model from ocrUtils2.ocrUtils import get_cfg,recognition_ocr,strLabelConverter if opt['business'] == 'ocr2': par={ 'image_dir':'images/ocr_en', 'outtxt':'images/results', 'weights':'../AIlib2/weights/conf/ocr2/crnn_448X32.pth', #'weights':'../weights/2080Ti/AIlib2/ocr2/crnn_2080Ti_fp16_448X32.engine', 'device':'cuda:0', 'cfg':'../AIlib2/weights/conf/ocr2/360CC_config.yaml', 'char_file':'../AIlib2/weights/conf/ocr2/chars.txt', 'imgH':32, 'imgW':448, 'workers':1 } image_dir=par['image_dir'] outtxt=par['outtxt'] workers=par['workers'] weights= par['weights'] device=par['device'] char_file=par['char_file'] imgH=par['imgH'] imgW=par['imgW'] cfg = par['cfg'] config = get_cfg(cfg, char_file) par['contextFlag']=False device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu') if weights.endswith('.pth'): model = crnn_model.get_crnn(config,weights=weights).to(device) par['model_mode']='pth' else: logger = trt.Logger(trt.Logger.ERROR) with open(weights, "rb") as f, trt.Runtime(logger) as runtime: model = runtime.deserialize_cuda_engine(f.read())# 输入trt本地文件，返回ICudaEngine对象 print('#####load TRT file:',weights,'success #####') context = model.create_execution_context() par['model_mode']='trt';par['contextFlag']=context converter = strLabelConverter(config.DATASET.ALPHABETS) img_urls=glob.glob('%s/*.jpg'%( image_dir )) img_urls.extend( glob.glob('%s/*.png'%( image_dir )) ) cnt=len(img_urls) print('%s has %d images'%(image_dir ,len(img_urls) ) ) # 准备数据 parList=[] for i in range(cnt): img_patch=cv2.imread( img_urls[i] , cv2.IMREAD_GRAYSCALE) started = time.time() img = cv2.imread(img_urls[i]) sim_pred = recognition_ocr(config, img, model, converter, device,par=par) finished = time.time() print('{0}: elapsed time: {1} prd:{2} '.format( os.path.basename( img_urls[i] ), finished - started, sim_pred )) def OBB_track_demo(opt): ###倾斜框（OBB）的ship目标检测 par={ 'obbModelPar':{ 'model_size':(608,608),'K':100,'conf_thresh':0.3, 'down_ratio':4,'num_classes':15,'dataset':'dota', 'heads': {'hm': None,'wh': 10,'reg': 2,'cls_theta': 1}, 'mean':(0.5, 0.5, 0.5),'std':(1, 1, 1), 'half': False,'decoder':None, 'weights':'../weights/%s/AIlib2/%s/obb_608X608_%s_fp16.engine'%(opt['gpu'],opt['business'],opt['gpu']), }, 'outpth': 'images/results', 'trackPar':{'sort_max_age':2,'sort_min_hits':3,'sort_iou_thresh':0.2,'det_cnt':10,'windowsize':29,'patchCnt':100}, 'device':"cuda:0", #'test_dir': '/mnt/thsw2/DSP2/videos/obbShips/DJI_20230208110806_0001_W_6M.MP4', 'test_dir':'/mnt/thsw2/DSP2/videos/obbShips/freighter2.mp4', 'test_flag':True, 'postFile': '../AIlib2/weights/conf/%s/para.json'%(opt['business'] ),###后处理参数文件 'drawBox':True,#####是否画框 'drawPar': { 'digitWordFont' :{'line_thickness':2,'boxLine_thickness':1,'wordSize':40, 'fontSize':1.0,'label_location':'leftTop'}} , 'labelnames':"../AIlib2/weights/conf/%s/labelnames.json"%(opt['business'] ), ###检测类别对照表 } #par['model_size'],par['mean'],par['std'],par['half'],par['saveType'],par['heads'],par['labelnames'],par['decoder'],par['down_ratio'],par['drawBox'] #par['rainbows'],par['label_array'],par['digitWordFont'] obbModelPar = par['obbModelPar'] ####加载模型 model,decoder2=load_model_decoder_OBB(obbModelPar) obbModelPar['decoder']=decoder2 names=get_labelnames(par['labelnames']);obbModelPar['labelnames']=names _,_,_,rainbows=get_postProcess_para(par['postFile']);par['drawPar']['rainbows']=rainbows label_arraylist = get_label_arrays(names,rainbows,outfontsize=par['drawPar']['digitWordFont']['wordSize'],fontpath="../AIlib2/conf/platech.ttf") #par['label_array']=label_arraylist trackPar=par['trackPar'] sort_tracker = OBB_Sort(max_age=trackPar['sort_max_age'], min_hits=trackPar['sort_min_hits'], iou_threshold=trackPar['sort_iou_thresh']) ##图像测试和视频 impth = par['test_dir'] img_urls=[]###获取文件里所有的图像 video_urls=[]###获取文件里所有的视频 img_postfixs = ['.jpg','.JPG','.PNG','.png']; vides_postfixs= ['.MP4','.mp4','.avi'] if os.path.isdir(impth): for postfix in img_postfixs: img_urls.extend(glob.glob('%s/*%s'%(impth,postfix )) ) for postfix in ['.MP4','.mp4','.avi']: video_urls.extend(glob.glob('%s/*%s'%(impth,postfix )) ) else: postfix = os.path.splitext(impth)[-1] if postfix in img_postfixs: img_urls=[ impth ] if postfix in vides_postfixs: video_urls = [impth ] parIn = {'obbModelPar':obbModelPar,'modelPar':{'obbmodel': model},'sort_tracker':sort_tracker,'outpth':par['outpth'],'trackPar':trackPar,'drawPar':par['drawPar']} par['drawPar']['label_array']=label_arraylist for img_url in img_urls: #print(img_url) ori_image=cv2.imread(img_url) #ori_image_list,infos = OBB_infer(model,ori_image,obbModelPar) ori_image_list,infos = OBB_tracker_batch([ori_image],[0],parIn['modelPar'],parIn['obbModelPar'],None,parIn['trackPar'],None) ori_image_list[1] = draw_obb(ori_image_list[2] ,ori_image_list[1],par['drawPar']) imgName = os.path.basename(img_url) saveFile = os.path.join(par['outpth'], imgName) ret=cv2.imwrite(saveFile, ori_image_list[1]) if not ret: print(saveFile, ' not created ') print( os.path.basename(img_url),':',infos,ori_image_list[2]) ###处理视频 for video_url in video_urls: process_video(video_url, parIn ,mode='obbTrack') def crowd_demo(opt): if opt['business']=='crowdCounting': from crowd import crowdModel as Model par={ 'mean':[0.485, 0.456, 0.406], 'std':[0.229, 0.224, 0.225],'threshold':0.5, '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/%s/AIlib2/%s/crowdCounting_%s_dynamic.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###检测模型路径 'testImgPath':'images/%s'%(opt['business'] ),###测试图像的位置 'testOutPath':'images/results/',###输出测试图像位置 } #weights='weights/best_mae.pth' cmodel = Model(par['weights'],par) img_path = par['testImgPath'] File = os.listdir(img_path) targetList = [] for file in File[0:]: COORlist = [] imgPath = img_path + os.sep + file img_raw = cv2.cvtColor(cv2.imread(imgPath),cv2.COLOR_BGR2RGB) # cmodel.eval--- # 输入读取的RGB数组 # 输出:list,0--原图，1-人头坐标list,2-对接OBB的格式数据，其中4个坐标均相同，2-格式如下： # [ [ [ (x0,y0),(x1,y1),(x2,y2),(x3,y3) ],score， cls ], [ [ (x0,y0),(x1,y1),(x2,y2),(x3,y3) ],score ，cls ],........ ] prets, infos = cmodel.eval(img_raw) print(file,infos,' 人数：',len(prets[1])) img_to_draw = cv2.cvtColor(np.array(img_raw), cv2.COLOR_RGB2BGR) # 打印预测图像中人头的个数 for p in prets[1]: 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['testOutPath'], file), img_to_draw) if __name__=="__main__": #jkm_demo() businessAll=['river', 'river2','highWay2','noParking','drowning','forest2','vehicle','pedestrian','smogfire' , 'AnglerSwimmer','channelEmergency', 'countryRoad','cityMangement','ship2','cityMangement2','cityRoad','illParking',"crowdCounting"] businessAll = ['crowdCounting'] for busi in businessAll: print('-'*40,'beg to test:',busi,'-'*40) opt={'gpu':'2080Ti','business':busi} if busi in ['ship2']: OBB_track_demo(opt) elif opt['business'] in ['crowdCounting'] : crowd_demo(opt) else: #if opt['business'] in ['river','highWay2','noParking','drowning','']: det_track_demo(opt )