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.

505 lines
25KB

  1. import cv2,os,time,json
  2. from models.experimental import attempt_load
  3. from segutils.segmodel import SegModel,get_largest_contours
  4. from segutils.trtUtils import segtrtEval,yolov5Trtforward,OcrTrtForward
  5. from segutils.trafficUtils import trafficPostProcessing,colour_code_segmentation,get_label_info,trafficPostProcessingV2,tracfficAccidentMixFunction
  6. from utils.torch_utils import select_device
  7. from utilsK.queRiver import get_labelnames,get_label_arrays,post_process_,img_pad,draw_painting_joint,detectDraw,getDetections,getDetectionsFromPreds
  8. from trackUtils.sort import moving_average_wang
  9. from utils.datasets import letterbox
  10. import numpy as np
  11. import torch
  12. import math
  13. from PIL import Image
  14. import torch.nn.functional as F
  15. from copy import deepcopy
  16. from scipy import interpolate
  17. def xywh2xyxy(box,iW=None,iH=None):
  18. xc,yc,w,h = box[0:4]
  19. x0 =max(0, xc-w/2.0)
  20. x1 =min(1, xc+w/2.0)
  21. y0=max(0, yc-h/2.0)
  22. y1=min(1,yc+h/2.0)
  23. if iW: x0,x1 = x0*iW,x1*iW
  24. if iH: y0,y1 = y0*iH,y1*iH
  25. return [x0,y0,x1,y1]
  26. def get_ms(t2,t1):
  27. return (t2-t1)*1000.0
  28. def get_postProcess_para(parfile):
  29. with open(parfile) as fp:
  30. par = json.load(fp)
  31. assert 'post_process' in par.keys(), ' parfile has not key word:post_process'
  32. parPost=par['post_process']
  33. return parPost["conf_thres"],parPost["iou_thres"],parPost["classes"],parPost["rainbows"]
  34. def get_postProcess_para_dic(parfile):
  35. with open(parfile) as fp:
  36. par = json.load(fp)
  37. parPost=par['post_process']
  38. return parPost
  39. def AI_process(im0s,model,segmodel,names,label_arraylist,rainbows,objectPar={ 'half':True,'device':'cuda:0' ,'conf_thres':0.25,'iou_thres':0.45,'allowedList':[0,1,2,3],'segRegionCnt':1, 'trtFlag_det':False,'trtFlag_seg':False }, font={ 'line_thickness':None, 'fontSize':None,'boxLine_thickness':None,'waterLineColor':(0,255,255),'waterLineWidth':3} ,segPar={'modelSize':(640,360),'mean':(0.485, 0.456, 0.406),'std' :(0.229, 0.224, 0.225),'numpy':False, 'RGB_convert_first':True},mode='others',postPar=None):
  40. #输入参数
  41. # im0s---原始图像列表
  42. # model---检测模型,segmodel---分割模型(如若没有用到,则为None)
  43. #
  44. #输出:两个元素(列表,字符)构成的元组,[im0s[0],im0,det_xywh,iframe],strout
  45. # [im0s[0],im0,det_xywh,iframe]中,
  46. # im0s[0]--原始图像,im0--AI处理后的图像,iframe--帧号/暂时不需用到。
  47. # det_xywh--检测结果,是一个列表。
  48. # 其中每一个元素表示一个目标构成如:[float(cls_c), xc,yc,w,h, float(conf_c)]
  49. # #cls_c--类别,如0,1,2,3; xc,yc,w,h--中心点坐标及宽;conf_c--得分, 取值范围在0-1之间
  50. # #strout---统计AI处理个环节的时间
  51. # Letterbox
  52. half,device,conf_thres,iou_thres,allowedList = objectPar['half'],objectPar['device'],objectPar['conf_thres'],objectPar['iou_thres'],objectPar['allowedList']
  53. trtFlag_det,trtFlag_seg,segRegionCnt = objectPar['trtFlag_det'],objectPar['trtFlag_seg'],objectPar['segRegionCnt']
  54. if 'ovlap_thres_crossCategory' in objectPar.keys():
  55. ovlap_thres = objectPar['ovlap_thres_crossCategory']
  56. else:
  57. ovlap_thres = None
  58. time0=time.time()
  59. if trtFlag_det:
  60. img, padInfos = img_pad(im0s[0], size=(640,640,3)) ;img = [img]
  61. else:
  62. img = [letterbox(x, 640, auto=True, stride=32)[0] for x in im0s];padInfos=None
  63. # Stack
  64. img = np.stack(img, 0)
  65. # Convert
  66. img = img[:, :, :, ::-1].transpose(0, 3, 1, 2) # BGR to RGB, to bsx3x416x416
  67. img = np.ascontiguousarray(img)
  68. img = torch.from_numpy(img).to(device)
  69. img = img.half() if half else img.float() # uint8 to fp16/32
  70. time01=time.time()
  71. if segmodel:
  72. if trtFlag_seg:
  73. seg_pred,segstr = segtrtEval(segmodel,im0s[0],par=segPar)
  74. else:
  75. seg_pred,segstr = segmodel.eval(im0s[0] )
  76. segFlag=True
  77. else:
  78. seg_pred = None;segFlag=False;segstr='Not implemented'
  79. #if mode=='highWay3.0':
  80. # seg_pred_mulcls = seg_pred.copy()
  81. # #seg_pred = (seg_pred==1).astype(np.uint8) ###把路提取出来,路的类别是1
  82. time1=time.time()
  83. if trtFlag_det:
  84. pred = yolov5Trtforward(model,img)
  85. else:
  86. pred = model(img,augment=False)[0]
  87. time2=time.time()
  88. #p_result,timeOut = post_process_(datas,conf_thres, iou_thres,names,label_arraylist,rainbows,10,ObjectPar=ObjectPar,font=font,padInfos=padInfos,ovlap_thres=ovlap_thres)
  89. p_result, timeOut = getDetectionsFromPreds(pred,img,im0s[0],conf_thres=conf_thres,iou_thres=iou_thres,ovlap_thres=ovlap_thres,padInfos=padInfos)
  90. #if mode=='highWay3.0':
  91. if segmodel:
  92. #assert postPar , ' postPar not implemented'
  93. #det_coords_original = tracfficAccidentMixFunction(p_result[2],seg_pred_mulcls,segPar['mixFunction']['pars'])
  94. #p_result[2] = det_coords_original
  95. mixFunction = segPar['mixFunction']['function']
  96. p_result[2] , timeMixPost= mixFunction(p_result[2], seg_pred, pars=segPar['mixFunction']['pars'] )
  97. p_result.append(seg_pred)
  98. else:
  99. timeMixPost=':0 ms'
  100. #print('#### line121: segstr:%s timeMixPost:%s timeOut:%s'%( segstr.strip(), timeMixPost,timeOut ))
  101. time_info = 'letterbox:%.1f, seg:%.1f , infer:%.1f,%s, seginfo:%s ,timeMixPost:%s '%( (time01-time0)*1000, (time1-time01)*1000 ,(time2-time1)*1000,timeOut , segstr.strip(),timeMixPost )
  102. #if mode=='highWay3.0':
  103. return p_result,time_info
  104. def AI_Seg_process(im0s,segmodel,digitWordFont,trtFlag_seg=True,segPar={'modelSize':(640,360),'mean':(0.485, 0.456, 0.406),'std' :(0.229, 0.224, 0.225),'numpy':False, 'RGB_convert_first':True},postPar= {'label_csv': './AIlib2/weights/conf/trafficAccident/class_dict.csv', 'speedRoadArea': 5100, 'vehicleArea': 100, 'speedRoadVehicleAngleMin': 15, 'speedRoadVehicleAngleMax': 75, 'vehicleLengthWidthThreshold': 4, 'vehicleSafeDistance': 7}):
  105. '''
  106. 输入参数
  107. im0s---原始图像列表
  108. segmodel---分割模型,segmodel---分割模型(如若没有用到,则为None)
  109. digitWordFont--显示字体,数字等参数
  110. trtFlag_seg--模型是否是TRT格式
  111. segPar--分割模型的参数
  112. postPar--后处理参数
  113. 输出
  114. seg_pred--返回语义分割的结果图(0,1,2...表示)
  115. img_draw--原图上带有矩形框的图
  116. segstr-----文本数据包括时间信息
  117. list1-----返回目标的坐标结果,每一个目标用[ cls, x0,y0,x1,y1,conf ]
  118. '''
  119. time1=time.time()
  120. H,W=im0s[0].shape[0:2]
  121. img_draw=im0s[0].copy()
  122. if trtFlag_seg:
  123. seg_pred,segstr = segtrtEval(segmodel,im0s[0],par=segPar)
  124. else:
  125. seg_pred,segstr = segmodel.eval(im0s[0] )
  126. time2 = time.time()
  127. label_info = get_label_info(postPar['label_csv'])
  128. postPar['CCS']=colour_code_segmentation(seg_pred.copy(), label_info)
  129. postPar['sourceImageSize'] = im0s[0].shape[0:2]
  130. postPar['seg_pred_size'] = seg_pred.shape[0:2]
  131. list1,post_time_infos = trafficPostProcessing(postPar)
  132. list2=[]
  133. cls=0
  134. label_arraylist=digitWordFont['label_arraylist']
  135. rainbows=digitWordFont['rainbows']
  136. for bpoints in list1:
  137. #print('###line104:',bpoints)
  138. bpoints=np.array(bpoints)
  139. x0=np.min( bpoints[:,0] )
  140. y0=np.min( bpoints[:,1] )
  141. x1=np.max( bpoints[:,0] )
  142. y1=np.max( bpoints[:,1] )
  143. conf= ((x0+x1)/W + (y0+y1)/H)/4.0;
  144. conf=1.0 - math.fabs((conf-0.5)/0.5)
  145. xyxy=[x0,y0,x1,y1]
  146. xyxy=[int(x+0.5) for x in xyxy]
  147. #float(cls_c), *xywh, float(conf_c)]
  148. list2.append( [ cls, x0,y0,x1,y1,conf ] )
  149. img_draw = draw_painting_joint(xyxy,img_draw,label_arraylist[int(cls)],score=conf,color=rainbows[int(cls)%20],font=digitWordFont)
  150. segstr = 'segInfer:%.2f %s '%( (time2-time1)*1000.0,post_time_infos )
  151. return seg_pred,img_draw,segstr,list2
  152. def AI_process_v2(im0s,model,segmodel,names,label_arraylist,rainbows,half=True,device=' cuda:0',conf_thres=0.25, iou_thres=0.45,allowedList=[0,1,2,3], font={ 'line_thickness':None, 'fontSize':None,'boxLine_thickness':None,'waterLineColor':(0,255,255),'waterLineWidth':3} ):
  153. #输入参数
  154. # im0s---原始图像列表
  155. # model---检测模型,segmodel---分割模型(如若没有用到,则为None)
  156. #输出:两个元素(列表,字符)构成的元组,[im0s[0],im0,det_xywh,iframe],strout
  157. # [im0s[0],im0,det_xywh,iframe]中,
  158. # im0s[0]--原始图像,im0--AI处理后的图像,iframe--帧号/暂时不需用到。
  159. # det_xywh--检测结果,是一个列表。
  160. # 其中每一个元素表示一个目标构成如:[float(cls_c), xc,yc,w,h, float(conf_c)]
  161. # #cls_c--类别,如0,1,2,3; xc,yc,w,h--中心点坐标及宽;conf_c--得分, 取值范围在0-1之间
  162. # #strout---统计AI处理个环节的时间
  163. # Letterbox
  164. time0=time.time()
  165. #img = [letterbox(x, 640, auto=True, stride=32)[0] for x in im0s]
  166. img, padInfos = img_pad(im0s[0], size=(640,640,3)) ;img = [img]
  167. # Stack
  168. img = np.stack(img, 0)
  169. # Convert
  170. img = img[:, :, :, ::-1].transpose(0, 3, 1, 2) # BGR to RGB, to bsx3x416x416
  171. img = np.ascontiguousarray(img)
  172. img = torch.from_numpy(img).to(device)
  173. img = img.half() if half else img.float() # uint8 to fp16/32
  174. time01=time.time()
  175. img /= 255.0 # 0 - 255 to 0.0 - 1.0
  176. if segmodel:
  177. seg_pred,segstr = segmodel.eval(im0s[0] )
  178. segFlag=True
  179. else:
  180. seg_pred = None;segFlag=False
  181. time1=time.time()
  182. pred = model(img,augment=False)
  183. time2=time.time()
  184. datas = [[''], img, im0s, None,pred,seg_pred,10]
  185. p_result,timeOut = post_process_(datas,conf_thres, iou_thres,names,label_arraylist,rainbows,10,object_config=allowedList,segmodel=segFlag,font=font,padInfos=padInfos)
  186. time_info = 'letterbox:%.1f, seg:%.1f , infer:%.1f,%s, seginfo:%s'%( (time01-time0)*1000, (time1-time01)*1000 ,(time2-time1)*1000,timeOut , segstr )
  187. return p_result,time_info
  188. def AI_process_forest(im0s,model,segmodel,names,label_arraylist,rainbows,half=True,device=' cuda:0',conf_thres=0.25, iou_thres=0.45,allowedList=[0,1,2,3], font={ 'line_thickness':None, 'fontSize':None,'boxLine_thickness':None,'waterLineColor':(0,255,255),'waterLineWidth':3} ,trtFlag_det=False,SecNms=None):
  189. #输入参数
  190. # im0s---原始图像列表
  191. # model---检测模型,segmodel---分割模型(如若没有用到,则为None)
  192. #输出:两个元素(列表,字符)构成的元组,[im0s[0],im0,det_xywh,iframe],strout
  193. # [im0s[0],im0,det_xywh,iframe]中,
  194. # im0s[0]--原始图像,im0--AI处理后的图像,iframe--帧号/暂时不需用到。
  195. # det_xywh--检测结果,是一个列表。
  196. # 其中每一个元素表示一个目标构成如:[float(cls_c), xc,yc,w,h, float(conf_c)]
  197. # #cls_c--类别,如0,1,2,3; xc,yc,w,h--中心点坐标及宽;conf_c--得分, 取值范围在0-1之间
  198. # #strout---统计AI处理个环节的时间
  199. # Letterbox
  200. time0=time.time()
  201. if trtFlag_det:
  202. img, padInfos = img_pad(im0s[0], size=(640,640,3)) ;img = [img]
  203. else:
  204. img = [letterbox(x, 640, auto=True, stride=32)[0] for x in im0s];padInfos=None
  205. #img = [letterbox(x, 640, auto=True, stride=32)[0] for x in im0s]
  206. # Stack
  207. img = np.stack(img, 0)
  208. # Convert
  209. img = img[:, :, :, ::-1].transpose(0, 3, 1, 2) # BGR to RGB, to bsx3x416x416
  210. img = np.ascontiguousarray(img)
  211. img = torch.from_numpy(img).to(device)
  212. img = img.half() if half else img.float() # uint8 to fp16/32
  213. img /= 255.0 # 0 - 255 to 0.0 - 1.0
  214. if segmodel:
  215. seg_pred,segstr = segmodel.eval(im0s[0] )
  216. segFlag=True
  217. else:
  218. seg_pred = None;segFlag=False
  219. time1=time.time()
  220. pred = yolov5Trtforward(model,img) if trtFlag_det else model(img,augment=False)[0]
  221. time2=time.time()
  222. datas = [[''], img, im0s, None,pred,seg_pred,10]
  223. ObjectPar={ 'object_config':allowedList, 'slopeIndex':[] ,'segmodel':segFlag,'segRegionCnt':0 }
  224. p_result,timeOut = post_process_(datas,conf_thres, iou_thres,names,label_arraylist,rainbows,10,ObjectPar=ObjectPar,font=font,padInfos=padInfos,ovlap_thres=SecNms)
  225. #p_result,timeOut = post_process_(datas,conf_thres, iou_thres,names,label_arraylist,rainbows,10,object_config=allowedList,segmodel=segFlag,font=font,padInfos=padInfos)
  226. time_info = 'letterbox:%.1f, infer:%.1f, '%( (time1-time0)*1000,(time2-time1)*1000 )
  227. return p_result,time_info+timeOut
  228. def AI_det_track( im0s_in,modelPar,processPar,sort_tracker,segPar=None):
  229. im0s,iframe=im0s_in[0],im0s_in[1]
  230. model = modelPar['det_Model']
  231. segmodel = modelPar['seg_Model']
  232. half,device,conf_thres, iou_thres,trtFlag_det = processPar['half'], processPar['device'], processPar['conf_thres'], processPar['iou_thres'],processPar['trtFlag_det']
  233. iou2nd = processPar['iou2nd']
  234. time0=time.time()
  235. if trtFlag_det:
  236. img, padInfos = img_pad(im0s[0], size=(640,640,3)) ;img = [img]
  237. else:
  238. img = [letterbox(x, 640, auto=True, stride=32)[0] for x in im0s];padInfos=None
  239. img = np.stack(img, 0)
  240. # Convert
  241. img = img[:, :, :, ::-1].transpose(0, 3, 1, 2) # BGR to RGB, to bsx3x416x416
  242. img = np.ascontiguousarray(img)
  243. img = torch.from_numpy(img).to(device)
  244. img = img.half() if half else img.float() # uint8 to fp16/32
  245. img /= 255.0 # 0 - 255 to 0.0 - 1.0
  246. seg_pred = None;segFlag=False
  247. time1=time.time()
  248. pred = yolov5Trtforward(model,img) if trtFlag_det else model(img,augment=False)[0]
  249. time2=time.time()
  250. #p_result,timeOut = getDetections(datas,conf_thres, iou_thres,names,label_arraylist,rainbows,10,ObjectPar=ObjectPar,font=font,padInfos=padInfos)
  251. p_result, timeOut = getDetectionsFromPreds(pred,img,im0s[0],conf_thres=conf_thres,iou_thres=iou_thres,ovlap_thres=iou2nd,padInfos=padInfos)
  252. if segmodel:
  253. seg_pred,segstr = segtrtEval(segmodel,im0s[0],par=segPar) if segPar['trtFlag_seg'] else segmodel.eval(im0s[0] )
  254. mixFunction = segPar['mixFunction']['function']
  255. p_result[2],timeInfos = mixFunction(p_result[2], seg_pred, pars=segPar['mixFunction']['pars'] )
  256. else:
  257. timeInfos = ' '
  258. #print('######line341:',seg_pred.shape,np.max(seg_pred),np.min(seg_pred) , len(p_result[2]) )
  259. time_info = 'letterbox:%.1f, detinfer:%.1f, '%( (time1-time0)*1000,(time2-time1)*1000 )
  260. if sort_tracker:
  261. #在这里增加设置调用追踪器的频率
  262. #..................USE TRACK FUNCTION....................
  263. #pass an empty array to sort
  264. dets_to_sort = np.empty((0,7), dtype=np.float32)
  265. # NOTE: We send in detected object class too
  266. for detclass,x1,y1,x2,y2,conf in p_result[2]:
  267. #print('#######line342:',x1,y1,x2,y2,img.shape,[x1, y1, x2, y2, conf, detclass,iframe])
  268. dets_to_sort = np.vstack((dets_to_sort,
  269. np.array([x1, y1, x2, y2, conf, detclass,iframe],dtype=np.float32) ))
  270. # Run SORT
  271. tracked_dets = deepcopy(sort_tracker.update(dets_to_sort) )
  272. tracks =sort_tracker.getTrackers()
  273. p_result.append(tracked_dets) ###index=4
  274. p_result.append(tracks) ###index=5
  275. return p_result,time_info+timeOut+timeInfos
  276. def AI_det_track_batch(imgarray_list, iframe_list ,modelPar,processPar,sort_tracker,trackPar,segPar=None):
  277. '''
  278. 输入:
  279. imgarray_list--图像列表
  280. iframe_list -- 帧号列表
  281. modelPar--模型参数,字典,modelPar={'det_Model':,'seg_Model':}
  282. processPar--字典,存放检测相关参数,'half', 'device', 'conf_thres', 'iou_thres','trtFlag_det'
  283. sort_tracker--对象,初始化的跟踪对象。为了保持一致,即使是单帧也要有。
  284. trackPar--跟踪参数,关键字包括:det_cnt,windowsize
  285. segPar--None,分割模型相关参数。如果用不到,则为None
  286. 输入:retResults,timeInfos
  287. retResults:list
  288. retResults[0]--imgarray_list
  289. retResults[1]--所有结果用numpy格式,所有的检测结果,包括8类,每列分别是x1, y1, x2, y2, conf, detclass,iframe,trackId
  290. retResults[2]--所有结果用list表示,其中每一个元素为一个list,表示每一帧的检测结果,每一个结果是由多个list构成,每个list表示一个框,格式为[ cls , x0 ,y0 ,x1 ,y1 ,conf,ifrmae,trackId ],如 retResults[2][j][k]表示第j帧的第k个框。
  291. '''
  292. det_cnt,windowsize = trackPar['det_cnt'] ,trackPar['windowsize']
  293. trackers_dic={}
  294. index_list = list(range( 0, len(iframe_list) ,det_cnt ));
  295. if len(index_list)>1 and index_list[-1]!= iframe_list[-1]:
  296. index_list.append( len(iframe_list) - 1 )
  297. if len(imgarray_list)==1: #如果是单帧图片,则不用跟踪
  298. retResults = []
  299. p_result,timeOut = AI_det_track( [ [imgarray_list[0]] ,iframe_list[0] ],modelPar,processPar,None,segPar )
  300. ##下面4行内容只是为了保持格式一致
  301. detArray = np.array(p_result[2])
  302. cnt = detArray.shape[0];trackIds=np.zeros((cnt,1));iframes = np.zeros((cnt,1)) + iframe_list[0]
  303. detArray = np.hstack( (detArray[:,1:5], detArray[:,5:6] ,detArray[:,0:1],iframes, trackIds ) )
  304. res = [[ b[6],b[0],b[1],b[2],b[3],b[4],b[5],b[7] ] for b in detArray ]
  305. retResults=[imgarray_list,detArray,res ]
  306. return retResults,timeOut
  307. else:
  308. t0 = time.time()
  309. timeInfos_track=''
  310. for iframe_index, index_frame in enumerate(index_list):
  311. p_result,timeOut = AI_det_track( [ [imgarray_list[index_frame]] ,iframe_list[index_frame] ],modelPar,processPar,sort_tracker,segPar )
  312. timeInfos_track='%s:%s'%(timeInfos_track,timeOut)
  313. for tracker in p_result[5]:
  314. trackers_dic[tracker.id]=deepcopy(tracker)
  315. t1 = time.time()
  316. track_det_result = np.empty((0,8))
  317. for trackId in trackers_dic.keys():
  318. tracker = trackers_dic[trackId]
  319. bbox_history = np.array(tracker.bbox_history)
  320. if len(bbox_history)<2: continue
  321. ###把(x0,y0,x1,y1)转换成(xc,yc,w,h)
  322. xcs_ycs = (bbox_history[:,0:2] + bbox_history[:,2:4] )/2
  323. whs = bbox_history[:,2:4] - bbox_history[:,0:2]
  324. bbox_history[:,0:2] = xcs_ycs;bbox_history[:,2:4] = whs;
  325. arrays_box = bbox_history[:,0:7].transpose();frames=bbox_history[:,6]
  326. #frame_min--表示该批次图片的起始帧,如该批次是[1,100],则frame_min=1,[101,200]--frame_min=101
  327. #frames[0]--表示该目标出现的起始帧,如[1,11,21,31,41],则frames[0]=1,frames[0]可能会在frame_min之前出现,即一个横跨了多个批次。
  328. ##如果要最好化插值范围,则取内区间[frame_min,则frame_max ]和[frames[0],frames[-1] ]的交集
  329. #inter_frame_min = int(max(frame_min, frames[0])); inter_frame_max = int(min( frame_max, frames[-1] )) ##
  330. ##如果要求得到完整的目标轨迹,则插值区间要以目标出现的起始点为准
  331. inter_frame_min=int(frames[0]);inter_frame_max=int(frames[-1])
  332. new_frames= np.linspace(inter_frame_min,inter_frame_max,inter_frame_max-inter_frame_min+1 )
  333. f_linear = interpolate.interp1d(frames,arrays_box); interpolation_x0s = (f_linear(new_frames)).transpose()
  334. move_cnt_use =(len(interpolation_x0s)+1)//2*2-1 if len(interpolation_x0s)<windowsize else windowsize
  335. for im in range(4):
  336. interpolation_x0s[:,im] = moving_average_wang(interpolation_x0s[:,im],move_cnt_use )
  337. cnt = inter_frame_max-inter_frame_min+1; trackIds = np.zeros((cnt,1)) + trackId
  338. interpolation_x0s = np.hstack( (interpolation_x0s, trackIds ) )
  339. track_det_result = np.vstack(( track_det_result, interpolation_x0s) )
  340. #print('#####line116:',trackId,frame_min,frame_max,'----------',interpolation_x0s.shape,track_det_result.shape ,'-----')
  341. ##将[xc,yc,w,h]转为[x0,y0,x1,y1]
  342. x0s = track_det_result[:,0] - track_det_result[:,2]/2 ; x1s = track_det_result[:,0] + track_det_result[:,2]/2
  343. y0s = track_det_result[:,1] - track_det_result[:,3]/2 ; y1s = track_det_result[:,1] + track_det_result[:,3]/2
  344. track_det_result[:,0] = x0s; track_det_result[:,1] = y0s;
  345. track_det_result[:,2] = x1s; track_det_result[:,3] = y1s;
  346. detResults=[]
  347. for iiframe in iframe_list:
  348. boxes_oneFrame = track_det_result[ track_det_result[:,6]==iiframe ]
  349. res = [[ b[6],b[0],b[1],b[2],b[3],b[4],b[5],b[7] ] for b in boxes_oneFrame ]
  350. #[ cls , x0 ,y0 ,x1 ,y1 ,conf,ifrmae,trackId ]
  351. detResults.append( res )
  352. retResults=[imgarray_list,track_det_result,detResults ]
  353. t2 = time.time()
  354. timeInfos = 'detTrack:%.1f TrackPost:%.1f, %s'%(get_ms(t1,t0),get_ms(t2,t1), timeInfos_track )
  355. return retResults,timeInfos
  356. def ocr_process(pars):
  357. img_patch,engine,context,converter,AlignCollate_normal,device=pars[0:6]
  358. time1 = time.time()
  359. img_tensor = AlignCollate_normal([ Image.fromarray(img_patch,'L') ])
  360. img_input = img_tensor.to('cuda:0')
  361. time2 = time.time()
  362. preds,trtstr=OcrTrtForward(engine,[img_input],context)
  363. time3 = time.time()
  364. batch_size = preds.size(0)
  365. preds_size = torch.IntTensor([preds.size(1)] * batch_size)
  366. ######## filter ignore_char, rebalance
  367. preds_prob = F.softmax(preds, dim=2)
  368. preds_prob = preds_prob.cpu().detach().numpy()
  369. pred_norm = preds_prob.sum(axis=2)
  370. preds_prob = preds_prob/np.expand_dims(pred_norm, axis=-1)
  371. preds_prob = torch.from_numpy(preds_prob).float().to(device)
  372. _, preds_index = preds_prob.max(2)
  373. preds_index = preds_index.view(-1)
  374. time4 = time.time()
  375. preds_str = converter.decode_greedy(preds_index.data.cpu().detach().numpy(), preds_size.data)
  376. time5 = time.time()
  377. info_str= ('pre-process:%.2f TRTforward:%.2f (%s) postProcess:%2.f decoder:%.2f, Total:%.2f , pred:%s'%(get_ms(time2,time1 ),get_ms(time3,time2 ),trtstr, get_ms(time4,time3 ), get_ms(time5,time4 ), get_ms(time5,time1 ), preds_str ) )
  378. return preds_str,info_str
  379. def main():
  380. ##预先设置的参数
  381. device_='1' ##选定模型,可选 cpu,'0','1'
  382. ##以下参数目前不可改
  383. Detweights = "weights/yolov5/class5/best_5classes.pt"
  384. seg_nclass = 2
  385. Segweights = "weights/BiSeNet/checkpoint.pth"
  386. conf_thres,iou_thres,classes= 0.25,0.45,5
  387. labelnames = "weights/yolov5/class5/labelnames.json"
  388. 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]]
  389. allowedList=[0,1,2,3]
  390. ##加载模型,准备好显示字符
  391. device = select_device(device_)
  392. names=get_labelnames(labelnames)
  393. label_arraylist = get_label_arrays(names,rainbows,outfontsize=40,fontpath="conf/platech.ttf")
  394. half = device.type != 'cpu' # half precision only supported on CUDA
  395. model = attempt_load(Detweights, map_location=device) # load FP32 model
  396. if half: model.half()
  397. segmodel = SegModel(nclass=seg_nclass,weights=Segweights,device=device)
  398. ##图像测试
  399. #url='images/examples/20220624_响水河_12300_1621.jpg'
  400. impth = 'images/examples/'
  401. outpth = 'images/results/'
  402. folders = os.listdir(impth)
  403. for i in range(len(folders)):
  404. imgpath = os.path.join(impth, folders[i])
  405. im0s=[cv2.imread(imgpath)]
  406. time00 = time.time()
  407. p_result,timeOut = AI_process(im0s,model,segmodel,names,label_arraylist,rainbows,half,device,conf_thres, iou_thres,allowedList,fontSize=1.0)
  408. time11 = time.time()
  409. image_array = p_result[1]
  410. cv2.imwrite( os.path.join( outpth,folders[i] ) ,image_array )
  411. #print('----process:%s'%(folders[i]), (time.time() - time11) * 1000)
  412. if __name__=="__main__":
  413. main()