import torch import numpy as np import cv2 import time import os import sys sys.path.extend(['../AIlib2/obbUtils']) import matplotlib.pyplot as plt import func_utils import time import torchvision.transforms as transforms from obbmodels import ctrbox_net import decoder import tensorrt as trt import onnx import onnxruntime as ort sys.path.extend(['../AIlib2/utils']) #sys.path.extend(['../AIlib2/utils']) from plots import draw_painting_joint from copy import deepcopy from scipy import interpolate def obbTohbb(obb): obbarray=np.array(obb) x0=np.min(obbarray[:,0]) x1=np.max(obbarray[:,0]) y0=np.min(obbarray[:,1]) y1=np.max(obbarray[:,1]) return [x0,y0,x1,y1] def trt_version(): return trt.__version__ def torch_device_from_trt(device): if device == trt.TensorLocation.DEVICE: return torch.device("cuda") elif device == trt.TensorLocation.HOST: return torch.device("cpu") else: return TypeError("%s is not supported by torch" % device) def torch_dtype_from_trt(dtype): if dtype == trt.int8: return torch.int8 elif trt_version() >= '7.0' and dtype == trt.bool: return torch.bool elif dtype == trt.int32: return torch.int32 elif dtype == trt.float16: return torch.float16 elif dtype == trt.float32: return torch.float32 else: raise TypeError("%s is not supported by torch" % dtype) def segTrtForward(engine,inputs,contextFlag=False): if not contextFlag: context = engine.create_execution_context() else: context=contextFlag #with engine.create_execution_context() as context: #input_names=['images'];output_names=['output'] namess=[ engine.get_binding_name(index) for index in range(engine.num_bindings) ] input_names = [namess[0]];output_names=namess[1:] batch_size = inputs[0].shape[0] bindings = [None] * (len(input_names) + len(output_names)) # 创建输出tensor，并分配内存 outputs = [None] * len(output_names) for i, output_name in enumerate(output_names): idx = engine.get_binding_index(output_name)#通过binding_name找到对应的input_id dtype = torch_dtype_from_trt(engine.get_binding_dtype(idx))#找到对应的数据类型 shape = (batch_size,) + tuple(engine.get_binding_shape(idx))#找到对应的形状大小 device = torch_device_from_trt(engine.get_location(idx)) output = torch.empty(size=shape, dtype=dtype, device=device) #print('&'*10,'batch_size:',batch_size , 'device:',device,'idx:',idx,'shape:',shape,'dtype:',dtype,' device:',output.get_device()) outputs[i] = output #print('###line65:',output_name,i,idx,dtype,shape) bindings[idx] = output.data_ptr()#绑定输出数据指针 for i, input_name in enumerate(input_names): idx =engine.get_binding_index(input_name) bindings[idx] = inputs[0].contiguous().data_ptr()#应当为inputs[i]，对应3个输入。但由于我们使用的是单张图片，所以将3个输入全设置为相同的图片。 #print('#'*10,'input_names:,', input_name,'idx:',idx, inputs[0].dtype,', inputs[0] device:',inputs[0].get_device()) context.execute_v2(bindings) # 执行推理 if len(outputs) == 1: outputs = outputs[0] return outputs[0] else: return outputs def apply_mask(image, mask, alpha=0.5): """Apply the given mask to the image. """ color = np.random.rand(3) for c in range(3): image[:, :, c] = np.where(mask == 1, image[:, :, c] * (1 - alpha) + alpha * color[c] * 255, image[:, :, c]) return image if not os.path.exists('output'): os.mkdir('output') saveDir = 'output' def get_ms(t2,t1): return (t2-t1)*1000.0 def draw_painting_joint_2(box,img,label_array,score=0.5,color=None,font={ 'line_thickness':None,'boxLine_thickness':None, 'fontSize':None},socre_location="leftTop"): ###先把中文类别字体赋值到img中 lh, lw, lc = label_array.shape imh, imw, imc = img.shape if socre_location=='leftTop': x0 , y1 = box[0][0],box[0][1] elif socre_location=='leftBottom': x0,y1=box[3][0],box[3][1] else: print('plot.py line217 ,label_location:%s not implemented '%( socre_location )) sys.exit(0) x1 , y0 = x0 + lw , y1 - lh if y0<0:y0=0;y1=y0+lh if y1>imh: y1=imh;y0=y1-lh if x0<0:x0=0;x1=x0+lw if x1>imw:x1=imw;x0=x1-lw img[y0:y1,x0:x1,:] = label_array pts_cls=[(x0,y0),(x1,y1) ] #把四边形的框画上 box_tl= font['boxLine_thickness'] or round(0.002 * (imh + imw) / 2) + 1 cv2.polylines(img, [box], True,color , box_tl) ####把英文字符score画到类别旁边 tl = font['line_thickness'] or round(0.002*(imh+imw)/2)+1#line/font thickness label = ' %.2f'%(score) tf = max(tl , 1) # font thickness fontScale = font['fontSize'] or tl * 0.33 t_size = cv2.getTextSize(label, 0, fontScale=fontScale , thickness=tf)[0] #if socre_location=='leftTop': p1,p2= (pts_cls[1][0], pts_cls[0][1]),(pts_cls[1][0]+t_size[0],pts_cls[1][1]) cv2.rectangle(img, p1 , p2, color, -1, cv2.LINE_AA) p3 = pts_cls[1][0],pts_cls[1][1]-(lh-t_size[1])//2 cv2.putText(img, label,p3, 0, fontScale, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA) return img def OBB_infer(model,ori_image,par): ''' 输出：[img_origin,ori_image, out_box,9999],infos img_origin---原图 ori_image---画框图 out_box---检测目标框 ---格式如下[ [ [ (x0,y0),(x1,y1),(x2,y2),(x3,y3) ],score， cls ], [ [ (x0,y0),(x1,y1),(x2,y2),(x3,y3) ],score ，cls ],........ ],etc ---[ [ [(1159, 297), [922, 615], [817, 591], [1054, 272]], 0.865605354309082，14], [[(1330, 0), [1289, 58], [1228, 50], [1270, 0]], 0.3928087651729584，14] #2023.08.03,修改输出格式 ] 9999---无意义，备用 ''' t1 = time.time() #ori_image = cv2.imread(impth+folders[i]) t2 = time.time() img= cv2.resize(ori_image, (par['model_size'])) img_origin = ori_image.copy() t3 = time.time() transf2 = transforms.Compose([transforms.ToTensor(), transforms.Normalize(mean=par['mean'], std=par['std'])]) img_tensor = transf2(img) img_tensor1=img_tensor.unsqueeze(0) #转成了需要的tensor格式，中心归一化及颜色通道匹配上 t4=time.time() #print('###line170: resize:%.1f ToTensor-Normal-Destd:%.1f '%(get_ms(t3,t2),get_ms(t4,t3) ), img_origin.shape,img_tensor1.size() ) #img_tensor1= img_tensor1.to( par['device']) #布置到cuda上 img_tensor1 = img_tensor1.cuda() t5 =time.time() img_tensor1 = img_tensor1.half() if par['half'] else img_tensor1 if par['saveType']=='trt': preds= segTrtForward(model,[img_tensor1]) preds=[x[0] for x in preds ] pr_decs={} heads=list(par['heads'].keys()) pr_decs={ heads[i]: preds[i] for i in range(len(heads)) } elif par['saveType']=='pth': with torch.no_grad(): # no Back propagation pr_decs = model(img_tensor1) # 前向传播一部分 elif par['saveType']=='onnx': img=img_tensor1.cpu().numpy().astype(np.float32) preds = model['sess'].run(None, {model['input_name']: img}) pr_decs={} heads=list(par['heads'].keys()) pr_decs={ heads[i]: torch.from_numpy(preds[i]) for i in range(len(heads)) } t6 = time.time() category=par['labelnames'] #torch.cuda.synchronize(par['device']) # 时间异步变同步 decoded_pts = [] decoded_scores = [] predictions = par['decoder'].ctdet_decode(pr_decs) # 解码 t6_1=time.time() pts0, scores0 = func_utils.decode_prediction(predictions, category,par['model_size'], par['down_ratio'],ori_image) # 改3 decoded_pts.append(pts0) decoded_scores.append(scores0) t7 = time.time() # nms results = {cat: [] for cat in category} # '''这里啊 for cat in category: if cat == 'background': continue pts_cat = [] scores_cat = [] for pts0, scores0 in zip(decoded_pts, decoded_scores): pts_cat.extend(pts0[cat]) scores_cat.extend(scores0[cat]) pts_cat = np.asarray(pts_cat, np.float32) scores_cat = np.asarray(scores_cat, np.float32) if pts_cat.shape[0]: nms_results = func_utils.non_maximum_suppression(pts_cat, scores_cat) results[cat].extend(nms_results) t8 = time.time() height, width, _ = ori_image.shape # nms out_box=[] for cat in category: if cat == 'background': continue result = results[cat] for pred in result: score = pred[-1] cls = category.index(cat) boxF=[ max(int(x),0) for x in pred[0:8]] #box_out=[ cls,[ ( boxF[0], boxF[1]),([boxF[2], boxF[3]]), ([boxF[4], boxF[5]]), ([boxF[6], boxF[7]]) ],score] box_out=[ [ ( boxF[0], boxF[1]),([boxF[2], boxF[3]]), ([boxF[4], boxF[5]]), ([boxF[6], boxF[7]]) ],score,cls] ''' if par['drawBox']: tl = np.asarray([pred[0], pred[1]], np.float32) tr = np.asarray([pred[2], pred[3]], np.float32) br = np.asarray([pred[4], pred[5]], np.float32) bl = np.asarray([pred[6], pred[7]], np.float32) box = np.asarray([tl, tr, br, bl], np.int32) bgColor=par['rainbows'][cls%len( par['rainbows'])] label_array =par['label_array'][cls] font=par['digitWordFont'] label_location=font['label_location'] ori_image=draw_painting_joint(box,ori_image,label_array,score=score,color=bgColor,font=font,socre_location=label_location) ''' out_box.append(box_out) t9 = time.time() t10 = time.time() infos=' preProcess:%.1f ToGPU:%.1f infer:%.1f decoder:%.1f, corr_change:%.1f nms:%.1f postProcess:%.1f, total process:%.1f '%( get_ms(t4,t2), get_ms(t5,t4),get_ms(t6,t5),get_ms(t6_1,t6),get_ms(t7,t6_1),get_ms(t8,t7) ,get_ms(t9,t8) ,get_ms(t9,t2) ) #'preProcess:%.1f ToGPU:%.1f infer:%.1f decoder:%.1f, corr_change:%.1f nms:%.1f postProcess:%.1f, total process:%.1f '% #( get_ms(t4,t2), get_ms(t5,t4),get_ms(t6,t5),get_ms(t6_1,t6),get_ms(t7,t6_1), get_ms(t8,t7) ,get_ms(t9,t8) , get_ms(t9,t2) ) if len(out_box) > 0: ret_4pts = np.array([ x[0] for x in out_box ] ) ret_4pts = rectangle_quadrangle_batch (ret_4pts) cnt = len(out_box ) for ii in range(cnt): out_box[ii][0] = ret_4pts[ii] return [img_origin,ori_image, out_box,9999],infos def draw_obb(preds,ori_image,par): for pred in preds: box = np.asarray(pred[0][0:4],np.int32) cls = int(pred[2]);score = pred[1] bgColor=par['rainbows'][cls%len( par['rainbows'])] label_array =par['label_array'][cls] font=par['digitWordFont'] label_location=font['label_location'] #print('###line285:',box,cls,score) ori_image=draw_painting_joint(box,ori_image,label_array,score=score,color=bgColor,font=font,socre_location=label_location) #cv2.imwrite( 'test.jpg',ori_image ) return ori_image def OBB_tracker(sort_tracker,hbbs,obbs,iframe): #sort_tracker--跟踪器 #hbbs--目标的水平框[x0,y0,x1,y1] #obbs--目标的倾斜框box = np.asarray([tl, tr, br, bl], np.int32) #返回值：sort_tracker，跟踪器 dets_to_sort = np.empty((0,7), dtype=np.float32) # NOTE: We send in detected object class too for x1,y1,x2,y2,conf, detclass in hbbs: #print('#######line342:',x1,y1,x2,y2,img.shape,[x1, y1, x2, y2, conf, detclass,iframe]) dets_to_sort = np.vstack((dets_to_sort, np.array([x1, y1, x2, y2, conf, detclass,iframe],dtype=np.float32) )) # Run SORT tracked_dets = deepcopy(sort_tracker.update(dets_to_sort,obbs) ) return tracked_dets def rectangle_quadrangle(vectors): ##输入的是四个点偏离中心点的向量，(M,4,2) ##输出：vectors--修正后的向量（M,4,2） # wh_thetas--矩形的向量 (M,1,3)[w,h,theta] distans = np.sqrt(np.sum(vectors**2,axis=2))#(M,4) mean_dis = np.mean( distans,axis=1 ).reshape(-1,1) #(M,1) mean_dis = np.tile(mean_dis,(1,4) ) #(M,4) scale_factors = mean_dis/distans #(M,4) scale_factors = np.expand_dims(scale_factors, axis=2 ) #(M,4,1) scale_factors = np.tile(scale_factors, (1,1,2) ) #M(M,4,2) vectors = vectors*scale_factors vectors = vectors.astype(np.int32) cnt = vectors.shape[0] boxes = [ cv2.minAreaRect( vectors[i] ) for i in range(cnt) ] wh_thetas = [[x[1][0],x[1][1],x[2] ] for x in boxes]#(M,3),[w,h,theta] wh_thetas = np.array(wh_thetas)##(M,3) return vectors,wh_thetas def adjust_pts_orders(vectors): #输入一系列(M,4,2)点 #输入原定框顺序的(M,4,2) #前后两个四边形框一次判定，调整下一个四边形框内四个点的顺序，保证与上一个一致。 cnt = vectors.shape[0] if cnt<=1: return vectors else: out=[];out.append(vectors[0]) for i in range(1,cnt): pts1 = out[-1] pts2 = vectors[i] diss,min_dis,min_index,pts2_adjust = pts_setDistance(pts1,pts2) #if min_index!=0: print(min_index,pts1,pts2 ) out.append(pts2_adjust) out = np.array(out) #if out[4,0,0]==53 and out[4,0,1]==10: #print('#line339:',out.shape ,' ','in ', vectors.reshape(-1,8) , ' out :',out.reshape(-1,8)) return out def pts_setDistance(pts1,pts2): #输入是两个四边形的坐标（4,2），pts1保持不变，pts2逐个调整顺序，找到与pts2最匹配的四个点。 #输出pts2 原始的距离，最匹配点的距离，最匹配的点的序号 pts3=np.vstack((pts2,pts2)) diss =[np.sum((pts1-pts3[i:i+4])**2) for i in range(4)] min_dis = min(diss) min_index = diss.index(min_dis) return diss[0],min_dis,min_index,pts3[min_index:min_index+4] def obbPointsConvert(obbs): obbArray = np.array(obbs)#( M,4,2) #计算中心点 middlePts = np.mean( obbArray,axis=1 )##中心点（M,2） middlePts = np.expand_dims(middlePts,axis=1)#(M,1,2) #将中心点扩展成（M,4,2） vectors = np.tile(middlePts,(1,4,1))#(M,4,2) #计算偏移向量 vectors = obbArray - vectors #(M,4,2) ##校正偏移向量 vectors,wh_thetas=rectangle_quadrangle(vectors) #vectors--(M,4,2) ##校正每一个框内四个点的顺序 vectors = adjust_pts_orders(vectors) # (M,4,2) #将中心点附在偏移向量后面 vectors = np.concatenate( (vectors,middlePts),axis=1 )#(M,5,2), #将数据拉平 vectors = vectors.reshape(-1,10)#(M,10) return vectors def rectangle_quadrangle_batch(obbs): ##输入出四边形的四个点(M,4,2) ##输出是矩形话后的4个点（M,4,2） obbArray = np.array(obbs)#( M,4,2) #计算中心点 middlePts = np.mean( obbArray,axis=1 )##中心点（M,2） middlePts = np.expand_dims(middlePts,axis=1)#(M,1,2) #将中心点扩展成（M,4,2） middlePts = np.tile(middlePts,(1,4,1))#(M,4,2) #vectors = np.tile(middlePts,(1,4,1))#(M,4,2) #计算偏移向量 vectors = obbArray - middlePts #(M,4,2) ##校正偏移向量 vectors,wh_thetas=rectangle_quadrangle(vectors) #vectors--(M,4,2) vectors = vectors + middlePts return vectors def obbPointsConvert_reverse(vectors): vectors = np.array(vectors)#(M,10) _vectors = vectors[:,:8] #(M,8) middlePts = vectors[:,8:10] #(M,2) middlePts = np.tile( middlePts,(1,4) ) #(M,8) _vectors += middlePts #(M,8) return _vectors def OBB_tracker_batch(imgarray_list,iframe_list,modelPar,obbModelPar,sort_tracker,trackPar,segPar=None): ''' 输入： imgarray_list--图像列表 iframe_list -- 帧号列表 modelPar--模型参数,字典,modelPar={'det_Model':,'seg_Model':} obbModelpar--字典，存放检测相关参数，'half', 'device', 'conf_thres', 'iou_thres','trtFlag_det' sort_tracker--对象，初始化的跟踪对象。为了保持一致，即使是单帧也要有。 trackPar--跟踪参数，关键字包括：det_cnt，windowsize segPar--None,分割模型相关参数。如果用不到，则为None 输入：[imgarray_list,track_det_result,detResults ] , timeInfos # timeInfos---时间信息 # imgarray_list--图像列表 # track_det_result--numpy 格式(M,14)--( 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10 , 11, 12, 13 ) # (x0,y0,x1,y1,x2,y2,x3,y3,xc,yc,conf, detclass,iframe, trackId) # detResults---给DSP的结果.每一帧是一个list，内部每一个框时一个list，格式为[ [(x0,y0),(x1,y1),(x2,y2),(x3,y3)],score,cls ] 2023.08.03,修改输出格式 ''' det_cnt,windowsize = trackPar['det_cnt'] ,trackPar['windowsize'] trackers_dic={} index_list = list(range( 0, len(iframe_list) ,det_cnt )); if len(index_list)>1 and index_list[-1]!= iframe_list[-1]: index_list.append( len(iframe_list) - 1 ) #print('###line349:',index_list ,iframe_list) if len(imgarray_list)==1: #如果是单帧图片，则不用跟踪 ori_image_list,infos = OBB_infer(modelPar['obbmodel'],imgarray_list[0],obbModelPar) #print('##'*20,'line405:',np.array(ori_image_list[2]),ret_4pts ) return ori_image_list,infos else: timeInfos_track='' t1=time.time() for iframe_index, index_frame in enumerate(index_list): ori_image_list,infos = OBB_infer(modelPar['obbmodel'],imgarray_list[index_frame],obbModelPar) obbs = [x[0] for x in ori_image_list[2] ];hbbs = [] for i in range(len(ori_image_list[2])): hbb=obbTohbb( ori_image_list[2][i][0] ); box=[ *hbb, ori_image_list[2][i][1],ori_image_list[2][i][2]] hbbs.append(box) tracked_dets = OBB_tracker(sort_tracker,hbbs,obbs,iframe_list[index_frame] ) tracks =sort_tracker.getTrackers() tt=[tracker.id for tracker in tracks] for tracker in tracks: trackers_dic[tracker.id]=deepcopy(tracker) t2=time.time() track_det_result = np.empty((0,14)) ###( 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10 , 11, 12, 13 ) ###(x0,y0,x1,y1,x2,y2,x3,y3,xc,yc,conf, detclass,iframe, trackId) trackIdIndex=13;frameIndex=12 #print('###line372:',list(trackers_dic.keys())) for trackId in trackers_dic.keys(): tracker = trackers_dic[trackId] obb_history = np.array(tracker.obb_history) hbb_history = np.array(tracker.bbox_history) #print('#'*20,obb_history.shape ) if len(obb_history)<2: #print('#'*20, trackId, ' trace Cnt:',len(obb_history)) continue #原来格式 np.asarray([tl, tr, br, bl], np.int32)--->中心点到tl, tr, br, bl的向量 #print('###line381: 插值转换前 obb_history:',obb_history.shape, ' trackId:',trackId, ' \n' ,obb_history.reshape(-1,8) ) obb_history = obbPointsConvert(obb_history) #(M,10) #print('###line381: 插值前 obb_history:',obb_history.shape , ' hbb_history[:,4:7]:',hbb_history[:,4:7].shape, ' trackId:',trackId,'\n',obb_history) arrays_box = np.concatenate( (obb_history,hbb_history[:,4:7]),axis=1) arrays_box = arrays_box.transpose();frames=hbb_history[:,6] #frame_min--表示该批次图片的起始帧，如该批次是[1,100],则frame_min=1，[101,200]--frame_min=101 #frames[0]--表示该目标出现的起始帧，如[1,11,21,31,41],则frames[0]=1，frames[0]可能会在frame_min之前出现，即一个横跨了多个批次。 ##如果要最小化插值范围，则取内区间[frame_min,则frame_max ]和[frames[0],frames[-1] ]的交集 #inter_frame_min = int(max(frame_min, frames[0])); inter_frame_max = int(min( frame_max, frames[-1] )) ## ##如果要求得到完整的目标轨迹，则插值区间要以目标出现的起始点为准 inter_frame_min=int(frames[0]);inter_frame_max=int(frames[-1]) new_frames= np.linspace(inter_frame_min,inter_frame_max,inter_frame_max-inter_frame_min+1 ) #print('###line389:',trackId, inter_frame_min,inter_frame_max ,frames) #print(' ##line396: 插值前：' ,arrays_box) f_linear = interpolate.interp1d(frames,arrays_box); interpolation_x0s = (f_linear(new_frames)).transpose() move_cnt_use =(len(interpolation_x0s)+1)//2*2-1 if len(interpolation_x0s)[tl, tr, br, bl] interpolation_x0s[:,0:8] = obbPointsConvert_reverse(interpolation_x0s[:,0:10] ) #print('##line403: 插值转换后: ',interpolation_x0s.shape, inter_frame_min,inter_frame_max,frames, '\n',interpolation_x0s ) #for im in range(10): # interpolation_x0s[:,im] = moving_average_wang(interpolation_x0s[:,im],move_cnt_use ) cnt = inter_frame_max-inter_frame_min+1; trackIds = np.zeros((cnt,1)) + trackId interpolation_x0s = np.hstack( (interpolation_x0s, trackIds ) ) track_det_result = np.vstack(( track_det_result, interpolation_x0s) ) detResults=[] for iiframe in iframe_list: boxes_oneFrame = track_det_result[ track_det_result[:,frameIndex]==iiframe ] ###( 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10 , 11, 12, 13 ) ###(x0,y0,x1,y1,x2,y2,x3,y3,xc,yc,conf, detclass,iframe, trackId) res = [ [ [(b[0],b[1]),(b[2],b[3]),(b[4],b[5]),(b[6],b[7])],b[10],b[11],b[12],b[13] ] for b in boxes_oneFrame] detResults.append( res ) t3 = time.time() timeInfos='%d frames,detect and track:%.1f ,interpolation:%.1f '%( len(index_list), get_ms(t2,t1),get_ms(t3,t2) ) retResults=[imgarray_list,track_det_result,detResults ] return retResults, timeInfos