Browse Source

add cityMangement3

master
wangjin0928 10 months ago
parent
commit
ee02ab0883
34 changed files with 803 additions and 437 deletions
  1. +176
    -2
      AI.py
  2. +1
    -2
      DMPR.py
  3. +1
    -1
      DMPRUtils/DMPR_process.py
  4. BIN
      DMPRUtils/__pycache__/DMPR_process.cpython-38.pyc
  5. BIN
      DMPRUtils/__pycache__/jointUtil.cpython-38.pyc
  6. +24
    -1
      DMPRUtils/jointUtil.py
  7. +2
    -2
      DMPRUtils/toTrt.sh
  8. BIN
      __pycache__/AI.cpython-38.pyc
  9. BIN
      __pycache__/DMPR.cpython-38.pyc
  10. BIN
      __pycache__/crowd.cpython-38.pyc
  11. BIN
      __pycache__/stdc.cpython-38.pyc
  12. BIN
      __pycache__/yolov5.cpython-38.pyc
  13. +1
    -1
      crowd.py
  14. +5
    -1
      readme.md
  15. BIN
      segutils/__pycache__/trafficUtils.cpython-38.pyc
  16. BIN
      segutils/__pycache__/trtUtils2.cpython-38.pyc
  17. +2
    -2
      segutils/run.sh
  18. +4
    -1
      segutils/trafficUtils.py
  19. +2
    -2
      segutils/trtUtils2.py
  20. +3
    -3
      stdc.py
  21. BIN
      utilsK/__pycache__/drownUtils.cpython-38.pyc
  22. BIN
      utilsK/__pycache__/illParkingUtils.cpython-38.pyc
  23. BIN
      utilsK/__pycache__/noParkingUtils.cpython-38.pyc
  24. BIN
      utilsK/__pycache__/queRiver.cpython-38.pyc
  25. +3
    -1
      utilsK/drownUtils.py
  26. +4
    -1
      utilsK/illParkingUtils.py
  27. +464
    -416
      utilsK/noParkingUtils.py
  28. +3
    -1
      utilsK/queRiver.py
  29. BIN
      weights/conf/cityMangement3/dmpr.pth
  30. +4
    -0
      weights/conf/cityMangement3/labelnames.json
  31. +7
    -0
      weights/conf/cityMangement3/para.json
  32. BIN
      weights/conf/cityMangement3/stdc_360X640.pth
  33. BIN
      weights/conf/cityMangement3/yolov5.pt
  34. +97
    -0
      yolov5.py

+ 176
- 2
AI.py View File

@@ -17,6 +17,24 @@ from PIL import Image
import torch.nn.functional as F
from copy import deepcopy
from scipy import interpolate
import glob
def get_images_videos(impth, imageFixs=['.jpg','.JPG','.PNG','.png'],videoFixs=['.MP4','.mp4','.avi']):
imgpaths=[];###获取文件里所有的图像
videopaths=[]###获取文件里所有的视频
if os.path.isdir(impth):
for postfix in imageFixs:
imgpaths.extend(glob.glob('%s/*%s'%(impth,postfix )) )
for postfix in videoFixs:
videopaths.extend(glob.glob('%s/*%s'%(impth,postfix )) )
else:
postfix = os.path.splitext(impth)[-1]
if postfix in imageFixs: imgpaths=[ impth ]
if postfix in videoFixs: videopaths = [impth ]
print('%s: test Images:%d , test videos:%d '%(impth, len(imgpaths), len(videopaths)))
return imgpaths,videopaths
def xywh2xyxy(box,iW=None,iH=None):
xc,yc,w,h = box[0:4]
@@ -141,6 +159,39 @@ def AI_process(im0s,model,segmodel,names,label_arraylist,rainbows,objectPar={ 'h
return p_result,time_info
def default_mix(predlist,par):
return predlist[0],''
def AI_process_N(im0s,modelList,postProcess):
#输入参数
# im0s---原始图像列表
# model---检测模型,segmodel---分割模型(如若没有用到,则为None)
#
#输出:两个元素(列表,字符)构成的元组,[im0s[0],im0,det_xywh,iframe],strout
# [im0s[0],im0,det_xywh,iframe]中,
# im0s[0]--原始图像,im0--AI处理后的图像,iframe--帧号/暂时不需用到。
# det_xywh--检测结果,是一个列表。
# 其中每一个元素表示一个目标构成如:[ xc,yc,w,h, float(conf_c),float(cls_c) ] ,2023.08.03修改输出格式
# #cls_c--类别,如0,1,2,3; xc,yc,w,h--中心点坐标及宽;conf_c--得分, 取值范围在0-1之间
# #strout---统计AI处理个环节的时间
# Letterbox
modelRets=[ model.eval(im0s[0]) for model in modelList]
timeInfos = [ x[1] for x in modelRets]
timeInfos=''.join(timeInfos)
timeInfos=timeInfos
mixFunction =postProcess['function']
predsList = [ modelRet[0] for modelRet in modelRets ]
H,W = im0s[0].shape[0:2]
postProcess['pars']['imgSize'] = (W,H)
ret = mixFunction( predsList, postProcess['pars'])
return ret[0],timeInfos+ret[1]
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):
#输入参数
@@ -198,7 +249,7 @@ def AI_det_track( im0s_in,modelPar,processPar,sort_tracker,segPar=None):
half,device,conf_thres, iou_thres,trtFlag_det = processPar['half'], processPar['device'], processPar['conf_thres'], processPar['iou_thres'],processPar['trtFlag_det']
if 'score_byClass' in processPar.keys(): score_byClass = processPar['score_byClass']
else: score_byClass = None
#print('-'*10,'line201:',score_byClass)
iou2nd = processPar['iou2nd']
time0=time.time()
@@ -382,6 +433,129 @@ def AI_det_track_batch(imgarray_list, iframe_list ,modelPar,processPar,sort_trac
t2 = time.time()
timeInfos = 'detTrack:%.1f TrackPost:%.1f, %s'%(get_ms(t1,t0),get_ms(t2,t1), timeInfos_track )
return retResults,timeInfos
def AI_det_track_N( im0s_in,modelList,postProcess,sort_tracker):
im0s,iframe=im0s_in[0],im0s_in[1]
dets = AI_process_N(im0s,modelList,postProcess)
p_result=[[],[],dets[0],[] ]
if sort_tracker:
#在这里增加设置调用追踪器的频率
#..................USE TRACK FUNCTION....................
#pass an empty array to sort
dets_to_sort = np.empty((0,7), dtype=np.float32)
# NOTE: We send in detected object class too
#for detclass,x1,y1,x2,y2,conf in p_result[2]:
for x1,y1,x2,y2,conf, detclass in p_result[2]:
#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) )
tracks =sort_tracker.getTrackers()
p_result.append(tracked_dets) ###index=4
p_result.append(tracks) ###index=5
return p_result,dets[1]
def AI_det_track_batch_N(imgarray_list, iframe_list ,modelList,postProcess,sort_tracker,trackPar):
'''
输入:
imgarray_list--图像列表
iframe_list -- 帧号列表
modelPar--模型参数,字典,modelPar={'det_Model':,'seg_Model':}
processPar--字典,存放检测相关参数,'half', 'device', 'conf_thres', 'iou_thres','trtFlag_det'
sort_tracker--对象,初始化的跟踪对象。为了保持一致,即使是单帧也要有。
trackPar--跟踪参数,关键字包括:det_cnt,windowsize
segPar--None,分割模型相关参数。如果用不到,则为None
输入:retResults,timeInfos
retResults:list
retResults[0]--imgarray_list
retResults[1]--所有结果用numpy格式,所有的检测结果,包括8类,每列分别是x1, y1, x2, y2, conf, detclass,iframe,trackId
retResults[2]--所有结果用list表示,其中每一个元素为一个list,表示每一帧的检测结果,每一个结果是由多个list构成,每个list表示一个框,格式为[ x0 ,y0 ,x1 ,y1 ,conf, cls ,ifrmae,trackId ],如 retResults[2][j][k]表示第j帧的第k个框。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 )
if len(imgarray_list)==1: #如果是单帧图片,则不用跟踪
retResults = []
p_result,timeOut = AI_det_track_N( [ [imgarray_list[0]] ,iframe_list[0] ],modelList,postProcess,None )
##下面4行内容只是为了保持格式一致
detArray = np.array(p_result[2])
if len(p_result[2])==0:res=[]
else:
cnt = detArray.shape[0];trackIds=np.zeros((cnt,1));iframes = np.zeros((cnt,1)) + iframe_list[0]
#detArray = np.hstack( (detArray[:,1:5], detArray[:,5:6] ,detArray[:,0:1],iframes, trackIds ) )
detArray = np.hstack( (detArray[:,0:4], detArray[:,4:6] ,iframes, trackIds ) ) ##2023.08.03 修改输入格式
res = [[ b[0],b[1],b[2],b[3],b[4],b[5],b[6],b[7] ] for b in detArray ]
retResults=[imgarray_list,detArray,res ]
#print('##line380:',retResults[2])
return retResults,timeOut
else:
t0 = time.time()
timeInfos_track=''
for iframe_index, index_frame in enumerate(index_list):
p_result,timeOut = AI_det_track_N( [ [imgarray_list[index_frame]] ,iframe_list[index_frame] ],modelList,postProcess,sort_tracker )
timeInfos_track='%s:%s'%(timeInfos_track,timeOut)
for tracker in p_result[5]:
trackers_dic[tracker.id]=deepcopy(tracker)
t1 = time.time()
track_det_result = np.empty((0,8))
for trackId in trackers_dic.keys():
tracker = trackers_dic[trackId]
bbox_history = np.array(tracker.bbox_history)
if len(bbox_history)<2: continue
###把(x0,y0,x1,y1)转换成(xc,yc,w,h)
xcs_ycs = (bbox_history[:,0:2] + bbox_history[:,2:4] )/2
whs = bbox_history[:,2:4] - bbox_history[:,0:2]
bbox_history[:,0:2] = xcs_ycs;bbox_history[:,2:4] = whs;
arrays_box = bbox_history[:,0:7].transpose();frames=bbox_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 )
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)<windowsize else windowsize
for im in range(4):
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) )
#print('#####line116:',trackId,frame_min,frame_max,'----------',interpolation_x0s.shape,track_det_result.shape ,'-----')
##将[xc,yc,w,h]转为[x0,y0,x1,y1]
x0s = track_det_result[:,0] - track_det_result[:,2]/2 ; x1s = track_det_result[:,0] + track_det_result[:,2]/2
y0s = track_det_result[:,1] - track_det_result[:,3]/2 ; y1s = track_det_result[:,1] + track_det_result[:,3]/2
track_det_result[:,0] = x0s; track_det_result[:,1] = y0s;
track_det_result[:,2] = x1s; track_det_result[:,3] = y1s;
detResults=[]
for iiframe in iframe_list:
boxes_oneFrame = track_det_result[ track_det_result[:,6]==iiframe ]
res = [[ b[0],b[1],b[2],b[3],b[4],b[5],b[6],b[7] ] for b in boxes_oneFrame ]
#[ x0 ,y0 ,x1 ,y1 ,conf,cls,ifrmae,trackId ]
#[ifrmae, x0 ,y0 ,x1 ,y1 ,conf,cls,trackId ]
detResults.append( res )
retResults=[imgarray_list,track_det_result,detResults ]
t2 = time.time()
timeInfos = 'detTrack:%.1f TrackPost:%.1f, %s'%(get_ms(t1,t0),get_ms(t2,t1), timeInfos_track )
return retResults,timeInfos
def ocr_process(pars):
@@ -456,4 +630,4 @@ def main():
if __name__=="__main__":
main()
main()

+ 1
- 2
DMPR.py View File

@@ -26,7 +26,6 @@ class DMPRModel(object):
logger = trt.Logger(trt.Logger.ERROR)
with open(weights, "rb") as f, trt.Runtime(logger) as runtime:
self.model=runtime.deserialize_cuda_engine(f.read())# 输入trt本地文件,返回ICudaEngine对象
print('############load seg model trt success: ',weights)
elif self.infer_type=='pth':
#self.model = DirectionalPointDetector(3, self.par['depth_factor'], self.par['NUM_FEATURE_MAP_CHANNEL']).to(self.device)
confUrl = os.path.join( os.path.dirname(__file__),'DMPRUtils','config','yolov5s.yaml' )
@@ -34,7 +33,7 @@ class DMPRModel(object):
self.model.load_state_dict(torch.load(weights))
print('#######load pt model:%s success '%(weights))
self.par['modelType']=self.infer_type
print('#########加载模型:',weights,' 类型:',self.infer_type)
def eval(self,image):
det,timeInfos = DMPR_process(image, self.model, self.device, self.par)
det = det.cpu().detach().numpy()

+ 1
- 1
DMPRUtils/DMPR_process.py View File

@@ -205,7 +205,7 @@ def DMPR_process(img0, model, device, DMPRmodelPar):
t3 = time.time()
timeInfos = 'lettbox:%.1f dectect:%.1f scaleBack:%.1f '%( (t1-t0)*1000,(t2-t1)*1000,(t3-t2)*1000, )
timeInfos = 'dmpr:%1.f (lettbox:%.1f dectect:%.1f scaleBack:%.1f) '%( (t3-t0)*1000,(t1-t0)*1000,(t2-t1)*1000,(t3-t2)*1000, )
return det,timeInfos

BIN
DMPRUtils/__pycache__/DMPR_process.cpython-38.pyc View File


BIN
DMPRUtils/__pycache__/jointUtil.cpython-38.pyc View File


+ 24
- 1
DMPRUtils/jointUtil.py View File

@@ -4,9 +4,11 @@ import numpy as np
import torch
import time
def dmpr_yolo( yolo_det, dmpr_det,pars):
#if len(yolo_det)==0 or len(dmpr_det)==0:
#print('line11:\n',yolo_det, dmpr_det,pars)
time1=time.time()
if len(yolo_det)==0:
return yolo_det,' No yolo detections'
@@ -109,4 +111,25 @@ def dmpr_yolo( yolo_det, dmpr_det,pars):
yolo_det_clone = yolo_det_clone[:,0:6]
time2=time.time()
return np.array(yolo_det_clone), '%.1f'%( (time2-time1)*1000 )
return np.array(yolo_det_clone), 'dmpr_yolo:%.1f'%( (time2-time1)*1000 )
def stdc_yolo(stdc_det, yolo_det):
im = np.uint8(stdc_det)
x_c = ((yolo_det[:, 0] + yolo_det[:, 2]) // 2).astype(int)
y_c = ((yolo_det[:, 1] + yolo_det[:, 3]) // 2).astype(int)
yolo_filted = yolo_det[im[y_c, x_c] == 0]
return yolo_filted
def dmpr_yolo_stdc(predsList,pars):
if len(predsList)==2:
yolo_det, dmpr_det = predsList[0:2]
else:
yolo_det, dmpr_det,stdc_det = predsList[0:3]
if len(yolo_det)==0:
return yolo_det,' No yolo detections'
if isinstance(yolo_det,list):
yolo_det = np.array(yolo_det)
if len(predsList)>2:
yolo_det = stdc_yolo(stdc_det, yolo_det)
return dmpr_yolo(yolo_det, dmpr_det,pars)

+ 2
- 2
DMPRUtils/toTrt.sh View File

@@ -1,5 +1,5 @@
weights=/mnt/thsw2/DSP2/weights/cityMangement2/dmpr
weights=/mnt/thsw2/DSP2/weights/cityMangement3/dmpr
#weights=/mnt/thsw2/DSP2/weights/cityMangement2_0916/weights/urbanManagement/DMPR/dp_detector_299
gpu=3090
gpu=2080Ti
python toTrt.py --weights ${weights}.pth
mv ${weights}.engine ${weights}_${gpu}.engine

BIN
__pycache__/AI.cpython-38.pyc View File


BIN
__pycache__/DMPR.cpython-38.pyc View File


BIN
__pycache__/crowd.cpython-38.pyc View File


BIN
__pycache__/stdc.cpython-38.pyc View File


BIN
__pycache__/yolov5.cpython-38.pyc View File


+ 1
- 1
crowd.py View File

@@ -64,7 +64,7 @@ class crowdModel(object):
elif self.infer_type=='pth':
self.anchors = AnchorPointsf(pyramid_levels=[3,], strides=None, row=self.modelPar.row, line=self.modelPar.line ,device='cuda:0')
print('#########加载模型:',weights,' 类型:',self.infer_type)
def preprocess(self,img):
tmpImg = preprocess(img,mean=self.mean, std=self.std,minShape=self.minShape,maxShape=self.maxShape)

+ 5
- 1
readme.md View File

@@ -88,4 +88,8 @@
1.更新cityRoad的模型权重
2023.10.20
1.增加crowdCounting模型,统计图片中的人群数量,返回的是人群的坐标。
2.增进trtUtils2 ,里面包括了动态onnx,trt模型推理使用方法。
2.增进trtUtils2 ,里面包括了动态onnx,trt模型推理使用方法。
2023.11.06
1.增加cityMangement3,和cityMangement2一样的检测目标。但用了三个模型。
2.所有的demo3.0.py采用模型列表的方式输入参数
3.从cityMangement3 没有更新A100的trt文件和权重,因为云端A100服务器不存在了

BIN
segutils/__pycache__/trafficUtils.cpython-38.pyc View File


BIN
segutils/__pycache__/trtUtils2.cpython-38.pyc View File


+ 2
- 2
segutils/run.sh View File

@@ -1,6 +1,6 @@
gpu=2080Ti
name=noParking
nclass=4
name=cityMangement3
nclass=2
mWidth=640
mHeight=360
python segmodel_trt.py --weights /mnt/thsw2/DSP2/weights/${name}/stdc_360X640.pth --nclass ${nclass} --mWidth ${mWidth} --mHeight ${mHeight}

+ 4
- 1
segutils/trafficUtils.py View File

@@ -514,4 +514,7 @@ def tracfficAccidentMixFunction(preds,seg_pred_mulcls,pars):
time_infos=" no tracfficAccidentMix process"
#p_result[2]= deepcopy(det_coords_original)
return deepcopy(det_coords_original),time_infos
return deepcopy(det_coords_original),time_infos
def tracfficAccidentMixFunction_N(predList,pars):
preds,seg_pred_mulcls = predList[0:2]
return tracfficAccidentMixFunction(preds,seg_pred_mulcls,pars)

+ 2
- 2
segutils/trtUtils2.py View File

@@ -62,7 +62,7 @@ import onnxruntime
import os,sys,cv2
#from model.u2net import U2NET
#cuda.init()
model_names = ['u2net.onnx','u2net_dynamic_batch.onnx','u2net_dynamic_hw.onnx','u2net_dynamic_batch-hw.onnx' ]
dynamic_batch = {'input':{0:'batch'},
'output0':{0:'batch'},
@@ -461,4 +461,4 @@ if __name__=='__main__':
print(e)
else:
print(f'Input:{i} on model:{model_name} succeed')
'''
'''

+ 3
- 3
stdc.py View File

@@ -28,10 +28,10 @@ class stdcModel(object):
logger = trt.Logger(trt.Logger.ERROR)
with open(weights, "rb") as f, trt.Runtime(logger) as runtime:
self.model=runtime.deserialize_cuda_engine(f.read())# 输入trt本地文件,返回ICudaEngine对象
print('############load seg model trt success: ',weights)
elif self.infer_type=='pth':
self.model = SegModel(nclass=par['seg_nclass'],weights=weights,device=self.device)
print('############load seg model pth success:',weights,self.model,seg_nclass )
print('############load seg model pth success:',weights,self.model,par['seg_nclass'] )
print('#########加载模型:',weights,' 类型:',self.infer_type)
def eval(self,image):
@@ -68,6 +68,6 @@ class stdcModel(object):
else:
pred = cv2.resize(pred.astype(np.uint8),(W,H))
time4_0 = time.time()
segInfoStr= 'pre-precess:%.1f ,infer:%.1f ,post-cpu-argmax:%.1f ,post-resize:%.1f, total:%.1f '%( self.get_ms(time1_0,time0_0),self.get_ms(time2_0,time1_0),self.get_ms(time3_0,time2_0),self.get_ms(time4_0,time3_0),self.get_ms(time4_0,time0_0) )
segInfoStr= 'stdc:%.1f (pre-precess:%.1f ,infer:%.1f ,post-cpu-argmax:%.1f ,post-resize:%.1f) '%(self.get_ms(time4_0,time0_0), self.get_ms(time1_0,time0_0),self.get_ms(time2_0,time1_0),self.get_ms(time3_0,time2_0),self.get_ms(time4_0,time3_0) )
return pred,segInfoStr

BIN
utilsK/__pycache__/drownUtils.cpython-38.pyc View File


BIN
utilsK/__pycache__/illParkingUtils.cpython-38.pyc View File


BIN
utilsK/__pycache__/noParkingUtils.cpython-38.pyc View File


BIN
utilsK/__pycache__/queRiver.cpython-38.pyc View File


+ 3
- 1
utilsK/drownUtils.py View File

@@ -220,5 +220,7 @@ def mixDrowing_water_postprocess(preds,_mask_cv,pars ):
# %((t11-t10) * 1000,(t10-t9) * 1000,(t9-t8) * 1000,(t8-t7) * 1000,(t7-t6) * 1000,(t6-t5) * 1000,(t5-t4) * 1000 ) )
timeInfos=' findMaxWater:%.1f releJudge:%.1f'%( ms(t6,t4) ,ms(t11,t6) )
return final_output_luoshui,timeInfos #返回最终绘制的结果图、最终落水人员(坐标、类别、置信度)
def mixDrowing_water_postprocess_N(predList,pars ):
preds,_mask_cv = predList[0:2]
return mixDrowing_water_postprocess(preds,_mask_cv,pars )

+ 4
- 1
utilsK/illParkingUtils.py View File

@@ -216,7 +216,10 @@ def illParking_postprocess(pred,cvMask,pars):
return outRes,timeStr #返回最终绘制的结果图、违停车辆(坐标、类别、置信度)
def illParking_postprocess_N(predList,pars):
pred=predList[0]
cvMask=None
return illParking_postprocess(pred,cvMask,pars)
def AI_process(model, args1,path1):
'''对原图进行目标检测'''
'''输入:检测模型、配置参数、路径

+ 464
- 416
utilsK/noParkingUtils.py View File

@@ -1,416 +1,464 @@
# 耗时最短代码
import numpy as np
import cv2,time,math
import matplotlib.pyplot as plt


def get_ms(time2, time1):
return (time2-time1)*1000.0


# 计算一点到二次函数曲线的距离,二次函数的表达式为x = a*(y**2) + b*y + c
def point2QF(a, b, c, y, x): # 坐标点(y, x)
distance = abs(x - a*(y**2) - b*y - c) / math.sqrt(1 + ((2*a*y + b)**2))
return distance


# 存储所有speedRoad的contours
def storageRoad(contours, pars):
allRoadCnt = [] # 存储所有speedRoad的contours
for cnt in contours: # 道路
if len(cnt) >= 6:
rect = cv2.minAreaRect(cnt)
if rect[1][0] * rect[1][1] > pars['RoadArea']: # 过滤掉面积小于阈值的speedRoad
allRoadCnt.append(cnt)
return allRoadCnt


# 用于判断lane数量是否大于2,当lane数量大于等于2时,返回laneNumber
def storageLane(contours, pars):
laneNumber = 0
for cnt in contours:
if len(cnt) >= 6:
rect = cv2.minAreaRect(cnt)
if rect[1][0] * rect[1][1] > pars['laneArea'] and min(rect[1]) / max(rect[1]) <= pars['roundness']:
laneNumber += 1
if laneNumber > 2:
break
return laneNumber


# 将contours中顶点数大于等于6的车辆信息(合格vehicle)和顶点数小于6的车辆信息(不合格vehicle)分别保存起来
def vehicleDivide(contours, vehicleBD, normVehicle, dets, count, i, unnormVehicle, normVehicleCOOR, centerCOOR):
if len(contours) >= 6:
vehicleBD.append(contours)
normVehicle.append(dets[count])
normVehicleCOOR.append(centerCOOR)
else:
dets[int(i / 2)].append(0)
dets[int(i / 2)].append(0)
unnormVehicle.append(dets[int(i / 2)])
return vehicleBD, normVehicle, unnormVehicle, normVehicleCOOR


# 存储所有vehicle的信息
def storageVehicle(pars, imgVehicle, dets):
"""
输入
pars:字典名
imgVehicle:分割图,只包含vehicle和背景
dets:是一个list,其中存储检测得到的各vehicle的信息,即[[x0, y0, x1, y1, 车辆得分, cls], ...]
输出
dets:存储合格vehicle的信息,即[x0, y0, x1, y1, 车辆得分, cls]
vehicleBD:存储合格vehicle的contours
unnormVehicle:存储不合格vehicle的信息,即[x0, y0, x1, y1, 车辆得分, cls]
normVehicleCOOR:存储合格vehicle的中心点坐标
说明
合格vehicle:contours中的顶点数大于等于6
不合格vehicle:contours中的顶点数小于6
"""
vehicleBD = [] # 存储一副图像中vehicles的contours
normVehicle = [] # 将合格vehicle的信息存储在normVehicle中
unnormVehicle = [] # 将不合格vehicle的信息存储在unnormVehicle中
normVehicleCOOR = [] # 存储合格vehicle的中心点坐标
img = cv2.cvtColor(imgVehicle, cv2.COLOR_BGR2GRAY)
count = 0
for i in range(0, len(pars['vehicleCOOR']), 2):
y1 = int(pars['vehicleCOOR'][i][1] * pars['ZoomFactor']['y'])
y2 = int(pars['vehicleCOOR'][i + 1][1] * pars['ZoomFactor']['y'])
x1 = int(pars['vehicleCOOR'][i][0] * pars['ZoomFactor']['x'])
x2 = int(pars['vehicleCOOR'][i + 1][0] * pars['ZoomFactor']['x'])
if y1 >= 2:
y1 = y1 - 2
if y2 <= (pars['modelSize'][1] - 2):
y2 = y2 + 2
if x1 >= 2:
x1 = x1 - 2
if x2 <= (pars['modelSize'][0] - 2):
x2 = x2 + 2
centerCOOR = (int((x1 + x2) / 2), int((y1 + y2) / 2))
img1 = img[y1:y2, x1:x2]
up = np.zeros((20, (x2 - x1)), dtype='uint8')
left = np.zeros(((40 + y2 - y1), 20), dtype='uint8')
img1 = np.concatenate((up, img1), axis=0)
img1 = np.concatenate((img1, up), axis=0)
img1 = np.concatenate((left, img1), axis=1)
img2 = np.concatenate((img1, left), axis=1)
contours2, hierarchy = cv2.findContours(img2, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
if len(contours2) != 0:
if len(contours2) > 1:
vehicleArea = [] # 存储vehicle的最小外接矩形的面积
for j in range(len(contours2)):
rect = cv2.minAreaRect(contours2[j])
vehicleArea.append(rect[1][0] * rect[1][1])
maxAreaIndex = vehicleArea.index(max(vehicleArea))
maxAreaContours = contours2[maxAreaIndex]
vehicleBD, normVehicle, unnormVehicle, normVehicleCOOR = vehicleDivide(maxAreaContours, vehicleBD, normVehicle, dets, count, i, unnormVehicle, normVehicleCOOR, centerCOOR)
elif len(contours2) == 1:
vehicleBD, normVehicle, unnormVehicle, normVehicleCOOR = vehicleDivide(contours2[0], vehicleBD, normVehicle, dets, count, i, unnormVehicle, normVehicleCOOR, centerCOOR)
else:
dets[int(i / 2)].append(0)
dets[int(i / 2)].append(0)
unnormVehicle.append(dets[int(i / 2)])
count += 1
dets = normVehicle
return dets, vehicleBD, unnormVehicle, normVehicleCOOR


# 计算违停得分
def IllegalParkScore(vehicleBD, allRoadCnt, dets, laneNumber, unnormVehicle, normVehicleCOOR, a_l, b_l, c_l, a_r, b_r, c_r):
"""
对vehicle是否在speedRoad上进行判断,并计算违章得分
输出targetList 其格式为:[[cls, x0, y0, x1, y1, score, 违章得分, 违章类别], ...]
"""
if len(vehicleBD) != 0:
for i in range(len(vehicleBD)):
rect = cv2.minAreaRect(vehicleBD[i])
center = normVehicleCOOR[i] # vehicle的中心点坐标
if len(allRoadCnt) != 0 and laneNumber >= 2: # 当车道线个数至少有两条时,才计算违章得分
for j in range(len(allRoadCnt)):
# 判断车辆矩形框的中心点坐标是否在道路矩形框的范围内
flag = cv2.pointPolygonTest(allRoadCnt[j], center, False)
if flag >= 0:
dets[i].append(0) # 给违章得分占位
dets[i].append(0) # 给违章类别占位
if center[0] < predict(a_l, b_l, c_l, center[1]):
distance = point2QF(a_l, b_l, c_l, center[1], center[0])
if distance >= min(rect[1]) / 2:
dets[i][6], dets[i][7] = 1, 1
else:
dets[i][6], dets[i][7] = distance / (min(rect[1]) / 2), 1
elif center[0] > predict(a_r, b_r, c_r, center[1]):
distance = point2QF(a_r, b_r, c_r, center[1], center[0])
if distance >= min(rect[1]) / 2:
dets[i][6], dets[i][7] = 1, 1
else:
dets[i][6], dets[i][7] = distance / (min(rect[1]) / 2), 1
else:
dets[i][6], dets[i][7] = 0, 0
break
# 如果分割图像中不存在speedRoad,则无法进行违章判定,将所有车辆的违章类别设为0,即没有违章
if len(dets[i]) < 8:
dets[i].append(0) # 违章得分为0
dets[i].append(0) # 0表示没有违章
targetList = dets
if len(unnormVehicle) != 0:
for i in range(len(unnormVehicle)):
targetList.append(unnormVehicle[i]) # 将所有车辆的信息合并到一起
else:
targetList = unnormVehicle
return targetList


# 找最左侧lane时,将要删除的右侧lane的序号存储在delRightLane。找最右侧lane时,将要删除的左侧lane的序号存储在delLeftLane。
def devideLane(laneInfo, i, m, delRightLane, delLeftLane, y):
index1 = np.where(laneInfo[i][3] == y)
index1 = index1[0].tolist()
index1.sort()
x_1 = laneInfo[i][5][index1[0]][0]
index2 = np.where(laneInfo[m][3] == y)
index2 = index2[0].tolist()
index2.sort()
x_2 = laneInfo[m][5][index2[0]][0]
if x_1 < x_2:
if i not in delLeftLane:
delLeftLane.append(i) # 保留右侧lane
if m not in delRightLane:
delRightLane.append(m)
else:
if m not in delLeftLane:
delLeftLane.append(m)
if i not in delRightLane:
delRightLane.append(i)

return delRightLane, delLeftLane


# 确定最左侧和最右侧的lane簇
def detLine(contours):
"""
输入
contours:各lane的contours
输出
laneInfo:存储各lane的信息,每条lane的信息为:[contours, y坐标范围, lane序号, arr_y, y坐标范围的长度, cnt]
delRightLane:在确定最左侧lane时,其存储需要删除的lane的序号
delLeftLane:在确定最右侧lane时,其存储需要删除的lane的序号
"""
mergList = []
for i in range(len(contours)):
cnt = np.squeeze(contours[i], 1)
arr_y = cnt[:, 1]
arrList = list(set(arr_y))
cnt_y = np.sort(np.array(arrList))
mergList.append([contours[i], cnt_y, i, arr_y, len(cnt_y), cnt])
laneInfo = sorted(mergList, key=(lambda x: x[4])) # [[contours[i], cnt_y, i, arr_y, len(cnt_y)],...]
delRightLane = [] # 求最左侧lane
delLeftLane = [] # 求最右侧lane
laneInfoNew = []
for i in range(len(laneInfo)):
laneInfoNew.append([laneInfo[i][1][0], laneInfo[i][1][-1], i]) # [[y_min, y_max, i],...]
laneInfoNew = np.array(laneInfoNew)
new1 = laneInfoNew[:, np.newaxis, :].repeat(laneInfoNew.shape[0], 1)
new2 = laneInfoNew[np.newaxis, ...].repeat(laneInfoNew.shape[0], 0)
new3 = np.concatenate((new1, new2), axis=2)
y_i_min, y_i_max, y_m_min, y_m_max = new3[..., 0], new3[..., 1], new3[..., 3], new3[..., 4]
mask1 = (y_i_min >= y_m_min) & (y_i_min <= y_m_max) & (y_i_max > y_m_max)
mask2 = (y_i_max >= y_m_min) & (y_i_max <= y_m_max) & (y_i_min < y_m_min)
mask3 = (y_i_min >= y_m_min) & (y_i_max <= y_m_max)
mask4 = (y_i_min < y_m_min) & (y_i_max > y_m_max)
if len(np.nonzero(mask1)[0]) != 0:
mask1 = np.triu(mask1, k=1)
serial_i = new3[mask1][..., 2]
serial_m = new3[mask1][..., 5]
for k in range(len(serial_i)):
if (serial_m[k] not in delLeftLane) or (serial_m[k] not in delRightLane) or (serial_i[k] not in delLeftLane) or (serial_i[k] not in delRightLane):
delRightLane, delLeftLane = devideLane(laneInfo, serial_i[k], serial_m[k], delRightLane, delLeftLane, laneInfo[serial_i[k]][1][0])

if len(np.nonzero(mask2)[0]) != 0:
mask2 = np.triu(mask2, k=1)
serial_i = new3[mask2][..., 2]
serial_m = new3[mask2][..., 5]
for k in range(len(serial_i)):
if (serial_m[k] not in delLeftLane) or (serial_m[k] not in delRightLane) or (serial_i[k] not in delLeftLane) or (serial_i[k] not in delRightLane):
delRightLane, delLeftLane = devideLane(laneInfo, serial_i[k], serial_m[k], delRightLane, delLeftLane, laneInfo[serial_i[k]][1][-1])

if len(np.nonzero(mask3)[0]) != 0:
mask3 = np.triu(mask3, k=1)
serial_i = new3[mask3][..., 2]
serial_m = new3[mask3][..., 5]
for k in range(len(serial_i)):
if (serial_m[k] not in delLeftLane) or (serial_m[k] not in delRightLane) or (serial_i[k] not in delLeftLane) or (serial_i[k] not in delRightLane):
delRightLane, delLeftLane = devideLane(laneInfo, serial_i[k], serial_m[k], delRightLane, delLeftLane, laneInfo[serial_i[k]][1][0])

if len(np.nonzero(mask4)[0]) != 0:
mask4 = np.triu(mask4, k=1)
serial_i = new3[mask4][..., 2]
serial_m = new3[mask4][..., 5]
for k in range(len(serial_i)):
if (serial_m[k] not in delLeftLane) or (serial_m[k] not in delRightLane) or (serial_i[k] not in delLeftLane) or (serial_i[k] not in delRightLane):
delRightLane, delLeftLane = devideLane(laneInfo, serial_i[k], serial_m[k], delRightLane, delLeftLane, laneInfo[serial_m[k]][1][0])
return laneInfo, delRightLane, delLeftLane


# 对lane中的y值坐标进行下采样
def downSample(cnt_y):
# number = len(cnt_y) * 0.0125
# cnt_y = np.random.choice(cnt_y, size=number, replace=False)
if len(cnt_y) >= 1000:
cnt_y = cnt_y[1::80]
elif len(cnt_y) >= 900 and len(cnt_y) < 1000:
cnt_y = cnt_y[1::75]
elif len(cnt_y) >= 800 and len(cnt_y) < 900:
cnt_y = cnt_y[1::70]
elif len(cnt_y) >= 700 and len(cnt_y) < 800:
cnt_y = cnt_y[1::65]
elif len(cnt_y) >= 600 and len(cnt_y) < 700:
cnt_y = cnt_y[1::60]
elif len(cnt_y) >= 500 and len(cnt_y) < 600:
cnt_y = cnt_y[1::55]
elif len(cnt_y) >= 400 and len(cnt_y) < 500:
cnt_y = cnt_y[1::40]
elif len(cnt_y) >= 300 and len(cnt_y) < 400:
cnt_y = cnt_y[1::45]
elif len(cnt_y) >= 200 and len(cnt_y) < 300:
cnt_y = cnt_y[1::40]
elif len(cnt_y) >= 100 and len(cnt_y) < 200:
cnt_y = cnt_y[1::35]
elif len(cnt_y) >= 50 and len(cnt_y) < 100:
cnt_y = cnt_y[1::20]
elif len(cnt_y) >= 20 and len(cnt_y) < 50:
cnt_y = cnt_y[1::6]
else:
cnt_y = cnt_y[1::5]
return cnt_y


# 求最左侧lane或最右侧lane中的各点坐标
def targetCOOR(laneInfo, delLane):
"""
输入
laneInfo:存储各lane的信息,每条lane的信息为:[contours, y坐标范围, lane序号, arr_y, y坐标范围的长度, cnt]
delLane:在确定最左侧lane或最右侧lane时,其存储需要删除的lane的序号。
输出
laneCOOR:存储最左侧或最右侧lane簇中各点的坐标
"""
laneCOOR = [] # 存储lane中各点的坐标
centerSort = [] # 存储各lane按照中心点的y坐标排序后的结果
for j in range(len(laneInfo)):
if j not in delLane:
cnt = laneInfo[j][0]
rect = cv2.minAreaRect(cnt)
cnt = np.squeeze(cnt, 1)
cnt_y = laneInfo[j][1]
cnt_y = downSample(cnt_y)
centerSort.append([rect[0][1], cnt_y, laneInfo[j][3], cnt, j])
centerSort = sorted(centerSort, key=(lambda x: x[0]))
for i in range(len(centerSort)):
centerCoordinate = []
for j in range(len(centerSort[i][1])):
index = np.where(centerSort[i][2] == centerSort[i][1][j])
indexList = index[0].tolist()
indexList.sort()
x = (centerSort[i][3][indexList[0]][0] + centerSort[i][3][indexList[-1]][0]) / 2
y = (centerSort[i][3][indexList[0]][1] + centerSort[i][3][indexList[-1]][1]) / 2
centerCoordinate.append([x, y])
laneCOOR = laneCOOR + centerCoordinate
return laneCOOR


# 二次函数曲线表达式:x = a*(y**2) + b*y + c,根据图像中一点的y坐标求二次曲线中的x坐标
def predict(a, b, c, y):
x = a * (y**2) + b * y + c
return x


def mixNoParking_road_postprocess(dets, mask, pars):
"""
对于字典traffic_dict中的各个键,说明如下:
RoadArea:speedRoad的最小外接矩形的面积
vehicleCOOR:是一个列表,用于存储被检测出的vehicle的坐标(vehicle检测模型)
roundness:圆度 ,lane的长与宽的比率,作为判定是否为车道线的标准之一
laneArea:车道线的最小外接矩形的面积
ZoomFactor:图像在H和W方向上的缩放因子,其值小于1
fitOrder:多点拟合曲线的阶数
最终输出格式:[[x0, y0, x1, y1, 车辆得分, cls, 违章停车得分, 违章类别], ...]
违章类别:0表示正常车辆,1表示违章车辆
"""
det_cors = []
for bb in dets:
det_cors.append((int(bb[0]), int(bb[1])))
det_cors.append((int(bb[2]), int(bb[3])))
#print('#################line341:',det_cors)
pars['vehicleCOOR'] = det_cors
H, W = mask.shape[0:2] # mask的分辨率为360x640
scaleH = pars['modelSize'][1] / H # 自适应调整缩放比例
scaleW = pars['modelSize'][0] / W
pars['ZoomFactor'] = {'x': scaleW, 'y': scaleH}
new_hw = [int(H * scaleH), int(W * scaleW)]
mask = cv2.resize(mask, (new_hw[1], new_hw[0]))
if len(mask.shape) == 3:
mask = mask[:, :, 0]

t1 = time.time()
imgRoad = mask.copy()
imgVehicle = mask.copy()
lane_line = mask.copy()
# 将vehicle和lane过滤掉,只包含背景和speedRoad
imgRoad[imgRoad == 2] = 1
imgRoad[imgRoad == 3] = 1
# 将speedRoad和lane过滤掉,只保留vehicle和背景
imgVehicle[imgVehicle != 2] = 0
# 将speedRoad和vehicle过滤掉,只保留lane和背景
lane_line[lane_line < 3] = 0
imgRoad = cv2.cvtColor(np.uint8(imgRoad), cv2.COLOR_RGB2BGR) # 道路
imgVehicle = cv2.cvtColor(np.uint8(imgVehicle), cv2.COLOR_RGB2BGR) # 车辆
lane_line = cv2.cvtColor(np.uint8(lane_line), cv2.COLOR_RGB2BGR)

# 对车道线进行膨胀操作
# kernel = np.ones((3, 3), np.uint8) # 膨胀范围
# lane_line = cv2.dilate(lane_line, kernel, iterations=2) # 迭代次数为2

t2 = time.time()
img1 = cv2.cvtColor(imgRoad, cv2.COLOR_BGR2GRAY)
roadContours, hierarchy = cv2.findContours(img1, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
t3 = time.time()
# 存储所有speedRoad的信息
allRoadCnt = storageRoad(roadContours, pars)
t4 = time.time()
img3 = cv2.cvtColor(lane_line, cv2.COLOR_BGR2GRAY)
laneContours, hierarchy = cv2.findContours(img3, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
t5 = time.time()
laneInfo, delRightLane, delLeftLane = detLine(laneContours)
t6 = time.time()
# 存储所有lane的信息
laneNumber = storageLane(laneContours, pars)
t7 = time.time()
# 存储所有vehicle的信息
dets, vehicleBD, unnormVehicle, normVehicleCOOR = storageVehicle(pars, imgVehicle, dets)
t8 = time.time()
leftLaneCOOR = targetCOOR(laneInfo, delRightLane)
rightLaneCOOR = targetCOOR(laneInfo, delLeftLane)
rightLaneCOOR = np.array(rightLaneCOOR)
rightX = rightLaneCOOR[:, 0]
rightY = rightLaneCOOR[:, 1]
leftLaneCOOR = np.array(leftLaneCOOR)
leftX = leftLaneCOOR[:, 0]
leftY = leftLaneCOOR[:, 1]
# a_r,b_r,c_r分别是:最右侧车道线簇拟合的二次函数的二次项系数一次项系数,和常数项
a_r, b_r, c_r = np.polyfit(rightY, rightX, pars['fitOrder'])[0], np.polyfit(rightY, rightX, pars['fitOrder'])[1], np.polyfit(rightY, rightX, pars['fitOrder'])[2]
# a_l,b_l,c_l分别是:最左侧车道线簇拟合的二次函数的二次项系数,一次项系数,和常数项
a_l, b_l, c_l = np.polyfit(leftY, leftX, pars['fitOrder'])[0], np.polyfit(leftY, leftX, pars['fitOrder'])[1], np.polyfit(leftY, leftX, pars['fitOrder'])[2]

"""以下四行代码用于在后处理函数外画图"""
finalLane = []
abc = [a_l, b_l, c_l, a_r, b_r, c_r] # abc中存储的是最左侧和最右侧二次函数的各项系数
finalLane.append(rightLaneCOOR)
finalLane.append(leftLaneCOOR)
t9 = time.time()
# 计算违停得分
targetList = IllegalParkScore(vehicleBD, allRoadCnt, dets, laneNumber, unnormVehicle, normVehicleCOOR, a_l, b_l, c_l, a_r, b_r, c_r)
t10 = time.time()
time_infos = 'postTime:%.2f(分割时间:%.2f, findContours:%.2f, ruleJudge:%.2f, storageRoad:%.2f, detLane:%.2f, storageLane:%.2f, storageVehicle:%.2f, fitLine:%.2f, IllegalParkScore:%.2f)' % (
get_ms(t10, t1), get_ms(t2, t1), get_ms(t3, t2), get_ms(t8, t3), get_ms(t4, t3), get_ms(t6, t5), get_ms(t7, t6), get_ms(t8, t7), get_ms(t9, t8), get_ms(t10, t9))
#, finalLane, lane_line, abc
targetList = [ [ *b[0:4],b[6] if b[6]>0 else b[4], b[7] ] for b in targetList ]
#print('---------line415:',targetList)
return targetList, time_infos
# 耗时最短代码
import numpy as np
import cv2,time,math
import matplotlib.pyplot as plt
def get_ms(time2, time1):
return (time2-time1)*1000.0
# 计算一点到二次函数曲线的距离,二次函数的表达式为x = a*(y**2) + b*y + c
def point2QF(a, b, c, y, x): # 坐标点(y, x)
distance = abs(x - a*(y**2) - b*y - c) / math.sqrt(1 + ((2*a*y + b)**2))
return distance
# 存储所有speedRoad的contours
def storageRoad(contours, pars):
allRoadCnt = [] # 存储所有speedRoad的contours
for cnt in contours: # 道路
if len(cnt) >= 6:
rect = cv2.minAreaRect(cnt)
if rect[1][0] * rect[1][1] > pars['RoadArea']: # 过滤掉面积小于阈值的speedRoad
allRoadCnt.append(cnt)
return allRoadCnt
# 返回符合标准的lane的个数及contours
def storageLane(contours, pars):
"""
contours:lane分割后的原始contours
newLaneContours:符合标准的lane的contours
laneNumber:符合标准的lane的个数
符合标准的lane定义如下:
(1)contours中的坐标点个数不小于6
(2)lane最小外接矩形的面积大于阈值laneArea
(3)lane最小外接矩形的最短边与最长边的比值小于等于阈值roundness
"""
laneNumber = 0
newLaneContours = ()
for cnt in contours:
if len(cnt) >= 6:
rect = cv2.minAreaRect(cnt)
if rect[1][0] * rect[1][1] > pars['laneArea'] and min(rect[1]) / max(rect[1]) <= pars['roundness']:
laneNumber += 1
newLaneContours = newLaneContours + (cnt, )
return laneNumber, newLaneContours
# 将contours中顶点数大于等于6的车辆信息(合格vehicle)和顶点数小于6的车辆信息(不合格vehicle)分别保存起来
def vehicleDivide(contours, vehicleBD, normVehicle, dets, count, i, unnormVehicle, normVehicleCOOR, centerCOOR):
if len(contours) >= 6:
vehicleBD.append(contours)
normVehicle.append(dets[count])
normVehicleCOOR.append(centerCOOR)
else:
dets[int(i / 2)].append(0)
dets[int(i / 2)].append(0)
unnormVehicle.append(dets[int(i / 2)])
return vehicleBD, normVehicle, unnormVehicle, normVehicleCOOR
# 存储所有vehicle的信息
def storageVehicle(pars, imgVehicle, dets):
"""
输入
pars:字典名
imgVehicle:分割图,只包含vehicle和背景
dets:是一个list,其中存储检测得到的各vehicle的信息,即[[x0, y0, x1, y1, 车辆得分, cls], ...]
输出
dets:存储合格vehicle的信息,即[x0, y0, x1, y1, 车辆得分, cls]
vehicleBD:存储合格vehicle的contours
unnormVehicle:存储不合格vehicle的信息,即[x0, y0, x1, y1, 车辆得分, cls]
normVehicleCOOR:存储合格vehicle的中心点坐标
说明
合格vehicle:contours中的顶点数大于等于6
不合格vehicle:contours中的顶点数小于6
"""
vehicleBD = [] # 存储一副图像中vehicles的contours
normVehicle = [] # 将合格vehicle的信息存储在normVehicle中
unnormVehicle = [] # 将不合格vehicle的信息存储在unnormVehicle中
normVehicleCOOR = [] # 存储合格vehicle的中心点坐标
img = cv2.cvtColor(imgVehicle, cv2.COLOR_BGR2GRAY)
count = 0
for i in range(0, len(pars['vehicleCOOR']), 2):
y1 = int(pars['vehicleCOOR'][i][1] * pars['ZoomFactor']['y'])
y2 = int(pars['vehicleCOOR'][i + 1][1] * pars['ZoomFactor']['y'])
x1 = int(pars['vehicleCOOR'][i][0] * pars['ZoomFactor']['x'])
x2 = int(pars['vehicleCOOR'][i + 1][0] * pars['ZoomFactor']['x'])
if y1 >= 2:
y1 = y1 - 2
if y2 <= (pars['modelSize'][1] - 2):
y2 = y2 + 2
if x1 >= 2:
x1 = x1 - 2
if x2 <= (pars['modelSize'][0] - 2):
x2 = x2 + 2
centerCOOR = (int((x1 + x2) / 2), int((y1 + y2) / 2))
img1 = img[y1:y2, x1:x2]
up = np.zeros((20, (x2 - x1)), dtype='uint8')
left = np.zeros(((40 + y2 - y1), 20), dtype='uint8')
img1 = np.concatenate((up, img1), axis=0)
img1 = np.concatenate((img1, up), axis=0)
img1 = np.concatenate((left, img1), axis=1)
img2 = np.concatenate((img1, left), axis=1)
contours2, hierarchy = cv2.findContours(img2, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
if len(contours2) != 0:
if len(contours2) > 1:
vehicleArea = [] # 存储vehicle的最小外接矩形的面积
for j in range(len(contours2)):
rect = cv2.minAreaRect(contours2[j])
vehicleArea.append(rect[1][0] * rect[1][1])
maxAreaIndex = vehicleArea.index(max(vehicleArea))
maxAreaContours = contours2[maxAreaIndex]
vehicleBD, normVehicle, unnormVehicle, normVehicleCOOR = vehicleDivide(maxAreaContours, vehicleBD, normVehicle, dets, count, i, unnormVehicle, normVehicleCOOR, centerCOOR)
elif len(contours2) == 1:
vehicleBD, normVehicle, unnormVehicle, normVehicleCOOR = vehicleDivide(contours2[0], vehicleBD, normVehicle, dets, count, i, unnormVehicle, normVehicleCOOR, centerCOOR)
else:
dets[int(i / 2)].append(0)
dets[int(i / 2)].append(0)
unnormVehicle.append(dets[int(i / 2)])
count += 1
dets = normVehicle
return dets, vehicleBD, unnormVehicle, normVehicleCOOR
# 计算违停得分
def IllegalParkScore1(vehicleBD, allRoadCnt, dets, unnormVehicle, normVehicleCOOR, a_l, b_l, c_l, a_r, b_r, c_r):
"""
对vehicle是否在speedRoad上进行判断,并计算违章得分
输出targetList 其格式为:[[cls, x0, y0, x1, y1, score, 违章得分, 违章类别], ...]
"""
if len(vehicleBD) != 0:
for i in range(len(vehicleBD)):
rect = cv2.minAreaRect(vehicleBD[i])
center = normVehicleCOOR[i] # vehicle的中心点坐标
if len(allRoadCnt) != 0: # 当车道线个数至少有两条时,才计算违章得分
for j in range(len(allRoadCnt)):
# 判断车辆矩形框的中心点坐标是否在道路矩形框的范围内
flag = cv2.pointPolygonTest(allRoadCnt[j], center, False)
if flag >= 0:
dets[i].append(0) # 给违章得分占位
dets[i].append(0) # 给违章类别占位
if center[0] < predict(a_l, b_l, c_l, center[1]):
distance = point2QF(a_l, b_l, c_l, center[1], center[0])
if distance >= min(rect[1]) / 2:
dets[i][6], dets[i][7] = 1, 1
else:
dets[i][6], dets[i][7] = distance / (min(rect[1]) / 2), 1
elif center[0] > predict(a_r, b_r, c_r, center[1]):
distance = point2QF(a_r, b_r, c_r, center[1], center[0])
if distance >= min(rect[1]) / 2:
dets[i][6], dets[i][7] = 1, 1
else:
dets[i][6], dets[i][7] = distance / (min(rect[1]) / 2), 1
else:
dets[i][6], dets[i][7] = 0, 0
break
# 如果分割图像中不存在speedRoad,则无法进行违章判定,将所有车辆的违章类别设为0,即没有违章
if len(dets[i]) < 8:
dets[i].append(0) # 违章得分为0
dets[i].append(0) # 0表示没有违章
targetList = dets
if len(unnormVehicle) != 0:
for i in range(len(unnormVehicle)):
targetList.append(unnormVehicle[i]) # 将所有车辆的信息合并到一起
else:
targetList = unnormVehicle
return targetList
# 计算违停得分
def IllegalParkScore2(vehicleBD, dets, unnormVehicle):
"""
计算违章得分
输出targetList 其格式为:[[cls, x0, y0, x1, y1, score, 违章得分, 违章类别], ...]
"""
if len(vehicleBD) != 0:
for i in range(len(vehicleBD)):
if len(dets[i]) < 8:
dets[i].append(0) # 违章得分为0
dets[i].append(0) # 0表示没有违章
targetList = dets
if len(unnormVehicle) != 0:
for i in range(len(unnormVehicle)):
targetList.append(unnormVehicle[i]) # 将所有车辆的信息合并到一起
else:
targetList = unnormVehicle
return targetList
# 找最左侧lane时,将要删除的右侧lane的序号存储在delRightLane。找最右侧lane时,将要删除的左侧lane的序号存储在delLeftLane。
def devideLane(laneInfo, i, m, delRightLane, delLeftLane, y):
index1 = np.where(laneInfo[i][3] == y)
index1 = index1[0].tolist()
index1.sort()
x_1 = laneInfo[i][5][index1[0]][0]
index2 = np.where(laneInfo[m][3] == y)
index2 = index2[0].tolist()
index2.sort()
x_2 = laneInfo[m][5][index2[0]][0]
if x_1 < x_2:
if i not in delLeftLane:
delLeftLane.append(i) # 保留右侧lane
if m not in delRightLane:
delRightLane.append(m)
else:
if m not in delLeftLane:
delLeftLane.append(m)
if i not in delRightLane:
delRightLane.append(i)
return delRightLane, delLeftLane
# 确定最左侧和最右侧的lane簇
def detLine(contours):
"""
输入
contours:各lane的contours
输出
laneInfo:存储各lane的信息,每条lane的信息为:[contours, y坐标范围, lane序号, arr_y, y坐标范围的长度, cnt]
delRightLane:在确定最左侧lane时,其存储需要删除的lane的序号
delLeftLane:在确定最右侧lane时,其存储需要删除的lane的序号
"""
mergList = []
for i in range(len(contours)):
cnt = np.squeeze(contours[i], 1)
arr_y = cnt[:, 1]
arrList = list(set(arr_y))
cnt_y = np.sort(np.array(arrList))
mergList.append([contours[i], cnt_y, i, arr_y, len(cnt_y), cnt])
laneInfo = sorted(mergList, key=(lambda x: x[4])) # [[contours[i], cnt_y, i, arr_y, len(cnt_y)],...]
delRightLane = [] # 求最左侧lane
delLeftLane = [] # 求最右侧lane
laneInfoNew = []
for i in range(len(laneInfo)):
laneInfoNew.append([laneInfo[i][1][0], laneInfo[i][1][-1], i]) # [[y_min, y_max, i],...]
laneInfoNew = np.array(laneInfoNew)
new1 = laneInfoNew[:, np.newaxis, :].repeat(laneInfoNew.shape[0], 1)
new2 = laneInfoNew[np.newaxis, ...].repeat(laneInfoNew.shape[0], 0)
new3 = np.concatenate((new1, new2), axis=2)
y_i_min, y_i_max, y_m_min, y_m_max = new3[..., 0], new3[..., 1], new3[..., 3], new3[..., 4]
mask1 = (y_i_min >= y_m_min) & (y_i_min <= y_m_max) & (y_i_max > y_m_max)
mask2 = (y_i_max >= y_m_min) & (y_i_max <= y_m_max) & (y_i_min < y_m_min)
mask3 = (y_i_min >= y_m_min) & (y_i_max <= y_m_max)
mask4 = (y_i_min < y_m_min) & (y_i_max > y_m_max)
if len(np.nonzero(mask1)[0]) != 0:
mask1 = np.triu(mask1, k=1)
serial_i = new3[mask1][..., 2]
serial_m = new3[mask1][..., 5]
for k in range(len(serial_i)):
if (serial_m[k] not in delLeftLane) or (serial_m[k] not in delRightLane) or (serial_i[k] not in delLeftLane) or (serial_i[k] not in delRightLane):
delRightLane, delLeftLane = devideLane(laneInfo, serial_i[k], serial_m[k], delRightLane, delLeftLane, laneInfo[serial_i[k]][1][0])
if len(np.nonzero(mask2)[0]) != 0:
mask2 = np.triu(mask2, k=1)
serial_i = new3[mask2][..., 2]
serial_m = new3[mask2][..., 5]
for k in range(len(serial_i)):
if (serial_m[k] not in delLeftLane) or (serial_m[k] not in delRightLane) or (serial_i[k] not in delLeftLane) or (serial_i[k] not in delRightLane):
delRightLane, delLeftLane = devideLane(laneInfo, serial_i[k], serial_m[k], delRightLane, delLeftLane, laneInfo[serial_i[k]][1][-1])
if len(np.nonzero(mask3)[0]) != 0:
mask3 = np.triu(mask3, k=1)
serial_i = new3[mask3][..., 2]
serial_m = new3[mask3][..., 5]
for k in range(len(serial_i)):
if (serial_m[k] not in delLeftLane) or (serial_m[k] not in delRightLane) or (serial_i[k] not in delLeftLane) or (serial_i[k] not in delRightLane):
delRightLane, delLeftLane = devideLane(laneInfo, serial_i[k], serial_m[k], delRightLane, delLeftLane, laneInfo[serial_i[k]][1][0])
if len(np.nonzero(mask4)[0]) != 0:
mask4 = np.triu(mask4, k=1)
serial_i = new3[mask4][..., 2]
serial_m = new3[mask4][..., 5]
for k in range(len(serial_i)):
if (serial_m[k] not in delLeftLane) or (serial_m[k] not in delRightLane) or (serial_i[k] not in delLeftLane) or (serial_i[k] not in delRightLane):
delRightLane, delLeftLane = devideLane(laneInfo, serial_i[k], serial_m[k], delRightLane, delLeftLane, laneInfo[serial_m[k]][1][0])
return laneInfo, delRightLane, delLeftLane
# 对lane中的y值坐标进行下采样
def downSample(cnt_y):
# number = len(cnt_y) * 0.0125
# cnt_y = np.random.choice(cnt_y, size=number, replace=False)
if len(cnt_y) >= 1000:
cnt_y = cnt_y[1::80]
elif len(cnt_y) >= 900 and len(cnt_y) < 1000:
cnt_y = cnt_y[1::75]
elif len(cnt_y) >= 800 and len(cnt_y) < 900:
cnt_y = cnt_y[1::70]
elif len(cnt_y) >= 700 and len(cnt_y) < 800:
cnt_y = cnt_y[1::65]
elif len(cnt_y) >= 600 and len(cnt_y) < 700:
cnt_y = cnt_y[1::60]
elif len(cnt_y) >= 500 and len(cnt_y) < 600:
cnt_y = cnt_y[1::55]
elif len(cnt_y) >= 400 and len(cnt_y) < 500:
cnt_y = cnt_y[1::40]
elif len(cnt_y) >= 300 and len(cnt_y) < 400:
cnt_y = cnt_y[1::45]
elif len(cnt_y) >= 200 and len(cnt_y) < 300:
cnt_y = cnt_y[1::40]
elif len(cnt_y) >= 100 and len(cnt_y) < 200:
cnt_y = cnt_y[1::35]
elif len(cnt_y) >= 50 and len(cnt_y) < 100:
cnt_y = cnt_y[1::20]
elif len(cnt_y) >= 20 and len(cnt_y) < 50:
cnt_y = cnt_y[1::6]
else:
cnt_y = cnt_y[1::5]
return cnt_y
# 求最左侧lane或最右侧lane中的各点坐标
def targetCOOR(laneInfo, delLane):
"""
输入
laneInfo:存储各lane的信息,每条lane的信息为:[contours, y坐标范围, lane序号, arr_y, y坐标范围的长度, cnt]
delLane:在确定最左侧lane或最右侧lane时,其存储需要删除的lane的序号。
输出
laneCOOR:存储最左侧或最右侧lane簇中各点的坐标
"""
laneCOOR = [] # 存储lane中各点的坐标
centerSort = [] # 存储各lane按照中心点的y坐标排序后的结果
for j in range(len(laneInfo)):
if j not in delLane:
cnt = laneInfo[j][0]
rect = cv2.minAreaRect(cnt)
cnt = np.squeeze(cnt, 1)
cnt_y = laneInfo[j][1]
cnt_y = downSample(cnt_y)
centerSort.append([rect[0][1], cnt_y, laneInfo[j][3], cnt, j])
centerSort = sorted(centerSort, key=(lambda x: x[0]))
for i in range(len(centerSort)):
centerCoordinate = []
for j in range(len(centerSort[i][1])):
index = np.where(centerSort[i][2] == centerSort[i][1][j])
indexList = index[0].tolist()
indexList.sort()
x = (centerSort[i][3][indexList[0]][0] + centerSort[i][3][indexList[-1]][0]) / 2
y = (centerSort[i][3][indexList[0]][1] + centerSort[i][3][indexList[-1]][1]) / 2
centerCoordinate.append([x, y])
laneCOOR = laneCOOR + centerCoordinate
return laneCOOR
# 二次函数曲线表达式:x = a*(y**2) + b*y + c,根据图像中一点的y坐标求二次曲线中的x坐标
def predict(a, b, c, y):
x = a * (y**2) + b * y + c
return x
def mixNoParking_road_postprocess(dets, mask, pars):
"""
对于字典traffic_dict中的各个键,说明如下:
RoadArea:speedRoad的最小外接矩形的面积
vehicleCOOR:是一个列表,用于存储被检测出的vehicle的坐标(vehicle检测模型)
roundness:圆度 ,lane的长与宽的比率,作为判定是否为车道线的标准之一
laneArea:车道线的最小外接矩形的面积
ZoomFactor:图像在H和W方向上的缩放因子,其值小于1
fitOrder:多点拟合曲线的阶数
最终输出格式:[[x0, y0, x1, y1, 车辆得分, cls, 违章停车得分, 违章类别], ...]
违章类别:0表示正常车辆,1表示违章车辆
"""
det_cors = []
for bb in dets:
det_cors.append((int(bb[0]), int(bb[1])))
det_cors.append((int(bb[2]), int(bb[3])))
print('###line341:', det_cors)
pars['vehicleCOOR'] = det_cors
H, W = mask.shape[0:2] # mask的分辨率为360x640
scaleH = pars['modelSize'][1] / H # 自适应调整缩放比例
scaleW = pars['modelSize'][0] / W
pars['ZoomFactor'] = {'x': scaleW, 'y': scaleH}
new_hw = [int(H * scaleH), int(W * scaleW)]
mask = cv2.resize(mask, (new_hw[1], new_hw[0]))
if len(mask.shape) == 3:
mask = mask[:, :, 0]
t1 = time.time()
imgRoad = mask.copy()
imgVehicle = mask.copy()
lane_line = mask.copy()
# 将vehicle和lane过滤掉,只包含背景和speedRoad
imgRoad[imgRoad == 2] = 1
imgRoad[imgRoad == 3] = 1
# 将speedRoad和lane过滤掉,只保留vehicle和背景
imgVehicle[imgVehicle != 2] = 0
# 将speedRoad和vehicle过滤掉,只保留lane和背景
lane_line[lane_line < 3] = 0
imgRoad = cv2.cvtColor(np.uint8(imgRoad), cv2.COLOR_RGB2BGR) # 道路
imgVehicle = cv2.cvtColor(np.uint8(imgVehicle), cv2.COLOR_RGB2BGR) # 车辆
lane_line = cv2.cvtColor(np.uint8(lane_line), cv2.COLOR_RGB2BGR)
# 对车道线进行膨胀操作
# kernel = np.ones((3, 3), np.uint8) # 膨胀范围
# lane_line = cv2.dilate(lane_line, kernel, iterations=2) # 迭代次数为2
t2 = time.time()
img1 = cv2.cvtColor(imgRoad, cv2.COLOR_BGR2GRAY)
roadContours, hierarchy = cv2.findContours(img1, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
t3 = time.time()
# 存储所有speedRoad的信息
allRoadCnt = storageRoad(roadContours, pars)
t4 = time.time()
img3 = cv2.cvtColor(lane_line, cv2.COLOR_BGR2GRAY)
laneContours, hierarchy = cv2.findContours(img3, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
# 存储所有lane的信息
laneNumber, newLaneContours = storageLane(laneContours, pars)
t5 = time.time()
if laneNumber >= 2:
laneInfo, delRightLane, delLeftLane = detLine(newLaneContours)
t6 = time.time()
# 存储所有vehicle的信息
dets, vehicleBD, unnormVehicle, normVehicleCOOR = storageVehicle(pars, imgVehicle, dets)
t7 = time.time()
leftLaneCOOR = targetCOOR(laneInfo, delRightLane)
rightLaneCOOR = targetCOOR(laneInfo, delLeftLane)
rightLaneCOOR = np.array(rightLaneCOOR)
rightX = rightLaneCOOR[:, 0]
rightY = rightLaneCOOR[:, 1]
leftLaneCOOR = np.array(leftLaneCOOR)
leftX = leftLaneCOOR[:, 0]
leftY = leftLaneCOOR[:, 1]
# a_r,b_r,c_r分别是:最右侧车道线簇拟合的二次函数的二次项系数一次项系数,和常数项
a_r, b_r, c_r = np.polyfit(rightY, rightX, pars['fitOrder'])[0], np.polyfit(rightY, rightX, pars['fitOrder'])[1], np.polyfit(rightY, rightX, pars['fitOrder'])[2]
# a_l,b_l,c_l分别是:最左侧车道线簇拟合的二次函数的二次项系数,一次项系数,和常数项
a_l, b_l, c_l = np.polyfit(leftY, leftX, pars['fitOrder'])[0], np.polyfit(leftY, leftX, pars['fitOrder'])[1], np.polyfit(leftY, leftX, pars['fitOrder'])[2]
# """以下四行代码用于在后处理函数外画图"""
# finalLane = []
# abc = [a_l, b_l, c_l, a_r, b_r, c_r] # abc中存储的是最左侧和最右侧二次函数的各项系数
# finalLane.append(rightLaneCOOR)
# finalLane.append(leftLaneCOOR)
# 计算违停得分
t8 = time.time()
targetList = IllegalParkScore1(vehicleBD, allRoadCnt, dets, unnormVehicle, normVehicleCOOR, a_l, b_l, c_l, a_r, b_r, c_r)
t9 = time.time()
time_infos = 'postTime:%.2f(分割时间:%.2f, findContours:%.2f, ruleJudge:%.2f, storageRoad:%.2f, detLane:%.2f, storageLane:%.2f, storageVehicle:%.2f, fitLine:%.2f, IllegalParkScore1:%.2f)' % (
get_ms(t9, t1), get_ms(t2, t1), get_ms(t3, t2), get_ms(t9, t3), get_ms(t4, t3), get_ms(t6, t5),
get_ms(t5, t4), get_ms(t7, t6), get_ms(t8, t7), get_ms(t9, t8))
# print('####line445:', targetList)
# return targetList, time_infos, finalLane, lane_line, abc
targetList = [ [ *b[0:4],b[6] if b[6]>0 else b[4], b[7] ] for b in targetList ]
return targetList, time_infos
else:
dets, vehicleBD, unnormVehicle, normVehicleCOOR = storageVehicle(pars, imgVehicle, dets)
t6 = time.time()
targetList = IllegalParkScore2(vehicleBD, dets, unnormVehicle)
t7 = time.time()
time_infos = 'postTime:%.2f(分割时间:%.2f, findContours:%.2f, ruleJudge:%.2f, storageRoad:%.2f, storageLane:%.2f, storageVehicle:%.2f, IllegalParkScore2:%.2f)' % (
get_ms(t7, t1), get_ms(t2, t1), get_ms(t3, t2), get_ms(t7, t3), get_ms(t4, t3), get_ms(t5, t4), get_ms(t6, t5), get_ms(t7, t6))
# print('####line456:', targetList)
targetList = [ [ *b[0:4],b[6] if b[6]>0 else b[4], b[7] ] for b in targetList ]
return targetList, time_infos
def mixNoParking_road_postprocess_N(predList, pars):
dets, mask =predList[0:2]
return mixNoParking_road_postprocess(dets, mask, pars)

+ 3
- 1
utilsK/queRiver.py View File

@@ -270,7 +270,9 @@ def riverDetSegMixProcess(preds,water,pars={'slopeIndex':list(range(20)),'riverI
time1=time.time()
timeInfos = 'all: %.1f '%( (time1-time0) )
return det ,timeInfos
def riverDetSegMixProcess_N(predList,pars={'slopeIndex':list(range(20)),'riverIou':0.1}):
preds, water = predList[0:2]
return riverDetSegMixProcess(preds,water,pars=pars)
def getDetectionsFromPreds(pred,img,im0,conf_thres=0.2,iou_thres=0.45,ovlap_thres=0.6,padInfos=None):
'''

BIN
weights/conf/cityMangement3/dmpr.pth View File


+ 4
- 0
weights/conf/cityMangement3/labelnames.json View File

@@ -0,0 +1,4 @@
{
"labelnames":["车辆","垃圾","商贩","违停"],
"labelIndexs":["SL01","SL02","SL03","SL04"]
}

+ 7
- 0
weights/conf/cityMangement3/para.json View File

@@ -0,0 +1,7 @@
{


"post_process":{ "name":"post_process","conf_thres":0.25,"score_byClass":{"0":0.88,"1":0.3,"2":0.3,"3":0.3 } ,"iou_thres":0.45,"classes":5,"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]] }


}

BIN
weights/conf/cityMangement3/stdc_360X640.pth View File


BIN
weights/conf/cityMangement3/yolov5.pt View File


+ 97
- 0
yolov5.py View File

@@ -0,0 +1,97 @@
from models.experimental import attempt_load
import tensorrt as trt
import sys
from segutils.trtUtils import yolov5Trtforward
from utilsK.queRiver import getDetectionsFromPreds,img_pad
from utils.datasets import letterbox
import numpy as np
import torch,time
def score_filter_byClass(pdetections,score_para_2nd):
ret=[]
for det in pdetections:
score,cls = det[4],det[5]
if int(cls) in score_para_2nd.keys():
score_th = score_para_2nd[int(cls)]
elif str(int(cls)) in score_para_2nd.keys():
score_th = score_para_2nd[str(int(cls))]
else:
score_th = 0.7
if score > score_th:
ret.append(det)
return ret
class yolov5Model(object):
def __init__(self, weights=None,par={}):
self.par = par
self.device = par['device']
self.half =par['half']
if weights.endswith('.engine'):
self. infer_type ='trt'
elif weights.endswith('.pth') or weights.endswith('.pt') :
self. infer_type ='pth'
else:
print('#########ERROR:',weights,': no registered inference type, exit')
sys.exit(0)
if self.infer_type=='trt':
logger = trt.Logger(trt.Logger.ERROR)
with open(weights, "rb") as f, trt.Runtime(logger) as runtime:
self.model=runtime.deserialize_cuda_engine(f.read())# 输入trt本地文件,返回ICudaEngine对象
#print('####load TRT model :%s'%(weights))
elif self.infer_type=='pth':
self.model = attempt_load(weights, map_location=self.device) # load FP32 model
if self.half: self.model.half()
if 'score_byClass' in par.keys(): self.score_byClass = par['score_byClass']
else: self.score_byClass = None
print('#########加载模型:',weights,' 类型:',self.infer_type)
def eval(self,image):
t0=time.time()
img = self.preprocess_image(image)
t1=time.time()
if self.infer_type=='trt':
pred = yolov5Trtforward(self.model,img)
else:
pred = self.model(img,augment=False)[0]
t2=time.time()
if 'ovlap_thres_crossCategory' in self.par.keys():
ovlap_thres = self.par['ovlap_thres_crossCategory']
else:
ovlap_thres = None
p_result, timeOut = getDetectionsFromPreds(pred,img,image,conf_thres=self.par['conf_thres'],iou_thres=self.par['iou_thres'],ovlap_thres=ovlap_thres,padInfos=self.padInfos)
if self.score_byClass:
p_result[2] = score_filter_byClass(p_result[2],self.score_byClass)
t3=time.time()
timeOut = 'yolov5 :%.1f (pre-process:%.1f, inference:%.1f, post-process:%.1f) '%( self.get_ms(t3,t0) , self.get_ms(t1,t0) , self.get_ms(t2,t1) , self.get_ms(t3,t2) )
return p_result[2], timeOut
def get_ms(self,t1,t0):
return (t1-t0)*1000.0
def preprocess_image(self,image):
if self.infer_type=='trt':
img, padInfos = img_pad( image , size=(640,640,3)) ;img = [img]
self.padInfos =padInfos
else:
img = [letterbox(x, 640, auto=True, stride=32)[0] for x in [image]];
self.padInfos=None
# Stack
img = np.stack(img, 0)
# Convert
img = img[:, :, :, ::-1].transpose(0, 3, 1, 2) # BGR to RGB, to bsx3x416x416
img = np.ascontiguousarray(img)
img = torch.from_numpy(img).to(self.device)
img = img.half() if self.half else img.float() # uint8 to fp16/32
img /= 255.0
return img

Loading…
Cancel
Save