Nevar pievienot vairāk kā 25 tēmas Tēmai ir jāsākas ar burtu vai ciparu, tā var saturēt domu zīmes ('-') un var būt līdz 35 simboliem gara.

1141 rinda
67KB

  1. import sys, yaml
  2. from easydict import EasyDict as edict
  3. from concurrent.futures import ThreadPoolExecutor
  4. sys.path.extend(['..','../AIlib2' ])
  5. from AI import AI_process,AI_process_forest,get_postProcess_para,get_postProcess_para_dic,ocr_process,AI_det_track,AI_det_track_batch,AI_det_track_batch_N,get_images_videos,default_mix
  6. import cv2,os,time
  7. from segutils.segmodel import SegModel
  8. from yolov5 import yolov5Model
  9. from stdc import stdcModel
  10. from segutils.trafficUtils import tracfficAccidentMixFunction
  11. from models.experimental import attempt_load
  12. from utils.torch_utils import select_device
  13. from utilsK.queRiver import get_labelnames,get_label_arrays,save_problem_images,riverDetSegMixProcess,draw_painting_joint
  14. from ocrUtils.ocrUtils import CTCLabelConverter,AlignCollate
  15. from trackUtils.sort import Sort,track_draw_boxAndTrace,track_draw_trace_boxes,moving_average_wang,drawBoxTraceSimplied
  16. from trackUtils.sort_obb import OBB_Sort,obbTohbb,track_draw_all_boxes,track_draw_trace
  17. from obbUtils.shipUtils import OBB_infer,OBB_tracker,draw_obb,OBB_tracker_batch
  18. from utilsK.noParkingUtils import mixNoParking_road_postprocess
  19. from obbUtils.load_obb_model import load_model_decoder_OBB
  20. import numpy as np
  21. import torch,glob
  22. import tensorrt as trt
  23. from utilsK.masterUtils import get_needed_objectsIndex
  24. from copy import deepcopy
  25. from scipy import interpolate
  26. from utilsK.drownUtils import mixDrowing_water_postprocess
  27. #import warnings
  28. #warnings.filterwarnings("error")
  29. rainbows=[ [0,0,255],[0,255,0],[255,0,0],[255,0,255],[255,255,0],[255,129,0],[255,0,127],[127,255,0],[0,255,127],[0,127,255],[127,0,255],[255,127,255],[255,255,127],[127,255,255],[0,255,255],[255,127,255],[127,255,255], [0,127,0],[0,0,127],[0,255,255]]
  30. def drawAllBox(preds,imgDraw,label_arraylist,rainbows,font):
  31. for box in preds:
  32. #cls,conf,xyxy = box[0],box[5], box[1:5]
  33. #print('#'*20,'line47',box)
  34. cls,conf,xyxy = box[5],box[4], box[0:4] ##2023.08.03,修改了格式
  35. #print('#####line46 demo.py:', cls,conf,xyxy, len(label_arraylist),len(rainbows) )
  36. imgDraw = draw_painting_joint(xyxy,imgDraw,label_arraylist[int(cls)],score=conf,color=rainbows[int(cls)%20],font=font,socre_location="leftBottom")
  37. return imgDraw
  38. def get_drawPar(par):
  39. labelnames = par['labelnames']
  40. names=get_labelnames(labelnames)
  41. mode_paras=par['detModelpara']
  42. allowedList,allowedList_string=get_needed_objectsIndex(mode_paras)
  43. outfontsize=int(1080/1920*40);###
  44. label_arraylist = get_label_arrays(names,rainbows,outfontsize=par['txtFontSize'],fontpath="../AIlib2/conf/platech.ttf")
  45. drawPar={'names':names,'label_arraylist':label_arraylist,'rainbows':rainbows,'font': par['digitFont'],'allowedList':allowedList}
  46. return drawPar
  47. def view_bar(num, total,time1,prefix='prefix'):
  48. rate = num / total
  49. time_n=time.time()
  50. rate_num = int(rate * 30)
  51. rate_nums = np.round(rate * 100)
  52. r = '\r %s %d / %d [%s%s] %.2f s'%(prefix,num,total, ">" * rate_num, " " * (30 - rate_num), time_n-time1 )
  53. sys.stdout.write(r)
  54. sys.stdout.flush()
  55. '''
  56. 多线程
  57. '''
  58. def process_v1(frame):
  59. #try:
  60. print('demo.py beging to :',frame[8])
  61. time00 = time.time()
  62. H,W,C = frame[0][0].shape
  63. 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])
  64. time11 = time.time()
  65. image_array = p_result[1]
  66. cv2.imwrite(os.path.join('images/results/',frame[8] ) ,image_array)
  67. bname = frame[8].split('.')[0]
  68. if len(p_result)==5:
  69. image_mask = p_result[4]
  70. cv2.imwrite(os.path.join('images/results/',bname+'_mask.png' ) , (image_mask).astype(np.uint8))
  71. boxes=p_result[2]
  72. with open( os.path.join('images/results/',bname+'.txt' ),'w' ) as fp:
  73. for box in boxes:
  74. box_str=[str(x) for x in box]
  75. out_str=','.join(box_str)+'\n'
  76. fp.write(out_str)
  77. time22 = time.time()
  78. 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)
  79. return 'success'
  80. #except Exception as e:
  81. # return 'failed:'+str(e)
  82. def process_video(video,par0,mode='detSeg'):
  83. cap=cv2.VideoCapture(video)
  84. if not cap.isOpened():
  85. print('#####error url:',video)
  86. return False
  87. bname=os.path.basename(video).split('.')[0]
  88. fps = int(cap.get(cv2.CAP_PROP_FPS)+0.5)
  89. width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH )+0.5)
  90. height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)+0.5)
  91. framecnt=int(cap.get(7)+0.5)
  92. save_path_AI = os.path.join(par0['outpth'],os.path.basename(video))
  93. problem_image_dir= os.path.join( par0['outpth'], 'probleImages' )
  94. os.makedirs(problem_image_dir,exist_ok=True)
  95. vid_writer_AI = cv2.VideoWriter(save_path_AI, cv2.VideoWriter_fourcc(*'mp4v'), fps, (width,height))
  96. num=0
  97. iframe=0;post_results=[];fpsample=30*10
  98. imgarray_list = []; iframe_list = []
  99. patch_cnt = par0['trackPar']['patchCnt']
  100. ##windowsize 对逐帧插值后的结果做平滑,windowsize为平滑的长度,没隔det_cnt帧做一次跟踪。
  101. trackPar={'det_cnt':10,'windowsize':29 }
  102. ##track_det_result_update= np.empty((0,8)) ###每100帧跑出来的结果,放在track_det_result_update,只保留当前100帧里有的tracker Id.
  103. while cap.isOpened():
  104. ret, imgarray = cap.read() #读取摄像头画面
  105. iframe +=1
  106. if not ret:break
  107. if mode=='detSeg':
  108. 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'])
  109. elif mode == 'track':
  110. #sampleCount=10
  111. imgarray_list.append( imgarray )
  112. iframe_list.append(iframe )
  113. if iframe%patch_cnt==0:
  114. time_patch0 = time.time()
  115. #retResults,timeInfos = AI_det_track_batch(imgarray_list, iframe_list ,par0['modelPar'],par0['processPar'],par0['sort_tracker'] ,par0['trackPar'],segPar=par0['segPar'])
  116. retResults,timeInfos = AI_det_track_batch_N(imgarray_list, iframe_list ,par0['modelList'],par0['postProcess'],par0['sort_tracker'],par0['trackPar'])
  117. print('###line111:',retResults[2])
  118. ###需要保存成一个二维list,每一个list是一帧检测结果。
  119. ###track_det_result 内容格式:x1, y1, x2, y2, conf, cls,iframe,trackId
  120. time_patch2 = time.time()
  121. frame_min = iframe_list[0];frame_max=iframe_list[-1]
  122. for iiframe in range(frame_min,frame_max+1):
  123. img_draw = imgarray_list[ iiframe- frame_min ]
  124. img_draw = drawBoxTraceSimplied(retResults[1] ,iiframe, img_draw,rainbows=par0['drawPar']['rainbows'],boxFlag=True,traceFlag=True,names=par0['drawPar']['names'] )
  125. ret = vid_writer_AI.write(img_draw)
  126. view_bar(iiframe, framecnt,time.time(),prefix=os.path.basename(video))
  127. imgarray_list=[];iframe_list=[]
  128. elif mode =='obbTrack':
  129. imgarray_list.append( imgarray )
  130. iframe_list.append(iframe )
  131. if iframe%patch_cnt==0:
  132. time_patch0 = time.time()
  133. track_det_results, timeInfos = OBB_tracker_batch(imgarray_list,iframe_list,par0['modelPar'],par0['obbModelPar'],par0['sort_tracker'],par0['trackPar'],segPar=None)
  134. print( timeInfos )
  135. #对结果画图
  136. track_det_np = track_det_results[1]
  137. frame_min = iframe_list[0];frame_max=iframe_list[-1]
  138. for iiframe in range(frame_min,frame_max+1):
  139. img_draw = imgarray_list[ iiframe- frame_min ]
  140. if len( track_det_results[2][ iiframe- frame_min]) > 0:
  141. img_draw = draw_obb( track_det_results[2][iiframe- frame_min ] ,img_draw,par0['drawPar'])
  142. if True:
  143. frameIdex=12;trackIdex=13;
  144. boxes_oneFrame = track_det_np[ track_det_np[:,frameIdex]==iiframe ]
  145. ###在某一帧上,画上轨迹
  146. track_ids = boxes_oneFrame[:,trackIdex].tolist()
  147. boxes_before_oneFrame = track_det_np[ track_det_np[:,frameIdex]<=iiframe ]
  148. for trackId in track_ids:
  149. boxes_before_oneFrame_oneId = boxes_before_oneFrame[boxes_before_oneFrame[:,trackIdex]==trackId]
  150. xcs = boxes_before_oneFrame_oneId[:,8]
  151. ycs = boxes_before_oneFrame_oneId[:,9]
  152. [cv2.line(img_draw, ( int(xcs[i]) , int(ycs[i]) ),
  153. ( int(xcs[i+1]),int(ycs[i+1]) ),(255,0,0), thickness=2)
  154. for i,_ in enumerate(xcs) if i < len(xcs)-1 ]
  155. ret = vid_writer_AI.write(img_draw)
  156. #sys.exit(0)
  157. #print('vide writer ret:',ret)
  158. imgarray_list=[];iframe_list=[]
  159. view_bar(iframe, framecnt,time.time(),prefix=os.path.basename(video))
  160. else:
  161. 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'])
  162. if mode not in [ 'track','obbTrack']:
  163. image_array = p_result[1];num+=1
  164. ret = vid_writer_AI.write(image_array)
  165. view_bar(num, framecnt,time.time(),prefix=os.path.basename(video))
  166. ##每隔 fpsample帧处理一次,如果有问题就保存图片
  167. if (iframe % fpsample == 0) and (len(post_results)>0) :
  168. parImage=save_problem_images(post_results,iframe,par0['names'],streamName=bname,outImaDir=problem_image_dir,imageTxtFile=False)
  169. post_results=[]
  170. if len(p_result[2] )>0:
  171. post_results.append(p_result)
  172. vid_writer_AI.release();
  173. def det_track_demo_N(business ):
  174. ####河道巡检的跟踪模型参数
  175. if opt['business'] == 'river' or opt['business'] == 'river2' or opt['business'] == 'riverT':
  176. from utilsK.queRiver import riverDetSegMixProcess_N
  177. par={
  178. 'device':'0', ###显卡号,如果用TRT模型,只支持0(单显卡)
  179. 'labelnames':"../AIlib2/weights/conf/%s/labelnames.json"%(opt['business']), ###检测类别对照表
  180. 'max_workers':1, ###并行线程数
  181. 'trackPar':{'sort_max_age':2,'sort_min_hits':3,'sort_iou_thresh':0.2,'det_cnt':10,'windowsize':29,'patchCnt':100},
  182. 'postProcess':{'function':riverDetSegMixProcess_N,'pars':{'slopeIndex':[1,3,4,7], 'riverIou':0.1}}, #分割和检测混合处理的函数
  183. 'models':
  184. [
  185. {
  186. 'weight':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###检测模型路径
  187. #'weight':'../AIlib2/weights/conf/%s/yolov5.pt'%(opt['business'] ),
  188. 'name':'yolov5',
  189. 'model':yolov5Model,
  190. 'par':{ 'half':True,'device':'cuda:0' ,'conf_thres':0.25,'iou_thres':0.45,'allowedList':[0,1,2,3],'segRegionCnt':1, 'trtFlag_det':False,'trtFlag_seg':False, "score_byClass":{"0":0.25,"1":0.3,"2":0.3,"3":0.3 } },
  191. },
  192. {
  193. 'weight':"../weights/%s/AIlib2/%s/stdc_360X640_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###检测模型路径
  194. #'weight':'../AIlib2/weights/conf/%s/stdc_360X640.pth'%(opt['business'] ),
  195. 'par':{
  196. 'modelSize':(640,360),'mean':(0.485, 0.456, 0.406),'std' :(0.229, 0.224, 0.225),'numpy':False, 'RGB_convert_first':True,'seg_nclass':2},###分割模型预处理参数
  197. 'model':stdcModel,
  198. 'name':'stdc'
  199. }
  200. ],
  201. 'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,4,5,6] ],###控制哪些检测类别显示、输出
  202. 'postFile': '../AIlib2/weights/conf/%s/para.json'%( opt['business'] ),###后处理参数文件
  203. 'txtFontSize':80,###文本字符的大小
  204. 'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'waterLineWidth':3},###显示框、线设置
  205. #'testImgPath':'images/videos/river',###测试图像的位置
  206. 'testImgPath':'images/%s'%(opt['business']),###测试图像的位置
  207. 'testOutPath':'images/results/',###输出测试图像位置
  208. }
  209. if opt['business'] == 'highWay2':
  210. from segutils.trafficUtils import tracfficAccidentMixFunction_N
  211. par={
  212. 'device':'0', ###显卡号,如果用TRT模型,只支持0(单显卡)
  213. 'labelnames':"../AIlib2/weights/conf/%s/labelnames.json"%( opt['business'] ), ###检测类别对照表
  214. 'max_workers':1, ###并行线程数
  215. 'trackPar':{'sort_max_age':2,'sort_min_hits':3,'sort_iou_thresh':0.2,'det_cnt':5,'windowsize':29,'patchCnt':100},
  216. 'postProcess':{'function':tracfficAccidentMixFunction_N,
  217. 'pars':{ 'RoadArea': 16000, 'vehicleArea': 10, 'roadVehicleAngle': 15, 'speedRoadVehicleAngleMax': 75,'radius': 50 , 'roundness': 1.0, 'cls': 10, 'vehicleFactor': 0.1,'cls':9, 'confThres':0.25,'roadIou':0.6,'vehicleFlag':False,'distanceFlag': False }
  218. },
  219. 'models':
  220. [
  221. {
  222. 'weight':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###检测模型路径
  223. #'weight':'../AIlib2/weights/conf/%s/yolov5.pt'%(opt['business'] ),
  224. 'name':'yolov5',
  225. 'model':yolov5Model,
  226. 'par':{ 'half':True,'device':'cuda:0' ,'conf_thres':0.25,'iou_thres':0.45,'allowedList':[0,1,2,3],'segRegionCnt':1, 'trtFlag_det':False,'trtFlag_seg':False, "score_byClass":{"0":0.25,"1":0.3,"2":0.3,"3":0.3 } },
  227. },
  228. {
  229. #'weight':"../weights/%s/AIlib2/%s/dmpr_%s.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###检测模型路径
  230. 'weight':'../AIlib2/weights/conf/%s/stdc_360X640.pth'%(opt['business'] ),
  231. 'par':{
  232. '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,'seg_nclass':3},###分割模型预处理参数
  233. 'model':stdcModel,
  234. 'name':'stdc'
  235. }
  236. ],
  237. 'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,5,6,7,8,9] ],###控制哪些检测类别显示、输出
  238. 'postFile': '../AIlib2/weights/conf/%s/para.json'%(opt['business'] ),###后处理参数文件
  239. 'txtFontSize':20,###文本字符的大小
  240. 'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':0.5,'waterLineColor':(0,255,255),'segLineShow':True,'waterLineWidth':2},###显示框、线设置
  241. #'testImgPath':'images/trafficAccident/8.png',###测试图像的位置
  242. #'testImgPath':'/home/chenyukun/777-7-42.mp4',###测试图像的位置
  243. 'testImgPath':'/mnt/thsw2/DSP2/weights/highWay2/images/',
  244. 'testOutPath':'images/results/',###输出测试图像位置
  245. }
  246. par['postProcess']['pars']['modelSize'] = par['models'][1]['par']['modelSize']
  247. if opt['business'] == 'noParking':
  248. from utilsK.noParkingUtils import mixNoParking_road_postprocess_N
  249. par={
  250. 'device':'0', ###显卡号,如果用TRT模型,只支持0(单显卡)
  251. 'labelnames':"../AIlib2/weights/conf/%s/labelnames.json"%( opt['business'] ), ###检测类别对照表
  252. 'max_workers':1, ###并行线程数
  253. 'trackPar':{'sort_max_age':2,'sort_min_hits':3,'sort_iou_thresh':0.2,'det_cnt':10,'windowsize':29,'patchCnt':100},
  254. 'postProcess':{'function':mixNoParking_road_postprocess_N,
  255. 'pars': { 'roundness': 0.3, 'cls': 9, 'laneArea': 10, 'laneAngleCha': 5 ,'RoadArea': 16000,'fitOrder':2}
  256. } ,
  257. 'models':
  258. [
  259. {
  260. 'weight':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###检测模型路径
  261. #'weight':'../AIlib2/weights/conf/%s/yolov5.pt'%(opt['business'] ),
  262. 'name':'yolov5',
  263. 'model':yolov5Model,
  264. 'par':{ 'half':True,'device':'cuda:0' ,'conf_thres':0.25,'iou_thres':0.45,'allowedList':[0,1,2,3],'segRegionCnt':1, 'trtFlag_det':False,'trtFlag_seg':False, "score_byClass":{"0":0.25,"1":0.3,"2":0.3,"3":0.3 } },
  265. },
  266. {
  267. #'weight':"../weights/%s/AIlib2/%s/dmpr_%s.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###检测模型路径
  268. 'weight':'../AIlib2/weights/conf/%s/stdc_360X640.pth'%(opt['business'] ),
  269. 'par':{
  270. '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,'seg_nclass':4},###分割模型预处理参数
  271. 'model':stdcModel,
  272. 'name':'stdc'
  273. }
  274. ],
  275. 'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,5,6,7,8,9] ],###控制哪些检测类别显示、输出
  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['postProcess']['pars']['modelSize'] = par['models'][1]['par']['modelSize']
  283. if opt['business'] == 'drowning':
  284. from utilsK.drownUtils import mixDrowing_water_postprocess_N
  285. par={
  286. 'device':'0', ###显卡号,如果用TRT模型,只支持0(单显卡)
  287. 'labelnames':"../AIlib2/weights/conf/%s/labelnames.json"%( opt['business'] ), ###检测类别对照表
  288. 'max_workers':1, ###并行线程数
  289. 'trackPar':{'sort_max_age':2,'sort_min_hits':3,'sort_iou_thresh':0.2,'det_cnt':10,'windowsize':29,'patchCnt':100},
  290. 'postProcess':{'function':mixDrowing_water_postprocess_N,
  291. 'pars':{ }},
  292. 'models':
  293. [
  294. {
  295. 'weight':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###检测模型路径
  296. #'weight':'../AIlib2/weights/conf/%s/yolov5.pt'%(opt['business'] ),
  297. 'name':'yolov5',
  298. 'model':yolov5Model,
  299. 'par':{ 'half':True,'device':'cuda:0' ,'conf_thres':0.25,'iou_thres':0.45,'allowedList':[0,1,2,3],'segRegionCnt':1, 'trtFlag_det':False,'trtFlag_seg':False, "score_byClass":{"0":0.25,"1":0.3,"2":0.3,"3":0.3 } },
  300. },
  301. {
  302. #'weight':"../weights/%s/AIlib2/%s/dmpr_%s.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###检测模型路径
  303. 'weight':'../AIlib2/weights/conf/%s/stdc_360X640.pth'%(opt['business'] ),
  304. 'par':{
  305. '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,'seg_nclass':2},###分割模型预处理参数
  306. 'model':stdcModel,
  307. 'name':'stdc'
  308. }
  309. ],
  310. 'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,5,6,7,8,9] ],###控制哪些检测类别显示、输出
  311. 'postFile': '../AIlib2/weights/conf/%s/para.json'%('highWay2' ),###后处理参数文件
  312. 'txtFontSize':20,###文本字符的大小
  313. 'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'segLineShow':True,'waterLineWidth':2},###显示框、线设置
  314. 'testImgPath':'images/drowning/',###测试图像的位置
  315. 'testOutPath':'images/results/',###输出测试图像位置
  316. }
  317. par['postProcess']['pars']['modelSize'] = par['models'][1]['par']['modelSize']
  318. if opt['business'] == 'cityMangement2':
  319. from DMPR import DMPRModel
  320. from DMPRUtils.jointUtil import dmpr_yolo_stdc
  321. par={
  322. 'device':'0', ###显卡号,如果用TRT模型,只支持0(单显卡)
  323. 'labelnames':"../AIlib2/weights/conf/%s/labelnames.json"%(opt['business']), ###检测类别对照表
  324. 'max_workers':1, ###并行线程数
  325. 'trackPar':{'sort_max_age':2,'sort_min_hits':3,'sort_iou_thresh':0.2,'det_cnt':10,'windowsize':29,'patchCnt':100},
  326. 'postProcess':{
  327. 'function':dmpr_yolo_stdc,
  328. 'pars':{'carCls':0 ,'illCls':3,'scaleRatio':0.5,'border':80}
  329. },
  330. 'models':
  331. [
  332. {
  333. #'weight':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###检测模型路径
  334. 'weight':'../AIlib2/weights/conf/%s/yolov5.pt'%(opt['business'] ),
  335. 'name':'yolov5',
  336. 'model':yolov5Model,
  337. 'par':{ 'half':True,'device':'cuda:0' ,'conf_thres':0.25,'iou_thres':0.45,'allowedList':[0,1,2,3],'segRegionCnt':1, 'trtFlag_det':False,'trtFlag_seg':False, "score_byClass":{"0":0.25,"1":0.3,"2":0.3,"3":0.3 } }
  338. },
  339. {
  340. #'weight':"../weights/%s/AIlib2/%s/dmpr_%s.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###检测模型路径
  341. 'weight':'../AIlib2/weights/conf/%s/dmpr.pth'%(opt['business'] ),
  342. 'par':{
  343. 'depth_factor':32,'NUM_FEATURE_MAP_CHANNEL':6,'dmpr_thresh':0.3, 'dmprimg_size':640,
  344. 'name':'dmpr'
  345. },
  346. 'model':DMPRModel,
  347. 'name':'dmpr'
  348. }
  349. ] ,
  350. 'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,5,6,7,8,9] ],###控制哪些检测类别显示、输出
  351. 'postFile': '../AIlib2/weights/conf/%s/para.json'%(opt['business'] ),###后处理参数文件
  352. 'txtFontSize':20,###文本字符的大小
  353. 'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'segLineShow':True,'waterLineWidth':2},###显示框、线设置
  354. #'testImgPath':'/mnt/thsw2/DSP2/demoImages/illParking',###测试图像的位置
  355. 'testImgPath':'/mnt/thsw2/DSP2/weights/cityMangement2_0916/images/input',
  356. #'testImgPath':'images/cityMangement/',
  357. 'testOutPath':'images/results/',###输出测试图像位置
  358. }
  359. if opt['business'] == 'cityMangement3':
  360. from DMPR import DMPRModel
  361. from DMPRUtils.jointUtil import dmpr_yolo_stdc
  362. par={
  363. 'device':'0', ###显卡号,如果用TRT模型,只支持0(单显卡)
  364. 'labelnames':"../AIlib2/weights/conf/%s/labelnames.json"%(opt['business']), ###检测类别对照表
  365. 'max_workers':1, ###并行线程数
  366. 'trackPar':{'sort_max_age':2,'sort_min_hits':3,'sort_iou_thresh':0.2,'det_cnt':10,'windowsize':29,'patchCnt':100},
  367. 'postProcess':{
  368. 'function':dmpr_yolo_stdc,
  369. 'pars':{'carCls':0 ,'illCls':3,'scaleRatio':0.5,'border':80}
  370. },
  371. 'models':[
  372. {
  373. 'weight':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###检测模型路径
  374. #'weight':'../AIlib2/weights/conf/%s/yolov5.pt'%(opt['business'] ),
  375. 'name':'yolov5',
  376. 'model':yolov5Model,
  377. 'par':{ 'half':True,'device':'cuda:0' ,'conf_thres':0.25,'iou_thres':0.45,'allowedList':[0,1,2,3],'segRegionCnt':1, 'trtFlag_det':False,'trtFlag_seg':False, "score_byClass":{"0":0.25,"1":0.3,"2":0.3,"3":0.3 } }
  378. },
  379. {
  380. 'weight':"../weights/%s/AIlib2/%s/dmpr_%s.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###DMPR模型路径
  381. #'weight':'../AIlib2/weights/conf/%s/dmpr.pth'%(opt['business'] ),
  382. 'par':{
  383. 'depth_factor':32,'NUM_FEATURE_MAP_CHANNEL':6,'dmpr_thresh':0.1, 'dmprimg_size':640,
  384. 'name':'dmpr'
  385. },
  386. 'model':DMPRModel,
  387. 'name':'dmpr'
  388. },
  389. {
  390. 'weight':"../weights/%s/AIlib2/%s/stdc_360X640_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###分割模型路径
  391. #'weight':'../AIlib2/weights/conf/%s/stdc_360X640.pth'%(opt['business'] ),
  392. 'par':{
  393. '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,'seg_nclass':2},###分割模型预处理参数
  394. 'model':stdcModel,
  395. 'name':'stdc'
  396. }
  397. ],
  398. 'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,5,6,7,8,9] ],###控制哪些检测类别显示、输出
  399. 'postFile': '../AIlib2/weights/conf/%s/para.json'%(opt['business'] ),###后处理参数文件
  400. 'txtFontSize':20,###文本字符的大小
  401. 'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'segLineShow':True,'waterLineWidth':2},###显示框、线设置
  402. 'testImgPath':'images/%s/'%(opt['business']),
  403. #'testImgPath':'/mnt/thsw2/DSP2/weights/cityMangement3/images/',
  404. 'testOutPath':'images/results/',###输出测试图像位置
  405. }
  406. if opt['business'] == 'forest2':
  407. par={
  408. 'device':'0', ###显卡号,如果用TRT模型,只支持0(单显卡)
  409. 'labelnames':"../AIlib2/weights/conf/forest2/labelnames.json", ###检测类别对照表
  410. 'max_workers':1, ###并行线程数
  411. 'trackPar':{'sort_max_age':2,'sort_min_hits':3,'sort_iou_thresh':0.2,'det_cnt':10,'windowsize':29,'patchCnt':100},
  412. 'postProcess':{'function':default_mix,'pars':{ }},
  413. 'models':
  414. [
  415. {
  416. 'weight':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###检测模型路径
  417. #'weight':'../AIlib2/weights/conf/%s/yolov5.pt'%(opt['business'] ),
  418. 'name':'yolov5',
  419. 'model':yolov5Model,
  420. 'par':{ 'half':True,'device':'cuda:0' ,'conf_thres':0.25,'iou_thres':0.45,'allowedList':[0,1,2,3],'segRegionCnt':1, 'trtFlag_det':False,'trtFlag_seg':False, "score_byClass":{"0":0.25,"1":0.3,"2":0.3,"3":0.3 } },
  421. }
  422. ],
  423. 'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [] ],
  424. 'postFile': '../AIlib2/weights/conf/forest/para.json',###后处理参数文件
  425. 'txtFontSize':80,###文本字符的大小
  426. 'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'waterLineWidth':3},###显示框、线设置
  427. 'testImgPath':'../AIdemo2/images/forest2/',###测试图像的位置
  428. 'testOutPath':'images/results/',###输出测试图像位置
  429. }
  430. ###车辆巡检参数
  431. if opt['business'] == 'vehicle':
  432. par={
  433. 'device':'0', ###显卡号,如果用TRT模型,只支持0(单显卡)
  434. 'labelnames':"../AIlib2/weights/conf/vehicle/labelnames.json", ###检测类别对照表
  435. 'max_workers':1, ###并行线程数
  436. 'trackPar':{'sort_max_age':2,'sort_min_hits':3,'sort_iou_thresh':0.2,'det_cnt':10,'windowsize':29,'patchCnt':100},
  437. 'postProcess':{'function':default_mix,'pars':{ }},
  438. 'models':
  439. [
  440. {
  441. 'weight':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###检测模型路径
  442. #'weight':'../AIlib2/weights/conf/%s/yolov5.pt'%(opt['business'] ),
  443. 'name':'yolov5',
  444. 'model':yolov5Model,
  445. 'par':{ 'half':True,'device':'cuda:0' ,'conf_thres':0.25,'iou_thres':0.45,'allowedList':[0,1,2,3],'segRegionCnt':1, 'trtFlag_det':False,'trtFlag_seg':False, "score_byClass":{"0":0.25,"1":0.3,"2":0.3,"3":0.3 } },
  446. }
  447. ],
  448. 'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,4,5,6] ],###控制哪些检测类别显示、输出
  449. 'postFile': '../AIlib2/weights/conf/vehicle/para.json',###后处理参数文件
  450. 'txtFontSize':40,###文本字符的大小
  451. 'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'waterLineWidth':3},###显示框、线设置
  452. 'testImgPath':'images/vehicle/',###测试图像的位置
  453. 'testOutPath':'images/results/',###输出测试图像位置
  454. }
  455. ###行人检测模型
  456. if opt['business'] == 'pedestrian':
  457. par={
  458. 'device':'0', ###显卡号,如果用TRT模型,只支持0(单显卡)
  459. 'labelnames':"../AIlib2/weights/conf/pedestrian/labelnames.json", ###检测类别对照表
  460. 'max_workers':1, ###并行线程数
  461. 'trackPar':{'sort_max_age':2,'sort_min_hits':3,'sort_iou_thresh':0.2,'det_cnt':10,'windowsize':29,'patchCnt':100},
  462. 'postProcess':{'function':default_mix,'pars':{ }},
  463. 'models':
  464. [
  465. {
  466. 'weight':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###检测模型路径
  467. #'weight':'../AIlib2/weights/conf/%s/yolov5.pt'%(opt['business'] ),
  468. 'name':'yolov5',
  469. 'model':yolov5Model,
  470. 'par':{ 'half':True,'device':'cuda:0' ,'conf_thres':0.25,'iou_thres':0.45,'allowedList':[0,1,2,3],'segRegionCnt':1, 'trtFlag_det':False,'trtFlag_seg':False, "score_byClass":{"0":0.25,"1":0.3,"2":0.3,"3":0.3 } },
  471. }
  472. ],
  473. 'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,4,5,6] ],###控制哪些检测类别显示、输出
  474. 'postFile': '../AIlib2/weights/conf/pedestrian/para.json',###后处理参数文件
  475. 'txtFontSize':40,###文本字符的大小
  476. 'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'waterLineWidth':3},###显示框、线设置
  477. 'testImgPath':'../AIdemo2/images/pedestrian/',###测试图像的位置
  478. 'testOutPath':'images/results/',###输出测试图像位置
  479. }
  480. if opt['business'] == 'smogfire':
  481. par={
  482. 'device':'0', ###显卡号,如果用TRT模型,只支持0(单显卡)
  483. 'labelnames':"../AIlib2/weights/conf/smogfire/labelnames.json", ###检测类别对照表
  484. 'max_workers':1, ###并行线程数
  485. 'trackPar':{'sort_max_age':2,'sort_min_hits':3,'sort_iou_thresh':0.2,'det_cnt':10,'windowsize':29,'patchCnt':100},
  486. 'postProcess':{'function':default_mix,'pars':{ }},
  487. 'models':
  488. [
  489. {
  490. 'weight':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###检测模型路径
  491. #'weight':'../AIlib2/weights/conf/%s/yolov5.pt'%(opt['business'] ),
  492. 'name':'yolov5',
  493. 'model':yolov5Model,
  494. 'par':{ 'half':True,'device':'cuda:0' ,'conf_thres':0.25,'iou_thres':0.45,'allowedList':[0,1,2,3],'segRegionCnt':1, 'trtFlag_det':False,'trtFlag_seg':False, "score_byClass":{"0":0.25,"1":0.3,"2":0.3,"3":0.3 } },
  495. }
  496. ],
  497. 'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,4,5,6] ],###控制哪些检测类别显示、输出
  498. 'postFile': '../AIlib2/weights/conf/smogfire/para.json',###后处理参数文件
  499. 'txtFontSize':40,###文本字符的大小
  500. 'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'waterLineWidth':3},###显示框、线设置
  501. 'testImgPath':'../AIdemo2/images/smogfire/',###测试图像的位置
  502. 'testOutPath':'images/results/',###输出测试图像位置
  503. }
  504. ###钓鱼游泳检测
  505. if opt['business'] == 'AnglerSwimmer':
  506. par={
  507. 'device':'0', ###显卡号,如果用TRT模型,只支持0(单显卡)
  508. 'labelnames':"../AIlib2/weights/conf/AnglerSwimmer/labelnames.json", ###检测类别对照表
  509. 'max_workers':1, ###并行线程数
  510. 'trackPar':{'sort_max_age':2,'sort_min_hits':3,'sort_iou_thresh':0.2,'det_cnt':10,'windowsize':29,'patchCnt':100},
  511. 'postProcess':{'function':default_mix,'pars':{ }},
  512. 'models':
  513. [
  514. {
  515. 'weight':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###检测模型路径
  516. #'weight':'../AIlib2/weights/conf/%s/yolov5.pt'%(opt['business'] ),
  517. 'name':'yolov5',
  518. 'model':yolov5Model,
  519. 'par':{ 'half':True,'device':'cuda:0' ,'conf_thres':0.25,'iou_thres':0.45,'allowedList':[0,1,2,3],'segRegionCnt':1, 'trtFlag_det':False,'trtFlag_seg':False, "score_byClass":{"0":0.25,"1":0.3,"2":0.3,"3":0.3 } },
  520. }
  521. ],
  522. 'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,4,5,6] ],###控制哪些检测类别显示、输出
  523. 'postFile': '../AIlib2/weights/conf/AnglerSwimmer/para.json',###后处理参数文件
  524. 'txtFontSize':40,###文本字符的大小
  525. 'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'waterLineWidth':3},###显示框、线设置
  526. 'testImgPath':'../AIdemo2/images/AnglerSwimmer/',###测试图像的位置
  527. 'testOutPath':'images/results/',###输出测试图像位置
  528. }
  529. #单独的“坑槽”检测
  530. if opt['business'] == 'pothole':
  531. par={
  532. 'device':'0', ###显卡号,如果用TRT模型,只支持0(单显卡)
  533. 'labelnames':"../AIlib2/weights/conf/pothole/labelnames.json", ###检测类别对照表
  534. 'max_workers':1, ###并行线程数
  535. 'trackPar':{'sort_max_age':2,'sort_min_hits':3,'sort_iou_thresh':0.2,'det_cnt':10,'windowsize':29,'patchCnt':100},
  536. 'postProcess':{'function':default_mix,'pars':{ }},
  537. 'models':
  538. [
  539. {
  540. 'weight':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###检测模型路径
  541. #'weight':'../AIlib2/weights/conf/%s/yolov5.pt'%(opt['business'] ),
  542. 'name':'yolov5',
  543. 'model':yolov5Model,
  544. 'par':{ 'half':True,'device':'cuda:0' ,'conf_thres':0.25,'iou_thres':0.45,'allowedList':[0,1,2,3],'segRegionCnt':1, 'trtFlag_det':False,'trtFlag_seg':False, "score_byClass":{"0":0.25,"1":0.3,"2":0.3,"3":0.3 } },
  545. }
  546. ],
  547. 'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [] ],
  548. 'postFile': '../AIlib2/weights/conf/forest/para.json',###后处理参数文件
  549. 'txtFontSize':80,###文本字符的大小
  550. 'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'waterLineWidth':3},###显示框、线设置
  551. 'testImgPath':'../AIdemo2/images/pothole/',###测试图像的位置
  552. 'testOutPath':'images/results/',###输出测试图像位置
  553. }
  554. ###航道应急,做落水人员检测, channelEmergency
  555. if opt['business'] == 'channelEmergency':
  556. par={
  557. 'device':'0', ###显卡号,如果用TRT模型,只支持0(单显卡)
  558. 'labelnames':"../AIlib2/weights/conf/channelEmergency/labelnames.json", ###检测类别对照表
  559. 'max_workers':1, ###并行线程数
  560. 'trackPar':{'sort_max_age':2,'sort_min_hits':3,'sort_iou_thresh':0.2,'det_cnt':10,'windowsize':29,'patchCnt':100},
  561. 'postProcess':{'function':default_mix,'pars':{ }},
  562. 'models':
  563. [
  564. {
  565. 'weight':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###检测模型路径
  566. #'weight':'../AIlib2/weights/conf/%s/yolov5.pt'%(opt['business'] ),
  567. 'name':'yolov5',
  568. 'model':yolov5Model,
  569. 'par':{ 'half':True,'device':'cuda:0' ,'conf_thres':0.25,'iou_thres':0.45,'allowedList':[0,1,2,3],'segRegionCnt':1, 'trtFlag_det':False,'trtFlag_seg':False, "score_byClass":{"0":0.25,"1":0.3,"2":0.3,"3":0.3 } },
  570. }
  571. ],
  572. 'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [] ],###控制哪些检测类别显示、输出
  573. 'postFile': '../AIlib2/weights/conf/channelEmergency/para.json',###后处理参数文件
  574. 'txtFontSize':40,###文本字符的大小
  575. 'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'waterLineWidth':3},###显示框、线设置
  576. 'testImgPath':'../AIdemo2/images/channelEmergency/',###测试图像的位置
  577. 'testOutPath':'images/results/',###输出测试图像位置
  578. }
  579. ###乡村路违法种植
  580. if opt['business'] == 'countryRoad':
  581. par={
  582. 'device':'0', ###显卡号,如果用TRT模型,只支持0(单显卡)
  583. 'labelnames':"../AIlib2/weights/conf/countryRoad/labelnames.json", ###检测类别对照表
  584. 'max_workers':1, ###并行线程数
  585. 'trackPar':{'sort_max_age':2,'sort_min_hits':3,'sort_iou_thresh':0.2,'det_cnt':10,'windowsize':29,'patchCnt':100},
  586. 'postProcess':{'function':default_mix,'pars':{ }},
  587. 'models':
  588. [
  589. {
  590. 'weight':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###检测模型路径
  591. #'weight':'../AIlib2/weights/conf/%s/yolov5.pt'%(opt['business'] ),
  592. 'name':'yolov5',
  593. 'model':yolov5Model,
  594. 'par':{ 'half':True,'device':'cuda:0' ,'conf_thres':0.25,'iou_thres':0.45,'allowedList':[0,1,2,3],'segRegionCnt':1, 'trtFlag_det':False,'trtFlag_seg':False, "score_byClass":{"0":0.25,"1":0.3,"2":0.3,"3":0.3 } },
  595. }
  596. ],
  597. 'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,4,5,6] ],###控制哪些检测类别显示、输出
  598. 'postFile': '../AIlib2/weights/conf/countryRoad/para.json',###后处理参数文件
  599. 'txtFontSize':40,###文本字符的大小
  600. 'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'waterLineWidth':3},###显示框、线设置
  601. 'testImgPath':'../AIdemo2/images/countryRoad/',###测试图像的位置
  602. 'testOutPath':'images/results/',###输出测试图像位置
  603. }
  604. ###城管项目,检测城市垃圾和车辆
  605. if opt['business'] == 'cityMangement':
  606. par={
  607. 'device':'0', ###显卡号,如果用TRT模型,只支持0(单显卡)
  608. 'labelnames':"../AIlib2/weights/conf/%s/labelnames.json"%(opt['business']), ###检测类别对照表
  609. 'max_workers':1, ###并行线程数
  610. 'trackPar':{'sort_max_age':2,'sort_min_hits':3,'sort_iou_thresh':0.2,'det_cnt':10,'windowsize':29,'patchCnt':100},
  611. 'postProcess':{'function':default_mix,'pars':{ }},
  612. 'models':
  613. [
  614. {
  615. 'weight':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###检测模型路径
  616. #'weight':'../AIlib2/weights/conf/%s/yolov5.pt'%(opt['business'] ),
  617. 'name':'yolov5',
  618. 'model':yolov5Model,
  619. 'par':{ 'half':True,'device':'cuda:0' ,'conf_thres':0.25,'iou_thres':0.45,'allowedList':[0,1,2,3],'segRegionCnt':1, 'trtFlag_det':False,'trtFlag_seg':False, "score_byClass":{"0":0.25,"1":0.3,"2":0.3,"3":0.3 } },
  620. }
  621. ],
  622. 'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,4,5,6] ],###控制哪些检测类别显示、输出
  623. 'postFile': '../AIlib2/weights/conf/%s/para.json'%(opt['business']),###后处理参数文件
  624. 'txtFontSize':40,###文本字符的大小
  625. 'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'waterLineWidth':3},###显示框、线设置
  626. 'testImgPath':'images/cityMangement',###测试图像的位置
  627. 'testOutPath':'images/results/',###输出测试图像位置
  628. }
  629. ###城管项目,检测道路情况,输入类别为五个:"护栏","交通标志","非交通标志","施工","施工“(第4,第5类别合并,名称相同)
  630. ###实际模型检测输出的类别为:"护栏","交通标志","非交通标志","锥桶","水马"
  631. if opt['business'] == 'cityRoad':
  632. par={
  633. 'device':'0', ###显卡号,如果用TRT模型,只支持0(单显卡)
  634. 'labelnames':"../AIlib2/weights/conf/%s/labelnames.json"%(opt['business']), ###检测类别对照表
  635. 'max_workers':1, ###并行线程数
  636. 'trackPar':{'sort_max_age':2,'sort_min_hits':3,'sort_iou_thresh':0.2,'det_cnt':10,'windowsize':29,'patchCnt':100},
  637. 'postProcess':{'function':default_mix,'pars':{ }},
  638. 'models':
  639. [
  640. {
  641. 'weight':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###检测模型路径
  642. #'weight':'../AIlib2/weights/conf/%s/yolov5.pt'%(opt['business'] ),
  643. 'name':'yolov5',
  644. 'model':yolov5Model,
  645. 'par':{ 'half':True,'device':'cuda:0' ,'conf_thres':0.25,'iou_thres':0.45,'allowedList':[0,1,2,3],'segRegionCnt':1, 'trtFlag_det':False,'trtFlag_seg':False, "score_byClass":{"0":0.25,"1":0.3,"2":0.3,"3":0.3 } },
  646. }
  647. ],
  648. 'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,4,5,6] ],###控制哪些检测类别显示、输出
  649. 'postFile': '../AIlib2/weights/conf/%s/para.json'%(opt['business']),###后处理参数文件
  650. 'txtFontSize':40,###文本字符的大小
  651. 'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'waterLineWidth':3},###显示框、线设置
  652. 'testImgPath':'images/%s'%(opt['business']),###测试图像的位置
  653. 'testOutPath':'images/results/',###输出测试图像位置
  654. }
  655. if opt['business'] == 'illParking':
  656. from utilsK.illParkingUtils import illParking_postprocess_N
  657. par={
  658. 'device':'0', ###显卡号,如果用TRT模型,只支持0(单显卡)
  659. 'labelnames':"../AIlib2/weights/conf/%s/labelnames.json"%(opt['business']), ###检测类别对照表
  660. 'trackPar':{'sort_max_age':2,'sort_min_hits':3,'sort_iou_thresh':0.2,'det_cnt':10,'windowsize':29,'patchCnt':100},
  661. 'postProcess':{'function':illParking_postprocess_N, 'pars':{ }},
  662. 'models':
  663. [
  664. {
  665. 'weight':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###检测模型路径
  666. #'weight':'../AIlib2/weights/conf/%s/yolov5.pt'%(opt['business'] ),
  667. 'name':'yolov5',
  668. 'model':yolov5Model,
  669. 'par':{ 'half':True,'device':'cuda:0' ,'conf_thres':0.25,'iou_thres':0.45,'allowedList':[0,1,2,3],'segRegionCnt':1, 'trtFlag_det':False,'trtFlag_seg':False, "score_byClass":{"0":0.25,"1":0.3,"2":0.3,"3":0.3 } },
  670. }
  671. ],
  672. 'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,5,6,7,8,9] ],###控制哪些检测类别显示、输出
  673. 'postFile': '../AIlib2/weights/conf/%s/para.json'%(opt['business'] ),###后处理参数文件
  674. 'txtFontSize':20,###文本字符的大小
  675. 'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'waterLineWidth':2},###显示框、线设置
  676. 'testImgPath':'images/cityMangement',###测试图像的位置
  677. 'testOutPath':'images/results/',###输出测试图像位置
  678. }
  679. #第一步加载模型
  680. modelList=[ modelPar['model'](weights=modelPar['weight'],par=modelPar['par']) for modelPar in par['models'] ]
  681. #第二步准备跟踪参数
  682. trackPar=par['trackPar']
  683. sort_tracker = Sort(max_age=trackPar['sort_max_age'],
  684. min_hits=trackPar['sort_min_hits'],
  685. iou_threshold=trackPar['sort_iou_thresh'])
  686. #第三步准备画图字体
  687. drawPar = get_drawPar(par)
  688. #第四步获取图像测试及视频
  689. imgpaths,videopaths = get_images_videos( par['testImgPath'])
  690. #第五步开始测试
  691. for i in range(len(imgpaths)):
  692. imgpath = imgpaths[i]
  693. bname = os.path.basename(imgpath )
  694. im0s=[cv2.imread(imgpath)]
  695. time00 = time.time()
  696. retResults,timeOut = AI_det_track_batch_N(im0s, [i] ,modelList,par['postProcess'],sort_tracker ,trackPar)
  697. '''
  698. 返回值retResults[2] --list,其中每一个元素为一个list,表示每一帧的检测结果,每一个结果是由多个list构成,每个list表示一个框,格式为[ x0 ,y0 ,x1 ,y1 ,conf,cls ,ifrmae,trackId ]
  699. --etc. retResults[2][j][k]表示第j帧的第k个框。
  700. '''
  701. time11 = time.time()
  702. if len(retResults[1])>0:
  703. retResults[0][0] = drawBoxTraceSimplied(retResults[1],i, retResults[0][0],rainbows=rainbows,boxFlag=True,traceFlag=False,names=drawPar['names'])
  704. image_array = retResults[0][0]
  705. cv2.imwrite( os.path.join( par['testOutPath'],bname ) ,image_array )
  706. print('----image:%s, Allprocess:%s %s ,save:%s , objcnt:%d'%(bname,(time11-time00) * 1000, timeOut,(time.time() - time11) * 1000 ,len(retResults[2])) )
  707. ##第五步开始测试视频
  708. print('##begin to process videos, total %d videos'%( len(videopaths)))
  709. for i,video in enumerate(videopaths):
  710. print('process video%d :%s '%(i,video))
  711. par0={'modelList':modelList,'postProcess':par['postProcess'],'drawPar':drawPar,'outpth':par['testOutPath'], 'sort_tracker':sort_tracker,'trackPar':trackPar}
  712. process_video(video,par0,mode='track')
  713. def OCR_demo2(opt):
  714. from ocrUtils2 import crnn_model
  715. from ocrUtils2.ocrUtils import get_cfg,recognition_ocr,strLabelConverter
  716. if opt['business'] == 'ocr2':
  717. par={
  718. 'image_dir':'images/ocr_en',
  719. 'outtxt':'images/results',
  720. 'weights':'../AIlib2/weights/conf/ocr2/crnn_448X32.pth',
  721. #'weights':'../weights/2080Ti/AIlib2/ocr2/crnn_2080Ti_fp16_448X32.engine',
  722. 'device':'cuda:0',
  723. 'cfg':'../AIlib2/weights/conf/ocr2/360CC_config.yaml',
  724. 'char_file':'../AIlib2/weights/conf/ocr2/chars.txt',
  725. 'imgH':32,
  726. 'imgW':448,
  727. 'workers':1
  728. }
  729. image_dir=par['image_dir']
  730. outtxt=par['outtxt']
  731. workers=par['workers']
  732. weights= par['weights']
  733. device=par['device']
  734. char_file=par['char_file']
  735. imgH=par['imgH']
  736. imgW=par['imgW']
  737. cfg = par['cfg']
  738. config = get_cfg(cfg, char_file)
  739. par['contextFlag']=False
  740. device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')
  741. if weights.endswith('.pth'):
  742. model = crnn_model.get_crnn(config,weights=weights).to(device)
  743. par['model_mode']='pth'
  744. else:
  745. logger = trt.Logger(trt.Logger.ERROR)
  746. with open(weights, "rb") as f, trt.Runtime(logger) as runtime:
  747. model = runtime.deserialize_cuda_engine(f.read())# 输入trt本地文件,返回ICudaEngine对象
  748. print('#####load TRT file:',weights,'success #####')
  749. context = model.create_execution_context()
  750. par['model_mode']='trt';par['contextFlag']=context
  751. converter = strLabelConverter(config.DATASET.ALPHABETS)
  752. img_urls=glob.glob('%s/*.jpg'%( image_dir ))
  753. img_urls.extend( glob.glob('%s/*.png'%( image_dir )) )
  754. cnt=len(img_urls)
  755. print('%s has %d images'%(image_dir ,len(img_urls) ) )
  756. # 准备数据
  757. parList=[]
  758. for i in range(cnt):
  759. img_patch=cv2.imread( img_urls[i] , cv2.IMREAD_GRAYSCALE)
  760. started = time.time()
  761. img = cv2.imread(img_urls[i])
  762. sim_pred = recognition_ocr(config, img, model, converter, device,par=par)
  763. finished = time.time()
  764. print('{0}: elapsed time: {1} prd:{2} '.format( os.path.basename( img_urls[i] ), finished - started, sim_pred ))
  765. def customization_demo(opt):
  766. from AI import AI_process_C
  767. if opt['business'] == 'channel2':
  768. from ocr import ocrModel
  769. par={
  770. 'device':'0', ###显卡号,如果用TRT模型,只支持0(单显卡)
  771. 'labelnames':"../AIlib2/weights/conf/%s/labelnames.json"%(opt['business']), ###检测类别对照表
  772. 'max_workers':1, ###并行线程数
  773. 'postProcess':{'function':default_mix,'pars':{'objs':[2] },'name':'channel2',},
  774. 'models':[
  775. {
  776. #'weight':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###检测模型路径
  777. 'weight':'../AIlib2/weights/conf/%s/yolov5.pt'%(opt['business'] ),
  778. #'weight':'/mnt/thsw2/DSP2/weights/channel2/yolov5.pt',
  779. #'weight':'/mnt/thsw2/DSP2/weights/channel2/yolov5_2080Ti_fp16.engine',
  780. 'name':'yolov5',
  781. 'model':yolov5Model,
  782. 'par':{ 'half':True,'device':'cuda:0' ,'conf_thres':0.1,'iou_thres':0.45,'allowedList':list(range(20)),'segRegionCnt':1, 'trtFlag_det':False,'trtFlag_seg':False, "score_byClass":{str(x):0.1 for x in range(20) } }
  783. },
  784. {
  785. 'weight' : '../weights/%s/AIlib2/ocr2/crnn_ch_%s_fp16_192X32.engine'%(opt['gpu'], opt['gpu']),
  786. 'name':'ocr',
  787. 'model':ocrModel,
  788. 'par':{
  789. 'char_file':'../AIlib2/weights/conf/ocr2/chars.txt',
  790. 'mode':'ch',
  791. 'nc':3,
  792. 'imgH':32,
  793. 'imgW':192,
  794. 'hidden':256,
  795. 'mean':[0.5,0.5,0.5],
  796. 'std':[0.5,0.5,0.5],
  797. 'dynamic':False,
  798. },
  799. }
  800. ],
  801. 'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,5,6,7,8,9] ],###控制哪些检测类别显示、输出
  802. 'postFile': '../AIlib2/weights/conf/%s/para.json'%(opt['business'] ),###后处理参数文件
  803. 'txtFontSize':20,###文本字符的大小
  804. 'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'segLineShow':True,'waterLineWidth':2},###显示框、线设置
  805. 'testImgPath':'images/%s/'%(opt['business']),
  806. #'testImgPath':'/home/thsw2/WJ/src/OCR/shipNames',
  807. 'testOutPath':'images/results/',###输出测试图像位置
  808. }
  809. if opt['business'] == 'crackMeasurement':
  810. from utilsK.crackUtils import Crack_measure
  811. print( '%s 只能测试图像,不能测试视频%s'%('#'*20,'#'*20))
  812. par={
  813. 'device':'0', ###显卡号,如果用TRT模型,只支持0(单显卡)
  814. 'labelnames':"../AIlib2/weights/conf/%s/labelnames.json"%(opt['business']), ###检测类别对照表
  815. 'max_workers':1, ###并行线程数
  816. 'postProcess':{
  817. 'name':'crackMeasurement',
  818. 'function':Crack_measure,
  819. 'pars':{'dsx':(123-30)*1000/35*0.004387636 ,'objs':[0,1,2]}
  820. },
  821. 'models':[
  822. {
  823. #'weight':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###检测模型路径
  824. 'weight':'../AIlib2/weights/conf/%s/yolov5.pt'%(opt['business'] ),
  825. 'name':'yolov5',
  826. 'model':yolov5Model,
  827. 'par':{ 'half':True,'device':'cuda:0' ,'conf_thres':0.1,'iou_thres':0.45,'allowedList':[0,1,2,3],'segRegionCnt':1, 'trtFlag_det':False,'trtFlag_seg':False, "score_byClass":{"0":0.1,"1":0.1,"2":0.1 } }
  828. },
  829. {
  830. #'weight':"../weights/%s/AIlib2/%s/stdc_360X640_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###检测模型路径
  831. 'weight':'../AIlib2/weights/conf/%s/stdc_360X640.pth'%(opt['business'] ),
  832. 'par':{
  833. #'modelSize':(640,360),
  834. 'modelSize':(1920,1080),
  835. 'dynamic':True,
  836. 'mean':(0.485, 0.456, 0.406),'std' :(0.229, 0.224, 0.225),'predResize':True,'numpy':False, 'RGB_convert_first':True,'seg_nclass':2},###分割模型预处理参数
  837. 'model':stdcModel,
  838. 'name':'stdc'
  839. }
  840. ],
  841. 'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,5,6,7,8,9] ],###控制哪些检测类别显示、输出
  842. 'postFile': '../AIlib2/weights/conf/%s/para.json'%(opt['business'] ),###后处理参数文件
  843. 'txtFontSize':20,###文本字符的大小
  844. 'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'segLineShow':True,'waterLineWidth':2},###显示框、线设置
  845. 'testImgPath':'images/%s/'%(opt['business']),
  846. #'testImgPath':'/mnt/thsw2/DSP2/weights/cityMangement2_1102/images/debug',
  847. 'testOutPath':'images/results/',###输出测试图像位置
  848. }
  849. #第一步加载模型
  850. modelList=[ modelPar['model'](weights=modelPar['weight'],par=modelPar['par']) for modelPar in par['models'] ]
  851. print(' load moder over')
  852. #准备画图字体
  853. labelnames = par['labelnames'] ##对应类别表
  854. names=get_labelnames(labelnames)
  855. label_arraylist = get_label_arrays(names,rainbows,outfontsize=par['txtFontSize'],fontpath="../AIlib2/conf/platech.ttf")
  856. #图像测试
  857. imgpaths,videopaths = get_images_videos( par['testImgPath'])
  858. #开始测试
  859. for imgUrl in imgpaths[0:]:
  860. img = cv2.imread(imgUrl);bname = os.path.basename(imgUrl)
  861. if opt['business'] == 'crackMeasurement':
  862. ret,timeInfos = AI_process_C([img],modelList,par['postProcess'])
  863. #返回类型ret-list,[[ x0,y0,x1,y1,score,class,裂缝长度,平均宽度,最大宽度,最小宽度],[...],[...]]
  864. for re in ret:
  865. print('Summarized Cracklength = %.1f mm Mean crack width = %.1f mm Max crack width = %.1f mm Min crack width = %.1f mm '%( re[6], re[7],re[8],re[9] ) )
  866. elif opt['business'] == 'channel2':
  867. ret,timeInfos = AI_process_C([img],modelList,par['postProcess'])
  868. timeInfos=bname+':'+timeInfos
  869. print(timeInfos,ret )
  870. if len(ret)>0:
  871. img0 = drawAllBox(ret,img,label_arraylist,rainbows,par['digitFont'])
  872. else: img0= img
  873. cv2.imwrite(os.path.join('images/results/',bname ) ,img0)
  874. def OBB_track_demo(opt):
  875. ###倾斜框(OBB)的ship目标检测
  876. par={
  877. 'obbModelPar':{
  878. 'model_size':(608,608),'K':100,'conf_thresh':0.3, 'down_ratio':4,'num_classes':15,'dataset':'dota',
  879. 'heads': {'hm': None,'wh': 10,'reg': 2,'cls_theta': 1},
  880. 'mean':(0.5, 0.5, 0.5),'std':(1, 1, 1), 'half': False,'decoder':None,
  881. 'weights':'../weights/%s/AIlib2/%s/obb_608X608_%s_fp16.engine'%(opt['gpu'],opt['business'],opt['gpu']),
  882. },
  883. 'outpth': 'images/results',
  884. 'trackPar':{'sort_max_age':2,'sort_min_hits':3,'sort_iou_thresh':0.2,'det_cnt':10,'windowsize':29,'patchCnt':100},
  885. 'device':"cuda:0",
  886. #'test_dir': '/mnt/thsw2/DSP2/videos/obbShips/DJI_20230208110806_0001_W_6M.MP4',
  887. 'test_dir':'/mnt/thsw2/DSP2/videos/obbShips/freighter2.mp4',
  888. 'test_flag':True,
  889. 'postFile': '../AIlib2/weights/conf/%s/para.json'%(opt['business'] ),###后处理参数文件
  890. 'drawBox':True,#####是否画框
  891. 'drawPar': { 'digitWordFont' :{'line_thickness':2,'boxLine_thickness':1,'wordSize':40, 'fontSize':1.0,'label_location':'leftTop'}} ,
  892. 'labelnames':"../AIlib2/weights/conf/%s/labelnames.json"%(opt['business'] ), ###检测类别对照表
  893. }
  894. #par['model_size'],par['mean'],par['std'],par['half'],par['saveType'],par['heads'],par['labelnames'],par['decoder'],par['down_ratio'],par['drawBox']
  895. #par['rainbows'],par['label_array'],par['digitWordFont']
  896. obbModelPar = par['obbModelPar']
  897. ####加载模型
  898. model,decoder2=load_model_decoder_OBB(obbModelPar)
  899. obbModelPar['decoder']=decoder2
  900. names=get_labelnames(par['labelnames']);obbModelPar['labelnames']=names
  901. _,_,_,rainbows=get_postProcess_para(par['postFile']);par['drawPar']['rainbows']=rainbows
  902. label_arraylist = get_label_arrays(names,rainbows,outfontsize=par['drawPar']['digitWordFont']['wordSize'],fontpath="../AIlib2/conf/platech.ttf")
  903. #par['label_array']=label_arraylist
  904. trackPar=par['trackPar']
  905. sort_tracker = OBB_Sort(max_age=trackPar['sort_max_age'],
  906. min_hits=trackPar['sort_min_hits'],
  907. iou_threshold=trackPar['sort_iou_thresh'])
  908. ##图像测试和视频
  909. impth = par['test_dir']
  910. img_urls=[]###获取文件里所有的图像
  911. video_urls=[]###获取文件里所有的视频
  912. img_postfixs = ['.jpg','.JPG','.PNG','.png'];
  913. vides_postfixs= ['.MP4','.mp4','.avi']
  914. if os.path.isdir(impth):
  915. for postfix in img_postfixs:
  916. img_urls.extend(glob.glob('%s/*%s'%(impth,postfix )) )
  917. for postfix in ['.MP4','.mp4','.avi']:
  918. video_urls.extend(glob.glob('%s/*%s'%(impth,postfix )) )
  919. else:
  920. postfix = os.path.splitext(impth)[-1]
  921. if postfix in img_postfixs: img_urls=[ impth ]
  922. if postfix in vides_postfixs: video_urls = [impth ]
  923. parIn = {'obbModelPar':obbModelPar,'modelPar':{'obbmodel': model},'sort_tracker':sort_tracker,'outpth':par['outpth'],'trackPar':trackPar,'drawPar':par['drawPar']}
  924. par['drawPar']['label_array']=label_arraylist
  925. for img_url in img_urls:
  926. #print(img_url)
  927. ori_image=cv2.imread(img_url)
  928. #ori_image_list,infos = OBB_infer(model,ori_image,obbModelPar)
  929. ori_image_list,infos = OBB_tracker_batch([ori_image],[0],parIn['modelPar'],parIn['obbModelPar'],None,parIn['trackPar'],None)
  930. ori_image_list[1] = draw_obb(ori_image_list[2] ,ori_image_list[1],par['drawPar'])
  931. imgName = os.path.basename(img_url)
  932. saveFile = os.path.join(par['outpth'], imgName)
  933. ret=cv2.imwrite(saveFile, ori_image_list[1])
  934. if not ret:
  935. print(saveFile, ' not created ')
  936. print( os.path.basename(img_url),':',infos,ori_image_list[2])
  937. ###处理视频
  938. for video_url in video_urls:
  939. process_video(video_url, parIn ,mode='obbTrack')
  940. def crowd_demo(opt):
  941. if opt['business']=='crowdCounting':
  942. from crowd import crowdModel as Model
  943. par={
  944. 'mean':[0.485, 0.456, 0.406], 'std':[0.229, 0.224, 0.225],'threshold':0.5,
  945. 'input_profile_shapes':[(1,3,256,256),(1,3,1024,1024),(1,3,2048,2048)],
  946. 'modelPar':{'backbone':'vgg16_bn', 'gpu_id':0,'anchorFlag':False, 'width':None,'height':None ,'line':2, 'row':2},
  947. 'weights':"../weights/%s/AIlib2/%s/crowdCounting_%s_dynamic.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###检测模型路径
  948. 'testImgPath':'images/%s'%(opt['business'] ),###测试图像的位置
  949. 'testOutPath':'images/results/',###输出测试图像位置
  950. }
  951. #weights='weights/best_mae.pth'
  952. cmodel = Model(par['weights'],par)
  953. img_path = par['testImgPath']
  954. File = os.listdir(img_path)
  955. targetList = []
  956. for file in File[0:]:
  957. COORlist = []
  958. imgPath = img_path + os.sep + file
  959. img_raw = cv2.cvtColor(cv2.imread(imgPath),cv2.COLOR_BGR2RGB)
  960. # cmodel.eval---
  961. # 输入读取的RGB数组
  962. # 输出:list,0--原图,1-人头坐标list,2-对接OBB的格式数据,其中4个坐标均相同,2-格式如下:
  963. # [ [ [ (x0,y0),(x1,y1),(x2,y2),(x3,y3) ],score, cls ], [ [ (x0,y0),(x1,y1),(x2,y2),(x3,y3) ],score ,cls ],........ ]
  964. prets, infos = cmodel.eval(img_raw)
  965. print(file,infos,' 人数:',len(prets[1]))
  966. img_to_draw = cv2.cvtColor(np.array(img_raw), cv2.COLOR_RGB2BGR)
  967. # 打印预测图像中人头的个数
  968. for p in prets[1]:
  969. img_to_draw = cv2.circle(img_to_draw, (int(p[0]), int(p[1])), 2, (0, 255, 0), -1)
  970. COORlist.append((int(p[0]), int(p[1])))
  971. # 将各测试图像中的人头坐标存储在targetList中, 格式:[[(x1, y1),(x2, y2),...], [(X1, Y1),(X2, Y2),..], ...]
  972. targetList.append(COORlist)
  973. #time.sleep(2)
  974. # 保存预测图片
  975. cv2.imwrite(os.path.join(par['testOutPath'], file), img_to_draw)
  976. if __name__=="__main__":
  977. #jkm_demo()
  978. businessAll=['river', 'river2','highWay2','noParking','drowning','forest2','vehicle','pedestrian','smogfire' , 'AnglerSwimmer','channelEmergency', 'countryRoad','cityMangement','ship2','cityMangement2','cityRoad','illParking',"crowdCounting",'cityMangement3','crackMeasurement','pothole','channel2','riverT']
  979. businessAll = [ 'forest2' ]
  980. for busi in businessAll:
  981. print('-'*40,'beg to test:',busi,'-'*40)
  982. opt={'gpu':'2080Ti','business':busi}
  983. if busi in ['ship2']:
  984. OBB_track_demo(opt)
  985. elif opt['business'] in ['crowdCounting'] :
  986. crowd_demo(opt)
  987. elif opt['business'] in ['crackMeasurement','channel2'] :
  988. customization_demo(opt)
  989. else:
  990. #if opt['business'] in ['river','highWay2','noParking','drowning','']:
  991. det_track_demo_N(opt )