gitadmin
/
tuoheng_alg_move_20240715
關注 1
收藏 0
複製 0
程式碼 問題 0 合併請求 0 版本發佈 11 Wiki 活動
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
127 提交
8 分支
目錄樹: 4ffcbbb90d
分支列表 標籤
${ item.name }
建立分支 ${ searchTerm }
從 '4ffcbbb90d'
${ noResults }
tuoheng_alg_move_20240715/test/ffmpeg11/Test.py

843 line
45KB
原始文件 Blame 歷史記錄 

  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. from segutils.segmodel import SegModel

  11. from segutils.trafficUtils import tracfficAccidentMixFunction

  12. from models.experimental import attempt_load

  13. from utils.torch_utils import select_device

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

  15. from ocrUtils.ocrUtils import CTCLabelConverter,AlignCollate

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

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

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

  19. from utilsK.noParkingUtils import mixNoParking_road_postprocess

  20. from obbUtils.load_obb_model import load_model_decoder_OBB

  21. import numpy as np

  22. import torch,glob

  23. import tensorrt as trt

  24. from utilsK.masterUtils import get_needed_objectsIndex

  25. from copy import deepcopy

  26. from scipy import interpolate

  27. from utilsK.drownUtils import mixDrowing_water_postprocess

  28. #import warnings

  29. #warnings.filterwarnings("error")

  30. 

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

  32.     rate = num / total

  33.     time_n=time.time()

  34.     rate_num = int(rate * 30)

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

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

  37.     sys.stdout.write(r)

  38.     sys.stdout.flush()

  39. 

  40. 

  41. '''

  42.     多线程

  43. '''

  44. 

  45. def process_v1(frame):

  46.     #try:

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

  48.     time00 = time.time()

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

  50. 

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

  52. 

  53.     time11 = time.time()

  54.     image_array = p_result[1]

  55. 

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

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

  58.     if len(p_result)==5:

  59.         image_mask = p_result[4]

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

  61. 

  62.     boxes=p_result[2]

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

  64.         for box in boxes:

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

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

  67.             fp.write(out_str)

  68.     time22 = time.time()

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

  70.     return 'success'

  71.     #except Exception as e:

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

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

  74.     cap=cv2.VideoCapture(video)

  75.     if not  cap.isOpened():

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

  77.         return False

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

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

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

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

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

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

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

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

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

  87.     num=0

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

  89. 

  90.     imgarray_list = []; iframe_list = []

  91. 

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

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

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

  95. 

  96. 

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

  98.     while cap.isOpened():

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

  100.         iframe +=1

  101.         if not ret:break

  102.         if mode=='detSeg':

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

  104.         elif mode == 'track':

  105.             #sampleCount=10

  106.             imgarray_list.append( imgarray )

  107.             iframe_list.append(iframe )

  108.             if iframe%patch_cnt==0:

  109.                 time_patch0 = time.time()

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

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

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

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

  114.                 time_patch2 = time.time()

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

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

  117.                     img_draw = imgarray_list[ iiframe- frame_min  ]

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

  119.                     ret = vid_writer_AI.write(img_draw)

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

  121.                 imgarray_list=[];iframe_list=[]

  122.         elif mode =='obbTrack':

  123.             imgarray_list.append( imgarray )

  124.             iframe_list.append(iframe )

  125.             if iframe%patch_cnt==0:

  126.                 time_patch0 = time.time()

  127. 

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

  129.                 print(  timeInfos )

  130. 

  131.                 #对结果画图

  132.                 track_det_np = track_det_results[1]

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

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

  135.                     img_draw = imgarray_list[ iiframe- frame_min  ]

  136. 

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

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

  139.                         if  True:

  140.                             frameIdex=12;trackIdex=13;

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

  142. 

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

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

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

  146.                             for trackId in track_ids:

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

  148.                                 xcs = boxes_before_oneFrame_oneId[:,8]

  149.                                 ycs = boxes_before_oneFrame_oneId[:,9]

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

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

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

  153. 

  154.                     ret = vid_writer_AI.write(img_draw)

  155. 

  156.                     #sys.exit(0)

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

  158.                 imgarray_list=[];iframe_list=[]

  159. 

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

  161. 

  162.         else:

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

  164. 

  165. 

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

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

  168.             ret = vid_writer_AI.write(image_array)

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

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

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

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

  173.                 post_results=[]

  174. 

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

  176.                 post_results.append(p_result)

  177. 

  178.     vid_writer_AI.release();

  179. 

  180. def det_track_demo(business, videopaths):

  181.     '''

  182.     跟踪参数说明：

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

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

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

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

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

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

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

  190.     '''

  191. 

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

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

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

  195.         par={

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

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

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

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

  200.             'half':True,

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

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

  203. 

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

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

  206. 

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

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

  209.                       },

  210. 

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

  212. 

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

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

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

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

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

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

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

  220. 

  221.         }

  222. 

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

  224.         par={

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

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

  227.             'half':True,

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

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

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

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

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

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

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

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

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

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

  238.                                      }

  239.                       },

  240. 

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

  242.             'mode':'highWay3.0',

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

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

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

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

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

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

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

  250.         }

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

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

  253.         par={

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

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

  256.             'half':True,

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

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

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

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

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

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

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

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

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

  266.                                      'pars':

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

  268. 

  269.                                          {'RoadArea': 16000, 'roadVehicleAngle': 15,'radius': 50, 'distanceFlag': False, 'vehicleFlag': False}

  270.                                      }

  271.                       },

  272. 

  273. 

  274. 

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

  276.             'mode':'highWay3.0',

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

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

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

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

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

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

  283.         }

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

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

  286.         par={

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

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

  289.             'half':True,

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

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

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

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

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

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

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

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

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

  299.                                      'pars':{ }

  300.                                      }

  301.                       },

  302. 

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

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

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

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

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

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

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

  310.         }

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

  312. 

  313. 

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

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

  316.         par={

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

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

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

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

  321.             'half':True,

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

  337.         }

  338. 

  339.     ###车辆巡检参数

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

  341.         par={

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

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

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

  345.             'half':True,

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

  361.         }

  362.     ###行人检测模型

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

  364.         par={

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

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

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

  368.             'half':True,

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

  384.         }

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

  386.         par={

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

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

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

  390.             'half':True,

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

  406.         }

  407.     ###钓鱼游泳检测

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

  409.         par={

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

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

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

  413.             'half':True,

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

  429.         }

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

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

  432.         par={

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

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

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

  436.             'half':True,

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

  453.         }

  454. 

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

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

  457.         par={

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

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

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

  461.             'half':True,

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

  477.         }

  478. 

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

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

  481.         par={

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

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

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

  485.             'half':True,

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

  501.         }

  502. 

  503. 

  504. 

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

  506.     if par['Segweights']:

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

  508. 

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

  510.     ##预先设置的参数

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

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

  513. 

  514. 

  515.     device = select_device(device_)

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

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

  518. 

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

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

  521.     digitFont= par['digitFont']

  522. 

  523. 

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

  525.     if trtFlag_det:

  526.         Detweights=par['Detweights']

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

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

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

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

  531.     else:

  532.         Detweights=par['Detweights']

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

  534.         if half: model.half()

  535. 

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

  537.     seg_nclass = par['seg_nclass']

  538.     segPar=par['segPar']

  539.     if par['Segweights']:

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

  541.             Segweights = par['Segweights']

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

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

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

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

  546.         else:

  547.             Segweights = par['Segweights']

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

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

  550.     else:

  551.         segmodel=None

  552. 

  553. 

  554.     trackPar=par['trackPar']

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

  556.                         min_hits=trackPar['sort_min_hits'],

  557.                         iou_threshold=trackPar['sort_iou_thresh'])

  558. 

  559. 

  560.     labelnames = par['labelnames']

  561.     postFile= par['postFile']

  562.     print( Detweights,labelnames  )

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

  564. 

  565.     detPostPar = get_postProcess_para_dic(postFile)

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

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

  568.     else:iou2nd = None

  569. 

  570. 

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

  572.     mode_paras=par['detModelpara']

  573. 

  574.     allowedList,allowedList_string=get_needed_objectsIndex(mode_paras)

  575.     #slopeIndex = par['slopeIndex']

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

  577. 

  578.     names=get_labelnames(labelnames)

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

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

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

  582. 

  583. 

  584. 

  585. 

  586.     ##图像测试和视频

  587.     outpth = par['testOutPath']

  588.     impth = par['testImgPath']

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

  590.     videopaths=videopaths###获取文件里所有的视频

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

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

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

  594.         for postfix in img_postfixs:

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

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

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

  598.     else:

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

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

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

  602. 

  603. 

  604. 

  605. 

  606.     imgpaths.sort()

  607. 

  608. 

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

  610. 

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

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

  613. 

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

  615.         #for i in range(2):

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

  617.         imgpath = imgpaths[i]

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

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

  620.         time00 = time.time()

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

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

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

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

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

  626.         time11 = time.time()

  627.         image_array = retResults[0][0]

  628.         '''

  629.         

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

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

  632.         '''

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

  634. 

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

  636. 

  637.     ##process video

  638. 

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

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

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

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

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

  644. 

  645. 

  646. 

  647. def OCR_demo2(opt):

  648.     from ocrUtils2 import crnn_model

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

  650. 

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

  652.         par={

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

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

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

  656. 

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

  658.             'device':'cuda:0',

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

  660. 

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

  662.             'imgH':32,

  663.             'imgW':448,

  664.             'workers':1

  665. 

  666. 

  667.         }

  668.     image_dir=par['image_dir']

  669.     outtxt=par['outtxt']

  670.     workers=par['workers']

  671.     weights= par['weights']

  672.     device=par['device']

  673. 

  674.     char_file=par['char_file']

  675.     imgH=par['imgH']

  676.     imgW=par['imgW']

  677.     cfg = par['cfg']

  678. 

  679. 

  680. 

  681.     config = get_cfg(cfg, char_file)

  682. 

  683. 

  684.     par['contextFlag']=False

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

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

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

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

  689.     else:

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

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

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

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

  694.         context = model.create_execution_context()

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

  696. 

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

  698. 

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

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

  701.     cnt=len(img_urls)

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

  703.     # 准备数据

  704.     parList=[]

  705.     for i in range(cnt):

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

  707.         started = time.time()

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

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

  710.         finished = time.time()

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

  712. 

  713. 

  714. 

  715. 

  716. 

  717. def OBB_track_demo(opt):

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

  719.     '''

  720.     par={

  721.         'model_size':(608,608), #width,height

  722.         'K':100, #Maximum of objects'

  723.         'conf_thresh':0.18,##Confidence threshold, 0.1 for general evaluation

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

  725. 

  726.         'down_ratio':4,'num_classes':15,

  727.         #'weights':'../AIlib2/weights/conf/ship2/obb_608X608.engine',

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

  729.         'dataset':'dota',

  730.         'test_dir': '/mnt/thsw2/DSP2/videos/obbShips',

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

  732.         'half': False,

  733.         'mean':(0.5, 0.5, 0.5),

  734.         'std':(1, 1, 1),

  735.         'model_size':(608,608),##width,height

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

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

  738.         'decoder':None,

  739.         'test_flag':True,

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

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

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

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

  744.     }

  745.     '''

  746.     par={

  747. 

  748.         'obbModelPar':{

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

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

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

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

  753. 

  754.         },

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

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

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

  758. 

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

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

  761. 

  762.         'test_flag':True,

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

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

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

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

  767.     }

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

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

  770. 

  771.     obbModelPar = par['obbModelPar']

  772.     ####加载模型

  773.     model,decoder2=load_model_decoder_OBB(obbModelPar)

  774.     obbModelPar['decoder']=decoder2

  775. 

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

  777. 

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

  779. 

  780. 

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

  782.     #par['label_array']=label_arraylist

  783. 

  784.     trackPar=par['trackPar']

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

  786.                             min_hits=trackPar['sort_min_hits'],

  787.                             iou_threshold=trackPar['sort_iou_thresh'])

  788. 

  789. 

  790.     ##图像测试和视频

  791.     impth = par['test_dir']

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

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

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

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

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

  797.         for postfix in img_postfixs:

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

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

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

  801.     else:

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

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

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

  805. 

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

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

  808.     for img_url in img_urls:

  809.         #print(img_url)

  810.         ori_image=cv2.imread(img_url)

  811. 

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

  813. 

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

  815. 

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

  817. 

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

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

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

  821.         if not ret:

  822.             print(saveFile, ' not created ')

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

  824. 

  825.     ###处理视频

  826. 

  827.     for video_url in video_urls:

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

  829. 

  830. if __name__=="__main__":

  831. 

  832.     #jkm_demo()

  833.     #businessAll=['river', 'river2','highWay2','noParking','drowning','forest2','vehicle','pedestrian','smogfire' , 'AnglerSwimmer','channelEmergency', 'countryRoad','cityMangement','ship2']

  834.     businessAll = ['river2']

  835.     videopaths = ['/home/th/tuo_heng/dev/DJI_20211229100908_0002_S.mp4']

  836.     for busi in businessAll:

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

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

  839.         if busi in ['ship2']:

  840.             OBB_track_demo(opt)

  841.         else:

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

  843.             det_track_demo(opt, videopaths)