wangjin0928
/
AIdemo2
Sledovat 1
Oblíbit 0
Rozštěpit 0
Zdrojový kód Úkoly 0 Požadavky na natažení 0 Vydání 0 Wiki Aktivita
Nelze vybrat více než 25 témat Téma musí začínat písmenem nebo číslem, může obsahovat pomlčky („-“) a může být dlouhé až 35 znaků.
47 Revize
1 Větev
Strom: f152998111
Větve Značky
${ item.name }
Vytvořit větev ${ searchTerm }
z „f152998111“
${ noResults }
AIdemo2/demo3.0.py

962 lines
54KB
Surový Blame Historie 

  1. import sys, yaml

  2. from easydict import EasyDict as edict

  3. from concurrent.futures import ThreadPoolExecutor

  4. 

  5. sys.path.extend(['..','../AIlib2' ])

  6. 

  7. from AI import AI_process,AI_process_forest,get_postProcess_para,get_postProcess_para_dic,ocr_process,AI_det_track,AI_det_track_batch

  8. import cv2,os,time

  9. from segutils.segmodel import SegModel

  10. 

  11. from stdc import stdcModel

  12. from segutils.trafficUtils import tracfficAccidentMixFunction

  13. from models.experimental import attempt_load

  14. from utils.torch_utils import select_device

  15. from utilsK.queRiver import get_labelnames,get_label_arrays,save_problem_images,riverDetSegMixProcess

  16. from ocrUtils.ocrUtils import CTCLabelConverter,AlignCollate

  17. from trackUtils.sort import Sort,track_draw_boxAndTrace,track_draw_trace_boxes,moving_average_wang,drawBoxTraceSimplied

  18. from trackUtils.sort_obb import OBB_Sort,obbTohbb,track_draw_all_boxes,track_draw_trace

  19. from obbUtils.shipUtils import OBB_infer,OBB_tracker,draw_obb,OBB_tracker_batch

  20. from utilsK.noParkingUtils import mixNoParking_road_postprocess

  21. from obbUtils.load_obb_model import load_model_decoder_OBB

  22. import numpy as np

  23. import torch,glob

  24. import tensorrt as trt

  25. from utilsK.masterUtils import get_needed_objectsIndex

  26. from copy import deepcopy

  27. from scipy import interpolate

  28. from utilsK.drownUtils import mixDrowing_water_postprocess

  29. #import warnings

  30. #warnings.filterwarnings("error")

  31. 

  32. def view_bar(num, total,time1,prefix='prefix'):

  33.     rate = num / total

  34.     time_n=time.time()

  35.     rate_num = int(rate * 30)

  36.     rate_nums = np.round(rate * 100)

  37.     r = '\r %s %d / %d [%s%s] %.2f s'%(prefix,num,total,  ">" * rate_num, " " * (30 - rate_num), time_n-time1 )

  38.     sys.stdout.write(r)

  39.     sys.stdout.flush()

  40.         

  41. 

  42. '''

  43.     多线程

  44. '''

  45. 

  46. def process_v1(frame):

  47.     #try:

  48.     print('demo.py beging to :',frame[8])

  49.     time00 = time.time()

  50.     H,W,C = frame[0][0].shape

  51.     

  52.     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])

  53.      

  54.     time11 = time.time()

  55.     image_array = p_result[1]

  56.     

  57.     cv2.imwrite(os.path.join('images/results/',frame[8] ) ,image_array)

  58.     bname = frame[8].split('.')[0]

  59.     if len(p_result)==5:

  60.         image_mask = p_result[4]

  61.         cv2.imwrite(os.path.join('images/results/',bname+'_mask.png' ) , (image_mask).astype(np.uint8))

  62.     

  63.     boxes=p_result[2]

  64.     with open(  os.path.join('images/results/',bname+'.txt' ),'w'  ) as fp:

  65.         for box in boxes:

  66.             box_str=[str(x) for x in box]

  67.             out_str=','.join(box_str)+'\n'

  68.             fp.write(out_str)

  69.     time22 = time.time()

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

  71.     return 'success'

  72.     #except Exception as e:

  73.     #    return 'failed:'+str(e)        

  74. def process_video(video,par0,mode='detSeg'):

  75.     cap=cv2.VideoCapture(video)

  76.     if not  cap.isOpened():

  77.         print('#####error url:',video)

  78.         return False  

  79.     bname=os.path.basename(video).split('.')[0]        

  80.     fps = int(cap.get(cv2.CAP_PROP_FPS)+0.5)

  81.     width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH )+0.5)

  82.     height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)+0.5)

  83.     framecnt=int(cap.get(7)+0.5)

  84.     save_path_AI = os.path.join(par0['outpth'],os.path.basename(video))

  85.     problem_image_dir= os.path.join(  par0['outpth'], 'probleImages'   )

  86.     os.makedirs(problem_image_dir,exist_ok=True)

  87.     vid_writer_AI = cv2.VideoWriter(save_path_AI, cv2.VideoWriter_fourcc(*'mp4v'), fps, (width,height))

  88.     num=0

  89.     iframe=0;post_results=[];fpsample=30*10

  90.     

  91.     imgarray_list = []; iframe_list = []

  92.     

  93.     patch_cnt = par0['trackPar']['patchCnt']

  94.     ##windowsize 对逐帧插值后的结果做平滑，windowsize为平滑的长度,没隔det_cnt帧做一次跟踪。

  95.     trackPar={'det_cnt':10,'windowsize':29 }

  96. 

  97. 

  98.     ##track_det_result_update= np.empty((0,8)) ###每100帧跑出来的结果，放在track_det_result_update，只保留当前100帧里有的tracker Id.

  99.     while cap.isOpened():

  100.         ret, imgarray = cap.read() #读取摄像头画面

  101.         iframe +=1

  102.         if not ret:break

  103.         if mode=='detSeg':

  104.             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'])

  105.         elif mode == 'track':

  106.             #sampleCount=10

  107.             imgarray_list.append( imgarray )

  108.             iframe_list.append(iframe )    

  109.             if iframe%patch_cnt==0:

  110.                 time_patch0 = time.time()

  111.                 retResults,timeInfos = AI_det_track_batch(imgarray_list, iframe_list ,par0['modelPar'],par0['processPar'],par0['sort_tracker'] ,par0['trackPar'],segPar=par0['segPar'])

  112.                 #print('###line111:',retResults[2])

  113.                 ###需要保存成一个二维list，每一个list是一帧检测结果。

  114.                 ###track_det_result 内容格式：x1, y1, x2, y2, conf, cls,iframe,trackId

  115.                 time_patch2 = time.time()                    

  116.                 frame_min = iframe_list[0];frame_max=iframe_list[-1]

  117.                 for iiframe in range(frame_min,frame_max+1):

  118.                     img_draw = imgarray_list[ iiframe- frame_min  ]

  119.                     img_draw = drawBoxTraceSimplied(retResults[1] ,iiframe, img_draw,rainbows=par0['drawPar']['rainbows'],boxFlag=True,traceFlag=True,names=par0['drawPar']['names'] )

  120.                     ret = vid_writer_AI.write(img_draw)        

  121.                     view_bar(iiframe, framecnt,time.time(),prefix=os.path.basename(video))

  122.                 imgarray_list=[];iframe_list=[]     

  123.         elif mode =='obbTrack':

  124.             imgarray_list.append( imgarray )

  125.             iframe_list.append(iframe )    

  126.             if iframe%patch_cnt==0:

  127.                 time_patch0 = time.time()

  128.                     

  129.                 track_det_results,  timeInfos = OBB_tracker_batch(imgarray_list,iframe_list,par0['modelPar'],par0['obbModelPar'],par0['sort_tracker'],par0['trackPar'],segPar=None)

  130.                 print(  timeInfos )

  131.                 

  132.                 #对结果画图

  133.                 track_det_np = track_det_results[1]

  134.                 frame_min = iframe_list[0];frame_max=iframe_list[-1]

  135.                 for iiframe in range(frame_min,frame_max+1):

  136.                     img_draw = imgarray_list[ iiframe- frame_min  ]

  137.                     

  138.                     if len( track_det_results[2][  iiframe- frame_min]) > 0:

  139.                         img_draw =  draw_obb( track_det_results[2][iiframe- frame_min ] ,img_draw,par0['drawPar'])

  140.                         if  True: 

  141.                             frameIdex=12;trackIdex=13;

  142.                             boxes_oneFrame = track_det_np[ track_det_np[:,frameIdex]==iiframe   ]

  143.                             

  144.                             ###在某一帧上，画上轨迹

  145.                             track_ids = boxes_oneFrame[:,trackIdex].tolist()

  146.                             boxes_before_oneFrame = track_det_np[ track_det_np[:,frameIdex]<=iiframe   ]

  147.                             for trackId in track_ids:

  148.                                 boxes_before_oneFrame_oneId = boxes_before_oneFrame[boxes_before_oneFrame[:,trackIdex]==trackId]

  149.                                 xcs = boxes_before_oneFrame_oneId[:,8]

  150.                                 ycs = boxes_before_oneFrame_oneId[:,9]

  151.                                 [cv2.line(img_draw, (  int(xcs[i]) , int(ycs[i]) ), 

  152.                                     ( int(xcs[i+1]),int(ycs[i+1])  ),(255,0,0), thickness=2) 

  153.                                     for i,_ in  enumerate(xcs)  if i < len(xcs)-1 ]

  154. 

  155.                     ret = vid_writer_AI.write(img_draw)   

  156.                    

  157.                 #sys.exit(0)

  158.                 #print('vide  writer ret:',ret)

  159.                 imgarray_list=[];iframe_list=[]     

  160.   

  161.             view_bar(iframe, framecnt,time.time(),prefix=os.path.basename(video))    

  162.             

  163.         else:

  164.             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'])

  165.             

  166. 

  167.         if  mode not in [ 'track','obbTrack']:   

  168.             image_array = p_result[1];num+=1

  169.             ret = vid_writer_AI.write(image_array)    

  170.             view_bar(num, framecnt,time.time(),prefix=os.path.basename(video))

  171.             ##每隔 fpsample帧处理一次，如果有问题就保存图片  

  172.             if (iframe % fpsample == 0) and (len(post_results)>0) :

  173.                 parImage=save_problem_images(post_results,iframe,par0['names'],streamName=bname,outImaDir=problem_image_dir,imageTxtFile=False)

  174.                 post_results=[]   

  175.     

  176.             if len(p_result[2] )>0:                

  177.                 post_results.append(p_result)

  178.  

  179.     vid_writer_AI.release();       

  180.       

  181. def det_track_demo(business ):

  182.     '''

  183.     跟踪参数说明：

  184.     'trackPar':{'sort_max_age':2,'sort_min_hits':3,'sort_iou_thresh':0.2,'det_cnt':10,'windowsize':29,'patchCnt':100} 

  185.         sort_max_age--跟踪链断裂时允许目标消失最大的次数。超过之后，会认为是新的目标。

  186.         sort_min_hits--每隔目标连续出现的次数，超过这个次数才认为是一个目标。

  187.         sort_iou_thresh--检测最小的置信度。

  188.         det_cnt--每隔几次做一个跟踪和检测，默认10。

  189.         windowsize--轨迹平滑长度，一定是奇数，表示每隔几帧做一平滑，默认29。

  190.         patchCnt--每次送入图像的数量，不宜少于100帧。

  191.     '''

  192.     

  193.     ''' 以下是基于检测和分割的跟踪模型，分割用来修正检测的结果'''        

  194.     ####河道巡检的跟踪模型参数

  195.     if opt['business'] == 'river' or opt['business'] == 'river2' : 

  196.         par={

  197.             'device':'0', ###显卡号,如果用TRT模型，只支持0（单显卡）

  198.             'labelnames':"../AIlib2/weights/conf/%s/labelnames.json"%(opt['business']), ###检测类别对照表

  199.             'gpuname':'2080Ti',###显卡名称

  200.             'max_workers':1, ###并行线程数

  201.             'half':True,

  202.             'Detweights':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu']  ),###检测模型路径

  203.             'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,4,5,6]  ],###控制哪些检测类别显示、输出

  204.             

  205.             'seg_nclass':2,###分割模型类别数目，默认2类

  206.             'segRegionCnt':0,###分割模型结果需要保留的等值线数目

  207.             

  208.             'segPar':{'modelSize':(640,360),'mean':(0.485, 0.456, 0.406),'std' :(0.229, 0.224, 0.225),'numpy':False, 'RGB_convert_first':True,#分割模型预处理参数

  209.                         'mixFunction':{'function':riverDetSegMixProcess,'pars':{'slopeIndex':[1,3,4,7], 'riverIou':0.1}} #分割和检测混合处理的函数 

  210.                         },

  211.             

  212.             'trackPar':{'sort_max_age':2,'sort_min_hits':3,'sort_iou_thresh':0.2,'det_cnt':10,'windowsize':29,'patchCnt':100},

  213.             

  214.             'Segweights' : "../weights/%s/AIlib2/%s/stdc_360X640_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu']  ),###分割模型权重位置

  215.             'postFile': '../AIlib2/weights/conf/%s/para.json'%( opt['business'] ),###后处理参数文件

  216.             'txtFontSize':80,###文本字符的大小

  217.             'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'waterLineWidth':3},###显示框、线设置

  218.             #'testImgPath':'images/videos/river',###测试图像的位置

  219.             'testImgPath':'images/tt',###测试图像的位置

  220.             'testOutPath':'images/results/',###输出测试图像位置

  221.             

  222.         }

  223.     

  224.     if opt['business'] == 'highWay2':

  225.         par={

  226.             'device':'0', ###显卡号,如果用TRT模型，只支持0（单显卡）

  227.             'labelnames':"../AIlib2/weights/conf/%s/labelnames.json"%(  opt['business']  ), ###检测类别对照表

  228.             'half':True,

  229.             'gpuname':'3090',###显卡名称

  230.             'max_workers':1, ###并行线程数

  231.             'Detweights':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu']  ),###检测模型路径

  232.             #'Detweights':"../AIlib2/weights/conf/highWay2/yolov5.pt",

  233.             'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,5,6,7,8,9]  ],###控制哪些检测类别显示、输出

  234.             'seg_nclass':3,###分割模型类别数目，默认2类

  235.             'segRegionCnt':2,###分割模型结果需要保留的等值线数目

  236.             '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,###分割模型预处理参数

  237.                       'mixFunction':{'function':tracfficAccidentMixFunction,

  238.                         '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 }

  239.                         }

  240.                     },

  241. 

  242.             'trackPar':{'sort_max_age':2,'sort_min_hits':3,'sort_iou_thresh':0.2,'det_cnt':5,'windowsize':29,'patchCnt':100},

  243.             'mode':'highWay3.0',

  244.             'Segweights' : "../weights/%s/AIlib2/%s/stdc_360X640_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu']  ),###分割模型权重位置

  245.             'postFile': '../AIlib2/weights/conf/%s/para.json'%(opt['business'] ),###后处理参数文件

  246.             'txtFontSize':20,###文本字符的大小

  247.             'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':0.5,'waterLineColor':(0,255,255),'segLineShow':True,'waterLineWidth':2},###显示框、线设置

  248.             #'testImgPath':'images/trafficAccident/8.png',###测试图像的位置

  249.             'testImgPath':'/home/chenyukun/777-7-42.mp4',###测试图像的位置

  250.             'testOutPath':'images/results/',###输出测试图像位置

  251.         }

  252.         par['segPar']['mixFunction']['pars']['modelSize'] = par['segPar']['modelSize']

  253.     if opt['business'] == 'noParking':

  254.         par={

  255.             'device':'0', ###显卡号,如果用TRT模型，只支持0（单显卡）

  256.             'labelnames':"../AIlib2/weights/conf/%s/labelnames.json"%(  opt['business']  ), ###检测类别对照表

  257.             'half':True,

  258.             'gpuname':'3090',###显卡名称

  259.             'max_workers':1, ###并行线程数

  260.             'Detweights':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu']  ),###检测模型路径

  261.             #'Detweights':"../AIlib2/weights/conf/highWay2/yolov5.pt",

  262.             'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,5,6,7,8,9]  ],###控制哪些检测类别显示、输出

  263.             'seg_nclass':4,###分割模型类别数目，默认2类

  264.             'segRegionCnt':2,###分割模型结果需要保留的等值线数目

  265.             '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,###分割模型预处理参数

  266.                       'mixFunction':{'function':mixNoParking_road_postprocess,

  267.                         'pars': { 'roundness': 0.3, 'cls': 9,  'laneArea': 10, 'laneAngleCha': 5 ,'RoadArea': 16000,'fitOrder':2}

  268.                         }    

  269.                     },

  270.                     

  271.                     

  272. 

  273.             'trackPar':{'sort_max_age':2,'sort_min_hits':3,'sort_iou_thresh':0.2,'det_cnt':10,'windowsize':29,'patchCnt':100},

  274.             'mode':'highWay3.0',

  275.             'Segweights' : "../weights/%s/AIlib2/%s/stdc_360X640_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu']  ),###分割模型权重位置

  276.             'postFile': '../AIlib2/weights/conf/%s/para.json'%('highWay2' ),###后处理参数文件

  277.             'txtFontSize':20,###文本字符的大小

  278.             'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'segLineShow':True,'waterLineWidth':2},###显示框、线设置

  279.             'testImgPath':'images/noParking/',###测试图像的位置

  280.             'testOutPath':'images/results/',###输出测试图像位置

  281.         }

  282.         par['segPar']['mixFunction']['pars']['modelSize'] = par['segPar']['modelSize']

  283.     

  284.     if opt['business'] == 'cityMangement2':

  285.         from DMPR import DMPRModel

  286.         from DMPRUtils.jointUtil import  dmpr_yolo

  287. 

  288.         par={

  289.             'device':'0', ###显卡号,如果用TRT模型，只支持0（单显卡）

  290.             'labelnames':"../AIlib2/weights/conf/%s/labelnames.json"%(opt['business']), ###检测类别对照表

  291.             'max_workers':1, ###并行线程数

  292.             'half':True,

  293.             'Detweights':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu']  ),###检测模型路径

  294.             #'Detweights':"/mnt/thsw2/DSP2/weights/cityMangement2/weights/urbanManagement/yolo/best.pt",

  295.             'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,5,6,7,8,9]  ],###控制哪些检测类别显示、输出

  296.             'seg_nclass':4,###分割模型类别数目，默认2类

  297.             'segRegionCnt':2,###分割模型结果需要保留的等值线数目

  298. 

  299.             'segPar':{ 'depth_factor':32,'NUM_FEATURE_MAP_CHANNEL':6,'dmpr_thresh':0.3, 'dmprimg_size':640,

  300.                       'mixFunction':{'function':dmpr_yolo,

  301.                         'pars':{'carCls':0 ,'illCls':3,'scaleRatio':0.5,'border':80}

  302.                         }

  303.                     },

  304.             'trackPar':{'sort_max_age':2,'sort_min_hits':3,'sort_iou_thresh':0.2,'det_cnt':10,'windowsize':29,'patchCnt':100},       

  305.             #'Segweights' : '/mnt/thsw2/DSP2/weights/cityMangement2/weights/urbanManagement/DMPR/dp_detector_499.engine',###分割模型权重位置

  306.             'Segweights':"../weights/%s/AIlib2/%s/dmpr_%s.engine"%(opt['gpu'], opt['business'] ,opt['gpu']  ),###检测模型路径

  307.             'postFile': '../AIlib2/weights/conf/%s/para.json'%(opt['business'] ),###后处理参数文件

  308.             'txtFontSize':20,###文本字符的大小

  309.             'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'segLineShow':True,'waterLineWidth':2},###显示框、线设置

  310.             #'testImgPath':'/mnt/thsw2/DSP2/demoImages/illParking',###测试图像的位置

  311.             'testImgPath':'/mnt/thsw2/DSP2/weights/cityMangement2_0916/images/input',

  312.             #'testImgPath':'images/cityMangement/',

  313.             'testOutPath':'images/results/',###输出测试图像位置

  314.         }

  315.     if opt['business'] == 'drowning':

  316.         par={

  317.             'device':'0', ###显卡号,如果用TRT模型，只支持0（单显卡）

  318.             'labelnames':"../AIlib2/weights/conf/%s/labelnames.json"%(  opt['business']  ), ###检测类别对照表

  319.             'half':True,

  320.             'gpuname':'3090',###显卡名称

  321.             'max_workers':1, ###并行线程数

  322.             'Detweights':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu']  ),###检测模型路径

  323.             #'Detweights':"../AIlib2/weights/conf/highWay2/yolov5.pt",

  324.             'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,5,6,7,8,9]  ],###控制哪些检测类别显示、输出

  325.             'seg_nclass':2,###分割模型类别数目，默认2类

  326.             'segRegionCnt':2,###分割模型结果需要保留的等值线数目

  327.             '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,###分割模型预处理参数

  328.                       'mixFunction':{'function':mixDrowing_water_postprocess,

  329.                         'pars':{ }

  330.                         }

  331.                     },

  332. 

  333.             'trackPar':{'sort_max_age':2,'sort_min_hits':3,'sort_iou_thresh':0.2,'det_cnt':10,'windowsize':29,'patchCnt':100},

  334.             'Segweights' : "../weights/%s/AIlib2/%s/stdc_360X640_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu']  ),###分割模型权重位置

  335.             'postFile': '../AIlib2/weights/conf/%s/para.json'%('highWay2' ),###后处理参数文件

  336.             'txtFontSize':20,###文本字符的大小

  337.             'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'segLineShow':True,'waterLineWidth':2},###显示框、线设置

  338.             'testImgPath':'images/drowning/',###测试图像的位置

  339.             'testOutPath':'images/results/',###输出测试图像位置

  340.         }

  341.         par['segPar']['mixFunction']['pars']['modelSize'] = par['segPar']['modelSize']

  342.     

  343.           

  344.     ''' 以下是基于检测的跟踪模型，只有检测没有分割 '''

  345.     if opt['business'] == 'forest2':

  346.         par={

  347.             'device':'0', ###显卡号,如果用TRT模型，只支持0（单显卡）

  348.             'labelnames':"../AIlib2/weights/conf/forest2/labelnames.json", ###检测类别对照表

  349.             'gpuname':opt['gpu'],###显卡名称

  350.             'max_workers':1, ###并行线程数

  351.             'half':True,

  352.             'trtFlag_det':True,###检测模型是否采用TRT

  353.             'trtFlag_seg':False,###分割模型是否采用TRT

  354.             'Detweights':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu']  ),###检测模型路径

  355.             #'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,4,5,6]  ],###控制哪些检测类别显示、输出

  356.             'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in []  ],

  357.             'trackPar':{'sort_max_age':2,'sort_min_hits':3,'sort_iou_thresh':0.2,'det_cnt':10,'windowsize':29,'patchCnt':100},

  358.             'seg_nclass':2,###分割模型类别数目，默认2类

  359.             'segRegionCnt':0,###分割模型结果需要保留的等值线数目

  360.             'segPar':None,###分割模型预处理参数

  361.             'Segweights' : None,###分割模型权重位置

  362.             'postFile': '../AIlib2/weights/conf/forest/para.json',###后处理参数文件

  363.             'txtFontSize':80,###文本字符的大小

  364.             'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'waterLineWidth':3},###显示框、线设置

  365.             'testImgPath':'../AIdemo2/images/forest2/',###测试图像的位置

  366.             'testOutPath':'images/results/',###输出测试图像位置

  367.         }

  368.         

  369.     ###车辆巡检参数

  370.     if opt['business'] == 'vehicle': 

  371.         par={

  372.             'device':'0', ###显卡号,如果用TRT模型，只支持0（单显卡）

  373.             'labelnames':"../AIlib2/weights/conf/vehicle/labelnames.json", ###检测类别对照表

  374.             'gpuname':'2080T',###显卡名称

  375.             'half':True,

  376.             'max_workers':1, ###并行线程数

  377.             'trtFlag_det':True,###检测模型是否采用TRT

  378.             'trtFlag_seg':False,###分割模型是否采用TRT

  379.             'Detweights':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu']  ),###检测模型路径

  380.             'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,4,5,6]  ],###控制哪些检测类别显示、输出

  381.             'trackPar':{'sort_max_age':2,'sort_min_hits':3,'sort_iou_thresh':0.2,'det_cnt':10,'windowsize':29,'patchCnt':100},

  382.             'seg_nclass':2,###分割模型类别数目，默认2类

  383.             'segRegionCnt':0,###分割模型结果需要保留的等值线数目

  384.             'segPar':None,###分割模型预处理参数

  385.             'Segweights' : None,###分割模型权重位置

  386.             'postFile': '../AIlib2/weights/conf/vehicle/para.json',###后处理参数文件

  387.             'txtFontSize':40,###文本字符的大小

  388.             'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'waterLineWidth':3},###显示框、线设置

  389.             'testImgPath':'images/videos/vehicle/',###测试图像的位置

  390.             'testOutPath':'images/results/',###输出测试图像位置

  391.         }

  392.     ###行人检测模型

  393.     if opt['business'] == 'pedestrian': 

  394.         par={

  395.             'device':'0', ###显卡号,如果用TRT模型，只支持0（单显卡）

  396.             'labelnames':"../AIlib2/weights/conf/pedestrian/labelnames.json", ###检测类别对照表

  397.             'gpuname':'2080T',###显卡名称

  398.             'half':True,

  399.             'max_workers':1, ###并行线程数

  400.             'trtFlag_det':True,###检测模型是否采用TRT

  401.             'trtFlag_seg':False,###分割模型是否采用TRT

  402.             'Detweights':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu']  ),###检测模型路径

  403.             'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,4,5,6]  ],###控制哪些检测类别显示、输出

  404.             'trackPar':{'sort_max_age':2,'sort_min_hits':3,'sort_iou_thresh':0.2,'det_cnt':10,'windowsize':29,'patchCnt':100},

  405.             'seg_nclass':2,###分割模型类别数目，默认2类

  406.             'segRegionCnt':0,###分割模型结果需要保留的等值线数目

  407.             'segPar':None,###分割模型预处理参数

  408.             'Segweights' : None,###分割模型权重位置

  409.             'postFile': '../AIlib2/weights/conf/pedestrian/para.json',###后处理参数文件

  410.             'txtFontSize':40,###文本字符的大小

  411.             'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'waterLineWidth':3},###显示框、线设置

  412.             'testImgPath':'../AIdemo2/images/pedestrian/',###测试图像的位置

  413.             'testOutPath':'images/results/',###输出测试图像位置

  414.         }

  415.     if opt['business'] == 'smogfire': 

  416.         par={

  417.             'device':'0', ###显卡号,如果用TRT模型，只支持0（单显卡）

  418.             'labelnames':"../AIlib2/weights/conf/smogfire/labelnames.json", ###检测类别对照表

  419.             'gpuname':'2080T',###显卡名称

  420.             'half':True,

  421.             'max_workers':1, ###并行线程数

  422.             'trtFlag_det':True,###检测模型是否采用TRT

  423.             'trtFlag_seg':False,###分割模型是否采用TRT

  424.             'Detweights':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu']  ),###检测模型路径

  425.             'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,4,5,6]  ],###控制哪些检测类别显示、输出

  426.             'trackPar':{'sort_max_age':2,'sort_min_hits':3,'sort_iou_thresh':0.2,'det_cnt':10,'windowsize':29,'patchCnt':100},

  427.             'seg_nclass':2,###没有分割模型，此处不用

  428.             'segRegionCnt':0,###没有分割模型，此处不用

  429.             'segPar':None,###分割模型预处理参数

  430.             'Segweights' : None,###分割模型权重位置

  431.             'postFile': '../AIlib2/weights/conf/smogfire/para.json',###后处理参数文件

  432.             'txtFontSize':40,###文本字符的大小

  433.             'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'waterLineWidth':3},###显示框、线设置

  434.             'testImgPath':'../AIdemo2/images/smogfire/',###测试图像的位置

  435.             'testOutPath':'images/results/',###输出测试图像位置

  436.         }

  437.     ###钓鱼游泳检测

  438.     if opt['business'] == 'AnglerSwimmer': 

  439.         par={

  440.             'device':'0', ###显卡号,如果用TRT模型，只支持0（单显卡）

  441.             'labelnames':"../AIlib2/weights/conf/AnglerSwimmer/labelnames.json", ###检测类别对照表

  442.             'gpuname':'2080T',###显卡名称

  443.             'half':True,

  444.             'max_workers':1, ###并行线程数

  445.             'trtFlag_det':True,###检测模型是否采用TRT

  446.             'trtFlag_seg':False,###分割模型是否采用TRT

  447.             'Detweights':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu']  ),###检测模型路径

  448.             'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,4,5,6]  ],###控制哪些检测类别显示、输出

  449.             'trackPar':{'sort_max_age':2,'sort_min_hits':3,'sort_iou_thresh':0.2,'det_cnt':10,'windowsize':29,'patchCnt':100},

  450.             'seg_nclass':2,###没有分割模型，此处不用

  451.             'segRegionCnt':0,###没有分割模型，此处不用

  452.             'segPar':None,###分割模型预处理参数

  453.             'Segweights' : None,###分割模型权重位置

  454.             'postFile': '../AIlib2/weights/conf/AnglerSwimmer/para.json',###后处理参数文件

  455.             'txtFontSize':40,###文本字符的大小

  456.             'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'waterLineWidth':3},###显示框、线设置

  457.             'testImgPath':'../AIdemo2/images/AnglerSwimmer/',###测试图像的位置

  458.             'testOutPath':'images/results/',###输出测试图像位置

  459.         }

  460.     ###航道应急，做落水人员检测, channelEmergency    

  461.     if opt['business'] == 'channelEmergency': 

  462.         par={

  463.             'device':'0', ###显卡号,如果用TRT模型，只支持0（单显卡）

  464.             'labelnames':"../AIlib2/weights/conf/channelEmergency/labelnames.json", ###检测类别对照表

  465.             'gpuname':'2080T',###显卡名称

  466.             'half':True,

  467.             'max_workers':1, ###并行线程数

  468.             'trtFlag_det':True,###检测模型是否采用TRT

  469.             'trtFlag_seg':False,###分割模型是否采用TRT

  470.             'Detweights':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu']  ),###检测模型路径

  471.             #'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,4,5,6]  ],###控制哪些检测类别显示、输出

  472.             'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in []  ],###控制哪些检测类别显示、输出

  473.             'trackPar':{'sort_max_age':2,'sort_min_hits':3,'sort_iou_thresh':0.2,'det_cnt':10,'windowsize':29,'patchCnt':100},

  474.             'seg_nclass':2,###没有分割模型，此处不用

  475.             'segRegionCnt':0,###没有分割模型，此处不用

  476.             'segPar':None,###分割模型预处理参数

  477.             'Segweights' : None,###分割模型权重位置

  478.             'postFile': '../AIlib2/weights/conf/channelEmergency/para.json',###后处理参数文件

  479.             'txtFontSize':40,###文本字符的大小

  480.             'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'waterLineWidth':3},###显示框、线设置

  481.             'testImgPath':'../AIdemo2/images/channelEmergency/',###测试图像的位置

  482.             'testOutPath':'images/results/',###输出测试图像位置

  483.         }

  484.    

  485.    ###乡村路违法种植

  486.     if opt['business'] == 'countryRoad': 

  487.         par={

  488.             'device':'0', ###显卡号,如果用TRT模型，只支持0（单显卡）

  489.             'labelnames':"../AIlib2/weights/conf/countryRoad/labelnames.json", ###检测类别对照表

  490.             'gpuname':'2080T',###显卡名称

  491.             'half':True,

  492.             'max_workers':1, ###并行线程数

  493.             'trtFlag_det':True,###检测模型是否采用TRT

  494.             'trtFlag_seg':False,###分割模型是否采用TRT

  495.             'Detweights':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu']  ),###检测模型路径

  496.             'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,4,5,6]  ],###控制哪些检测类别显示、输出

  497.             'trackPar':{'sort_max_age':2,'sort_min_hits':3,'sort_iou_thresh':0.2,'det_cnt':10,'windowsize':29,'patchCnt':100},

  498.             'seg_nclass':2,###没有分割模型，此处不用

  499.             'segRegionCnt':0,###没有分割模型，此处不用

  500.             'segPar':None,###分割模型预处理参数

  501.             'Segweights' : None,###分割模型权重位置

  502.             'postFile': '../AIlib2/weights/conf/countryRoad/para.json',###后处理参数文件

  503.             'txtFontSize':40,###文本字符的大小

  504.             'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'waterLineWidth':3},###显示框、线设置

  505.             'testImgPath':'../AIdemo2/images/countryRoad/',###测试图像的位置

  506.             'testOutPath':'images/results/',###输出测试图像位置

  507.         }

  508.    

  509.     ###城管项目，检测城市垃圾和车辆

  510.     if opt['business'] == 'cityMangement':

  511.         par={

  512.             'device':'0', ###显卡号,如果用TRT模型，只支持0（单显卡）

  513.             'labelnames':"../AIlib2/weights/conf/%s/labelnames.json"%(opt['business']), ###检测类别对照表

  514.             'gpuname':'2080Ti',###显卡名称

  515.             'half':True,

  516.             'max_workers':1, ###并行线程数

  517.             'trtFlag_det':True,###检测模型是否采用TRT

  518.             'trtFlag_seg':False,###分割模型是否采用TRT

  519.             'Detweights':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu']  ),###检测模型路径

  520.             'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,4,5,6]  ],###控制哪些检测类别显示、输出

  521.             'trackPar':{'sort_max_age':2,'sort_min_hits':3,'sort_iou_thresh':0.2,'det_cnt':10,'windowsize':29,'patchCnt':100},

  522.             'seg_nclass':2,###没有分割模型，此处不用

  523.             'segRegionCnt':0,###没有分割模型，此处不用

  524.             'segPar':None,###分割模型预处理参数

  525.             'Segweights' : None,###分割模型权重位置

  526.             'postFile': '../AIlib2/weights/conf/%s/para.json'%(opt['business']),###后处理参数文件

  527.             'txtFontSize':40,###文本字符的大小

  528.             'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'waterLineWidth':3},###显示框、线设置

  529.             'testImgPath':'images/cityMangement',###测试图像的位置

  530.             'testOutPath':'images/results/',###输出测试图像位置

  531.         }

  532.     ###城管项目，检测道路情况，输入类别为五个："护栏","交通标志","非交通标志","施工","施工“(第4，第5类别合并，名称相同)

  533.     ###实际模型检测输出的类别为："护栏","交通标志","非交通标志","锥桶","水马"    

  534.     if opt['business'] == 'cityRoad':

  535.         par={

  536.             'device':'0', ###显卡号,如果用TRT模型，只支持0（单显卡）

  537.             'labelnames':"../AIlib2/weights/conf/%s/labelnames.json"%(opt['business']), ###检测类别对照表

  538.             'gpuname':'2080Ti',###显卡名称

  539.             'half':True,

  540.             'max_workers':1, ###并行线程数

  541.             'trtFlag_det':True,###检测模型是否采用TRT

  542.             'trtFlag_seg':False,###分割模型是否采用TRT

  543.             'Detweights':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu']  ),###检测模型路径

  544.             'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,4,5,6]  ],###控制哪些检测类别显示、输出

  545.             'trackPar':{'sort_max_age':2,'sort_min_hits':3,'sort_iou_thresh':0.2,'det_cnt':10,'windowsize':29,'patchCnt':100},

  546.             'seg_nclass':2,###没有分割模型，此处不用

  547.             'segRegionCnt':0,###没有分割模型，此处不用

  548.             'segPar':None,###分割模型预处理参数

  549.             'Segweights' : None,###分割模型权重位置

  550.             'postFile': '../AIlib2/weights/conf/%s/para.json'%(opt['business']),###后处理参数文件

  551.             'txtFontSize':40,###文本字符的大小

  552.             'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'waterLineWidth':3},###显示框、线设置

  553.             'testImgPath':'images/%s'%(opt['business']),###测试图像的位置

  554.             'testOutPath':'images/results/',###输出测试图像位置

  555.         }

  556.         

  557.     if opt['business'] == 'illParking':

  558.         from utilsK.illParkingUtils import illParking_postprocess

  559.         par={

  560.             'device':'0', ###显卡号,如果用TRT模型，只支持0（单显卡）

  561.             'half':True,

  562.             'labelnames':"../AIlib2/weights/conf/%s/labelnames.json"%(opt['business']), ###检测类别对照表

  563.             'max_workers':1, ###并行线程数

  564.             'Detweights':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu']  ),###检测模型路径

  565.             'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,5,6,7,8,9]  ],###控制哪些检测类别显示、输出

  566.             'trackPar':{'sort_max_age':2,'sort_min_hits':3,'sort_iou_thresh':0.2,'det_cnt':10,'windowsize':29,'patchCnt':100},

  567.             'seg_nclass':4,###没有分割模型，此处不用

  568.             'segRegionCnt':2,###没有分割模型，此处不用

  569.        

  570.             'segPar':{

  571.                       'mixFunction':{'function':illParking_postprocess,

  572.                         'pars':{ }

  573.                         }

  574.                     },

  575.                     

  576.             'Segweights' : None,###分割模型权重位置

  577.             'postFile': '../AIlib2/weights/conf/%s/para.json'%(opt['business'] ),###后处理参数文件

  578.             'txtFontSize':20,###文本字符的大小

  579.             'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'waterLineWidth':2},###显示框、线设置

  580.             'testImgPath':'images/cityMangement',###测试图像的位置

  581.             'testOutPath':'images/results/',###输出测试图像位置

  582.         }

  583.         

  584.         

  585.     par['trtFlag_det']=True if par['Detweights'].endswith('.engine') else False

  586.     if par['Segweights']:

  587.         par['segPar']['trtFlag_seg']=True if  par['Segweights'].endswith('.engine') else False

  588.         

  589.     ##使用森林，道路模型,business 控制['forest','road']

  590.     ##预先设置的参数

  591.     #gpuname=par['gpuname']#如果用trt就需要此参数，只能是"3090" "2080Ti"

  592.     device_=par['device']  ##选定模型，可选 cpu,'0','1'

  593.     

  594.     

  595.     device = select_device(device_)

  596.     half = device.type != 'cpu'  # half precision only supported on CUDA

  597.     trtFlag_det=par['trtFlag_det'] ###是否采用TRT模型加速

  598.     

  599.     ##以下参数目前不可改

  600.     imageW=1080     ####道路模型

  601.     digitFont= par['digitFont']

  602.     

  603.     

  604.     ####加载检测模型

  605.     if trtFlag_det:

  606.         Detweights=par['Detweights']

  607.         logger = trt.Logger(trt.Logger.ERROR)

  608.         with open(Detweights, "rb") as f, trt.Runtime(logger) as runtime:

  609.             model=runtime.deserialize_cuda_engine(f.read())# 输入trt本地文件，返回ICudaEngine对象

  610.         print('####load TRT model :%s'%(Detweights))    

  611.     else:

  612.         Detweights=par['Detweights']

  613.         model = attempt_load(Detweights, map_location=device)  # load FP32 model

  614.         if half: model.half()

  615. 

  616.     ####加载分割模型

  617.     seg_nclass = par['seg_nclass']

  618.     segPar=par['segPar']

  619.     if par['Segweights']:

  620.         if opt['business'] == 'cityMangement2':

  621.             segmodel = DMPRModel(weights=par['Segweights'], par = par['segPar'])

  622.         else:

  623.             segmodel = stdcModel(weights=par['Segweights'], par = par['segPar'])

  624.         '''    

  625.         if par['segPar']['trtFlag_seg']:

  626.             Segweights = par['Segweights']

  627.             logger = trt.Logger(trt.Logger.ERROR)

  628.             with open(Segweights, "rb") as f, trt.Runtime(logger) as runtime:

  629.                 segmodel=runtime.deserialize_cuda_engine(f.read())# 输入trt本地文件，返回ICudaEngine对象

  630.             print('############locad seg model trt success: ',Segweights)

  631.         else:

  632.             Segweights = par['Segweights']

  633.             segmodel = SegModel(nclass=seg_nclass,weights=Segweights,device=device)

  634.             print('############locad seg model pth success:',Segweights)

  635.         '''    

  636.     else:

  637.         segmodel=None

  638.         

  639. 

  640.     trackPar=par['trackPar']

  641.     sort_tracker = Sort(max_age=trackPar['sort_max_age'],

  642.                        min_hits=trackPar['sort_min_hits'],

  643.                        iou_threshold=trackPar['sort_iou_thresh'])

  644.                        

  645.         

  646.     labelnames = par['labelnames']

  647.     postFile= par['postFile']

  648.     print( Detweights,labelnames  ) 

  649.     conf_thres,iou_thres,classes,rainbows=get_postProcess_para(postFile)

  650.     

  651.     detPostPar = get_postProcess_para_dic(postFile)

  652.     conf_thres,iou_thres,classes,rainbows =  detPostPar["conf_thres"],detPostPar["iou_thres"],detPostPar["classes"],detPostPar["rainbows"]

  653.     if 'ovlap_thres_crossCategory' in detPostPar.keys(): iou2nd=detPostPar['ovlap_thres_crossCategory']

  654.     else:iou2nd = None

  655.     if 'score_byClass' in detPostPar.keys(): score_byClass=detPostPar['score_byClass']

  656.     else: score_byClass = None

  657.     

  658.     

  659.     

  660.     ####模型选择参数用如下：

  661.     mode_paras=par['detModelpara']

  662.     

  663.     allowedList,allowedList_string=get_needed_objectsIndex(mode_paras)    

  664.     #slopeIndex = par['slopeIndex']

  665.     ##只加载检测模型，准备好显示字符

  666.     

  667.     names=get_labelnames(labelnames)

  668.     #imageW=4915;###默认是1920，在森林巡检的高清图像中是4920

  669.     outfontsize=int(imageW/1920*40);###

  670.     label_arraylist = get_label_arrays(names,rainbows,outfontsize=par['txtFontSize'],fontpath="../AIlib2/conf/platech.ttf")

  671. 

  672. 

  673.     

  674.     

  675.     ##图像测试和视频

  676.     outpth = par['testOutPath'] 

  677.     impth = par['testImgPath']

  678.     imgpaths=[]###获取文件里所有的图像

  679.     videopaths=[]###获取文件里所有的视频

  680.     img_postfixs = ['.jpg','.JPG','.PNG','.png'];

  681.     vides_postfixs= ['.MP4','.mp4','.avi']

  682.     if os.path.isdir(impth):

  683.         for postfix in img_postfixs:

  684.             imgpaths.extend(glob.glob('%s/*%s'%(impth,postfix ))  )

  685.         for postfix in ['.MP4','.mp4','.avi']:

  686.             videopaths.extend(glob.glob('%s/*%s'%(impth,postfix ))  )

  687.     else:

  688.         postfix = os.path.splitext(impth)[-1]

  689.         if postfix in img_postfixs: imgpaths=[ impth ]

  690.         if postfix in vides_postfixs: videopaths = [impth ]

  691. 

  692.         

  693.         

  694.         

  695.     imgpaths.sort()  

  696.     

  697. 

  698. 

  699.     modelPar={ 'det_Model': model,'seg_Model':segmodel   }

  700. 

  701.     processPar={'half':par['half'],'device':device,'conf_thres':conf_thres,'iou_thres':iou_thres,'trtFlag_det':trtFlag_det,'iou2nd':iou2nd,'score_byClass':score_byClass}

  702.     drawPar={'names':names,'label_arraylist':label_arraylist,'rainbows':rainbows,'font': par['digitFont'],'allowedList':allowedList}

  703. 

  704.     for i in range(len(imgpaths)):

  705.     #for i in range(2):

  706.         #imgpath = os.path.join(impth, folders[i])

  707.         imgpath = imgpaths[i]

  708.         bname = os.path.basename(imgpath )

  709.         im0s=[cv2.imread(imgpath)]

  710.         time00 = time.time()

  711.         retResults,timeOut = AI_det_track_batch(im0s, [i] ,modelPar,processPar,sort_tracker ,trackPar,segPar)

  712.         #print('###line627:',retResults[2])

  713.         #retResults,timeInfos = AI_det_track_batch(imgarray_list, iframe_list ,par0['modelPar'],par0['processPar'],par0['sort_tracker'] ,par0['trackPar'],segPar=par0['segPar'])

  714.         if len(retResults[1])>0:

  715.             retResults[0][0] = drawBoxTraceSimplied(retResults[1],i, retResults[0][0],rainbows=rainbows,boxFlag=True,traceFlag=False,names=drawPar['names'])

  716.         time11 = time.time()

  717.         image_array = retResults[0][0]

  718.         '''

  719.         

  720.         返回值retResults[2] --list,其中每一个元素为一个list，表示每一帧的检测结果，每一个结果是由多个list构成，每个list表示一个框，格式为[ cls , x0 ,y0  ,x1  ,y1  ,conf,ifrmae,trackId ]

  721.                             --etc. retResults[2][j][k]表示第j帧的第k个框。

  722.         '''

  723.         cv2.imwrite(  os.path.join( outpth,bname ) ,image_array )

  724.         

  725.         print('----image:%s, process:%s ( %s ),save:%s'%(bname,(time11-time00) * 1000, timeOut,(time.time() - time11) * 1000)  )

  726.         

  727.     ##process video

  728. 

  729.     print('##begin to process videos, total %d videos'%( len(videopaths)))

  730.     for i,video in enumerate(videopaths):

  731.         print('process video%d :%s '%(i,video))

  732.         par0={'modelPar':modelPar,'processPar':processPar,'drawPar':drawPar,'outpth':par['testOutPath'], 'sort_tracker':sort_tracker,'trackPar':trackPar,'segPar':segPar}

  733.         process_video(video,par0,mode='track')  

  734.         

  735.       

  736. 

  737. def OCR_demo2(opt):

  738.     from ocrUtils2 import crnn_model

  739.     from ocrUtils2.ocrUtils import get_cfg,recognition_ocr,strLabelConverter

  740.     

  741.     if opt['business'] == 'ocr2':

  742.          par={

  743.             'image_dir':'images/ocr_en',

  744.             'outtxt':'images/results',

  745.             'weights':'../AIlib2/weights/conf/ocr2/crnn_448X32.pth',

  746. 

  747.             #'weights':'../weights/2080Ti/AIlib2/ocr2/crnn_2080Ti_fp16_448X32.engine',

  748.             'device':'cuda:0',

  749.             'cfg':'../AIlib2/weights/conf/ocr2/360CC_config.yaml',

  750.             

  751.             'char_file':'../AIlib2/weights/conf/ocr2/chars.txt',

  752.             'imgH':32,

  753.             'imgW':448,

  754.             'workers':1

  755. 

  756.             

  757.         }  

  758.     image_dir=par['image_dir']

  759.     outtxt=par['outtxt']

  760.     workers=par['workers']

  761.     weights= par['weights']

  762.     device=par['device']

  763. 

  764.     char_file=par['char_file']

  765.     imgH=par['imgH']

  766.     imgW=par['imgW']

  767.     cfg = par['cfg']

  768.     

  769.    

  770.     

  771.     config = get_cfg(cfg, char_file)    

  772. 

  773. 

  774.     par['contextFlag']=False

  775.     device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu') 

  776.     if weights.endswith('.pth'):

  777.         model = crnn_model.get_crnn(config,weights=weights).to(device)

  778.         par['model_mode']='pth'

  779.     else:    

  780.         logger = trt.Logger(trt.Logger.ERROR)

  781.         with open(weights, "rb") as f, trt.Runtime(logger) as runtime:

  782.             model = runtime.deserialize_cuda_engine(f.read())# 输入trt本地文件，返回ICudaEngine对象

  783.             print('#####load TRT file:',weights,'success #####')

  784.         context = model.create_execution_context() 

  785.         par['model_mode']='trt';par['contextFlag']=context

  786. 

  787.     converter = strLabelConverter(config.DATASET.ALPHABETS)

  788. 

  789.     img_urls=glob.glob('%s/*.jpg'%( image_dir ))

  790.     img_urls.extend( glob.glob('%s/*.png'%( image_dir ))  )

  791.     cnt=len(img_urls)

  792.     print('%s  has %d images'%(image_dir ,len(img_urls) )  )

  793.     # 准备数据

  794.     parList=[]

  795.     for i in range(cnt):

  796.         img_patch=cv2.imread( img_urls[i] ,   cv2.IMREAD_GRAYSCALE) 

  797.         started = time.time()

  798.         img = cv2.imread(img_urls[i])

  799.         sim_pred = recognition_ocr(config, img, model, converter, device,par=par)

  800.         finished = time.time()

  801.         print('{0}: elapsed time: {1} prd:{2} '.format( os.path.basename( img_urls[i] ), finished - started, sim_pred ))

  802.         

  803. 

  804.     

  805.     

  806.                 

  807. def OBB_track_demo(opt):

  808.     ###倾斜框（OBB）的ship目标检测

  809.     par={

  810.      

  811.         'obbModelPar':{

  812.             'model_size':(608,608),'K':100,'conf_thresh':0.3, 'down_ratio':4,'num_classes':15,'dataset':'dota',

  813.             'heads': {'hm': None,'wh': 10,'reg': 2,'cls_theta': 1},

  814.             'mean':(0.5, 0.5, 0.5),'std':(1, 1, 1), 'half': False,'decoder':None,

  815.             'weights':'../weights/%s/AIlib2/%s/obb_608X608_%s_fp16.engine'%(opt['gpu'],opt['business'],opt['gpu']),

  816.             

  817.         },

  818.         'outpth': 'images/results',

  819.         'trackPar':{'sort_max_age':2,'sort_min_hits':3,'sort_iou_thresh':0.2,'det_cnt':10,'windowsize':29,'patchCnt':100},

  820.         'device':"cuda:0",

  821.         

  822.         #'test_dir': '/mnt/thsw2/DSP2/videos/obbShips/DJI_20230208110806_0001_W_6M.MP4',

  823.         'test_dir':'/mnt/thsw2/DSP2/videos/obbShips/freighter2.mp4',

  824.         

  825.         'test_flag':True,

  826.         'postFile': '../AIlib2/weights/conf/%s/para.json'%(opt['business'] ),###后处理参数文件

  827.         'drawBox':True,#####是否画框

  828.         'drawPar': { 'digitWordFont' :{'line_thickness':2,'boxLine_thickness':1,'wordSize':40, 'fontSize':1.0,'label_location':'leftTop'}} ,  

  829.         'labelnames':"../AIlib2/weights/conf/%s/labelnames.json"%(opt['business'] ), ###检测类别对照表

  830.     }

  831.     #par['model_size'],par['mean'],par['std'],par['half'],par['saveType'],par['heads'],par['labelnames'],par['decoder'],par['down_ratio'],par['drawBox']

  832.     #par['rainbows'],par['label_array'],par['digitWordFont']

  833. 

  834.     obbModelPar = par['obbModelPar']

  835.     ####加载模型

  836.     model,decoder2=load_model_decoder_OBB(obbModelPar)

  837.     obbModelPar['decoder']=decoder2

  838. 

  839.     names=get_labelnames(par['labelnames']);obbModelPar['labelnames']=names

  840.     

  841.     _,_,_,rainbows=get_postProcess_para(par['postFile']);par['drawPar']['rainbows']=rainbows

  842. 

  843.     

  844.     label_arraylist = get_label_arrays(names,rainbows,outfontsize=par['drawPar']['digitWordFont']['wordSize'],fontpath="../AIlib2/conf/platech.ttf")

  845.     #par['label_array']=label_arraylist

  846. 

  847.     trackPar=par['trackPar']

  848.     sort_tracker = OBB_Sort(max_age=trackPar['sort_max_age'],

  849.                        min_hits=trackPar['sort_min_hits'],

  850.                        iou_threshold=trackPar['sort_iou_thresh'])

  851.                        

  852.     

  853.      ##图像测试和视频

  854.     impth = par['test_dir']

  855.     img_urls=[]###获取文件里所有的图像

  856.     video_urls=[]###获取文件里所有的视频

  857.     img_postfixs = ['.jpg','.JPG','.PNG','.png'];

  858.     vides_postfixs= ['.MP4','.mp4','.avi']

  859.     if os.path.isdir(impth):

  860.         for postfix in img_postfixs:

  861.             img_urls.extend(glob.glob('%s/*%s'%(impth,postfix ))  )

  862.         for postfix in ['.MP4','.mp4','.avi']:

  863.             video_urls.extend(glob.glob('%s/*%s'%(impth,postfix ))  )

  864.     else:

  865.         postfix = os.path.splitext(impth)[-1]

  866.         if postfix in img_postfixs: img_urls=[ impth ]

  867.         if postfix in vides_postfixs: video_urls = [impth ]

  868.         

  869.     parIn = {'obbModelPar':obbModelPar,'modelPar':{'obbmodel': model},'sort_tracker':sort_tracker,'outpth':par['outpth'],'trackPar':trackPar,'drawPar':par['drawPar']}

  870.     par['drawPar']['label_array']=label_arraylist        

  871.     for img_url in img_urls:

  872.         #print(img_url)

  873.         ori_image=cv2.imread(img_url)     

  874.         

  875.         #ori_image_list,infos = OBB_infer(model,ori_image,obbModelPar)         

  876.         

  877.         ori_image_list,infos  = OBB_tracker_batch([ori_image],[0],parIn['modelPar'],parIn['obbModelPar'],None,parIn['trackPar'],None)

  878.      

  879.         ori_image_list[1] =  draw_obb(ori_image_list[2] ,ori_image_list[1],par['drawPar'])

  880.         

  881.         imgName = os.path.basename(img_url)

  882.         saveFile = os.path.join(par['outpth'], imgName)

  883.         ret=cv2.imwrite(saveFile, ori_image_list[1])

  884.         if not ret:

  885.             print(saveFile, ' not created ')

  886.         print(  os.path.basename(img_url),':',infos,ori_image_list[2])

  887.         

  888.     ###处理视频

  889.     

  890.     for video_url in video_urls:

  891.         process_video(video_url, parIn ,mode='obbTrack')

  892. 

  893. def crowd_demo(opt):

  894.     if opt['business']=='crowdCounting':

  895.      

  896.         from crowd import crowdModel as Model

  897.         par={

  898.             'mean':[0.485, 0.456, 0.406], 'std':[0.229, 0.224, 0.225],'threshold':0.5, 

  899.             'input_profile_shapes':[(1,3,256,256),(1,3,1024,1024),(1,3,2048,2048)],

  900.             'modelPar':{'backbone':'vgg16_bn', 'gpu_id':0,'anchorFlag':False, 'width':None,'height':None ,'line':2, 'row':2},

  901. 

  902.             'weights':"../weights/%s/AIlib2/%s/crowdCounting_%s_dynamic.engine"%(opt['gpu'], opt['business'] ,opt['gpu']  ),###检测模型路径

  903.             'testImgPath':'images/%s'%(opt['business'] ),###测试图像的位置

  904.             'testOutPath':'images/results/',###输出测试图像位置

  905.             

  906.             }

  907.      

  908.     

  909.     #weights='weights/best_mae.pth'

  910.     cmodel = Model(par['weights'],par)

  911. 

  912.     img_path = par['testImgPath']

  913.     File = os.listdir(img_path)

  914.     targetList = []

  915.     for file in File[0:]:

  916.         COORlist = []

  917.         imgPath = img_path + os.sep + file

  918.         

  919.         

  920.         img_raw = cv2.cvtColor(cv2.imread(imgPath),cv2.COLOR_BGR2RGB)

  921.         # cmodel.eval---

  922.         # 输入读取的RGB数组

  923.         # 输出:list,0--原图，1-人头坐标list,2-对接OBB的格式数据，其中4个坐标均相同，2-格式如下：

  924.         # [ [ [ (x0,y0),(x1,y1),(x2,y2),(x3,y3)  ],score， cls ],  [ [ (x0,y0),(x1,y1),(x2,y2),(x3,y3)  ],score ，cls ],........  ]

  925.         

  926.         prets, infos = cmodel.eval(img_raw)

  927.         

  928.         

  929.         print(file,infos,' 人数：',len(prets[1]))

  930.         

  931.  

  932.         img_to_draw = cv2.cvtColor(np.array(img_raw), cv2.COLOR_RGB2BGR)

  933.         # 打印预测图像中人头的个数

  934.         for p in prets[1]:

  935.             img_to_draw = cv2.circle(img_to_draw, (int(p[0]), int(p[1])), 2, (0, 255, 0), -1)

  936.             COORlist.append((int(p[0]), int(p[1])))

  937.         # 将各测试图像中的人头坐标存储在targetList中, 格式：[[(x1, y1),(x2, y2),...], [(X1, Y1),(X2, Y2),..], ...]

  938.         targetList.append(COORlist)

  939.         #time.sleep(2)

  940.         # 保存预测图片

  941.         cv2.imwrite(os.path.join(par['testOutPath'], file), img_to_draw)

  942.  

  943. if __name__=="__main__":

  944. 

  945.     #jkm_demo()

  946.     businessAll=['river', 'river2','highWay2','noParking','drowning','forest2','vehicle','pedestrian','smogfire' , 'AnglerSwimmer','channelEmergency', 'countryRoad','cityMangement','ship2','cityMangement2','cityRoad','illParking',"crowdCounting"]

  947.     businessAll = ['crowdCounting']

  948.     

  949.     for busi in businessAll:

  950.         print('-'*40,'beg to test:',busi,'-'*40)

  951.         opt={'gpu':'2080Ti','business':busi}

  952.         if busi in ['ship2']:

  953.             OBB_track_demo(opt)

  954.         elif  opt['business'] in ['crowdCounting'] :

  955.             crowd_demo(opt)    

  956.         else:

  957.         #if opt['business'] in ['river','highWay2','noParking','drowning','']:

  958.             det_track_demo(opt ) 

  959.             

  960.   

  961.