970 lines
43 KiB
Python
970 lines
43 KiB
Python
import cv2,os,time,json
|
||
from models.experimental import attempt_load
|
||
from segutils.segmodel import SegModel,get_largest_contours
|
||
from segutils.trtUtils import segtrtEval,yolov5Trtforward,OcrTrtForward
|
||
from segutils.trafficUtils import tracfficAccidentMixFunction
|
||
|
||
from utils.torch_utils import select_device
|
||
from utilsK.queRiver import get_labelnames, img_pad, getDetections, getDetectionsFromPreds, scale_back
|
||
from utilsK.jkmUtils import pre_process, post_process, get_return_data
|
||
from trackUtils.sort import moving_average_wang
|
||
|
||
from utils.datasets import letterbox
|
||
from utils.general import non_max_suppression, scale_coords,xyxy2xywh,overlap_box_suppression
|
||
from utils.plots import draw_painting_joint,get_label_arrays
|
||
import numpy as np
|
||
import torch
|
||
import math
|
||
from PIL import Image
|
||
import torch.nn.functional as F
|
||
from copy import deepcopy
|
||
from scipy import interpolate
|
||
import glob
|
||
from loguru import logger
|
||
|
||
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 xywh2xy(box,iW=None,iH=None):
|
||
xc,yc,w,h = box[0:4]
|
||
x0 =max(0, xc-w/2.0)
|
||
x1 =min(1, xc+w/2.0)
|
||
y0=max(0, yc-h/2.0)
|
||
y1=min(1,yc+h/2.0)
|
||
if iW: x0,x1 = x0*iW,x1*iW
|
||
if iH: y0,y1 = y0*iH,y1*iH
|
||
return [x0,y0,x1,y1]
|
||
|
||
def get_ms(t2,t1):
|
||
return (t2-t1)*1000.0
|
||
def get_postProcess_para(parfile):
|
||
with open(parfile) as fp:
|
||
par = json.load(fp)
|
||
assert 'post_process' in par.keys(), ' parfile has not key word:post_process'
|
||
parPost=par['post_process']
|
||
|
||
return parPost["conf_thres"],parPost["iou_thres"],parPost["classes"],parPost["rainbows"]
|
||
def get_postProcess_para_dic(parfile):
|
||
with open(parfile) as fp:
|
||
par = json.load(fp)
|
||
parPost=par['post_process']
|
||
return parPost
|
||
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.25
|
||
if score > score_th:
|
||
ret.append(det)
|
||
return ret
|
||
# 按类过滤
|
||
def filter_byClass(pdetections, fiterList):
|
||
ret = []
|
||
for det in pdetections:
|
||
score, cls = det[4], det[5]
|
||
if int(cls) in fiterList:
|
||
continue
|
||
elif str(int(cls)) in fiterList:
|
||
continue
|
||
else:
|
||
ret.append(det)
|
||
|
||
return ret
|
||
|
||
# 对ocr识别车牌格式化处理
|
||
def plat_format(ocr):
|
||
carDct = ['黑','吉','辽','冀','晋','陕','甘','青','鲁','苏','浙','皖','闽','赣','豫','鄂',\
|
||
'湘','粤','琼','川','贵','云','蒙','藏','宁','新','桂','京','津','沪','渝','使','领']
|
||
label = ocr[0]
|
||
# print(label)
|
||
label = list(filter(lambda x: (ord(x) > 19968 and ord(x) < 63865) or (ord(x) > 96 and ord(x) < 123)
|
||
or (ord(x) > 47 and ord(x) < 58) or (ord(x) in [33, 73, 65281]), label))
|
||
def spt(x):
|
||
if x in ['I', 'i', '!', '！']:
|
||
return '1'
|
||
else:
|
||
return x
|
||
|
||
label = list(map(spt, label))
|
||
if len(label) < 7 or len(label) >8:
|
||
return None
|
||
if not label[0] in carDct:
|
||
return None
|
||
|
||
label.insert(2, '・')
|
||
label = ' '.join(label)
|
||
# label = label.split('I','1').split('!','1').split('i','1').split('！','1')
|
||
# label = label.split('I','1').split('!','1').split('i','1').split('！','1
|
||
|
||
return label.upper()
|
||
|
||
def post_process_det(pred,padInfos,img,im0s,conf_thres,iou_thres,label_arraylist,rainbows,font,score_byClass,fiterList,ovlap_thres=None):
|
||
time0 = time.time()
|
||
pred = non_max_suppression(pred, conf_thres, iou_thres, classes=None, agnostic=False)
|
||
if ovlap_thres:
|
||
pred = overlap_box_suppression(pred, ovlap_thres)
|
||
time1 = time.time()
|
||
det = pred[0] ###一次检测一张图片
|
||
det_xywh = [];
|
||
im0 = im0s.copy()
|
||
#im0 = im0s[0]
|
||
if len(det) > 0:
|
||
# Rescale boxes from img_size to im0 size
|
||
if not padInfos:
|
||
det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()
|
||
else:
|
||
# print('####line131:',det[:, :])
|
||
det[:, :4] = scale_back(det[:, :4], padInfos).round()
|
||
|
||
|
||
for *xyxy, conf, cls in reversed(det):
|
||
cls_c = cls.cpu().numpy()
|
||
conf_c = conf.cpu().numpy()
|
||
tt = [int(x.cpu()) for x in xyxy]
|
||
if fiterList:
|
||
if int(cls) in fiterList: ###如果不是所需要的目标，则不显示
|
||
continue
|
||
if score_byClass:
|
||
if int(cls) in score_byClass.keys():
|
||
if conf < score_byClass[int(cls)]:
|
||
continue
|
||
line = [*tt, float(conf_c), float(cls_c)] # label format
|
||
det_xywh.append(line)
|
||
time2 = time.time()
|
||
strout='nms:%s ,detDraw:%s '%(get_ms(time0,time1), get_ms(time1,time2) )
|
||
return [im0s[0],im0s[0], det_xywh, 10],strout
|
||
|
||
def post_process_seg(im0s,segmodel,boxes,ksize):
|
||
time0 = time.time()
|
||
im0 = im0s[0].copy()
|
||
segmodel.set_image(im0s[0])
|
||
# # 创建一个空白掩码用于保存所有火焰
|
||
# combined_mask = np.zeros((im0.shape[0], im0.shape[1]), dtype=np.uint8)
|
||
# # 创建边缘可视化图像
|
||
# edge_image = np.zeros_like(im0s[0])
|
||
# 处理每个火焰检测框
|
||
det_xywhP = []
|
||
for box in boxes:
|
||
x_min, y_min, x_max, y_max = box[:4]
|
||
# 转换为SAM需要的格式
|
||
input_box = np.array([x_min, y_min, x_max, y_max])
|
||
|
||
# 使用框提示进行分割
|
||
masks, _, _ = segmodel.predict(
|
||
box=input_box,
|
||
multimask_output=False # 只返回最佳掩码
|
||
)
|
||
|
||
# 获取分割掩码
|
||
flame_mask = masks[0].astype(np.uint8)
|
||
# 使用形态学操作填充小孔洞
|
||
filled_mask = cv2.morphologyEx(flame_mask, cv2.MORPH_CLOSE, ksize)
|
||
# 查找所有轮廓（包括内部小点）
|
||
contours, _ = cv2.findContours(filled_mask.astype(np.uint8),
|
||
cv2.RETR_TREE,
|
||
cv2.CHAIN_APPROX_SIMPLE)
|
||
if len(contours) == 0:
|
||
continue
|
||
largest_contour = max(contours, key=cv2.contourArea)
|
||
# 通过轮廓填充。
|
||
#cv2.drawContours(im0, [largest_contour], -1, (0, 0, 255), 2)
|
||
box.append(largest_contour)
|
||
det_xywhP.append(box)
|
||
time1 = time.time()
|
||
strout = 'segDraw:%s ' % get_ms(time0, time1)
|
||
return [im0s[0], im0s[0], det_xywhP, 10], strout
|
||
|
||
def AI_process(im0s, model, segmodel, names, label_arraylist, rainbows,
|
||
objectPar={'half': True, 'device': 'cuda:0', 'conf_thres': 0.25, 'iou_thres': 0.45,
|
||
'segRegionCnt': 1, 'trtFlag_det': False,'trtFlag_seg': False,'score_byClass':None,'fiterList':[]},
|
||
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):
|
||
# 输入参数
|
||
# 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
|
||
|
||
half, device, conf_thres, iou_thres, fiterList,score_byClass = objectPar['half'], objectPar['device'], objectPar['conf_thres'], \
|
||
objectPar['iou_thres'], objectPar['fiterList'], objectPar['score_byClass']
|
||
|
||
trtFlag_det, trtFlag_seg, segRegionCnt = objectPar['trtFlag_det'], objectPar['trtFlag_seg'], objectPar[
|
||
'segRegionCnt']
|
||
|
||
time0 = time.time()
|
||
if trtFlag_det:
|
||
img, padInfos = img_pad(im0s[0], size=(640,640,3)) ;img = [img]
|
||
else:
|
||
#print('####line72:',im0s[0][10:12,10:12,2])
|
||
img = [letterbox(x, 640, auto=True, stride=32)[0] for x in im0s];padInfos=None
|
||
#print('####line74:',img[0][10:12,10:12,2])
|
||
# 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(device)
|
||
img = img.half() if half else img.float() # uint8 to fp16/32
|
||
img /= 255.0
|
||
time01 = time.time()
|
||
|
||
if segmodel:
|
||
seg_pred,segstr = segmodel.eval(im0s[0])
|
||
segFlag=True
|
||
else:
|
||
seg_pred = None;segFlag=False;segstr='Not implemented'
|
||
|
||
|
||
time1=time.time()
|
||
if trtFlag_det:
|
||
pred = yolov5Trtforward(model,img)
|
||
else:
|
||
#print('####line96:',img[0,0,10:12,10:12])
|
||
pred = model(img,augment=False)[0]
|
||
|
||
time2=time.time()
|
||
|
||
|
||
p_result, timeOut = getDetectionsFromPreds(pred, img, im0s[0], conf_thres=conf_thres, iou_thres=iou_thres,
|
||
ovlap_thres=None, padInfos=padInfos)
|
||
# if mode=='highWay3.0':
|
||
# if segmodel:
|
||
if segPar and segPar['mixFunction']['function']:
|
||
|
||
mixFunction = segPar['mixFunction']['function'];H,W = im0s[0].shape[0:2]
|
||
parMix = segPar['mixFunction']['pars'];#print('###line117:',parMix,p_result[2])
|
||
parMix['imgSize'] = (W,H)
|
||
#print(' -----------line110: ',p_result[2] ,'\n', seg_pred)
|
||
p_result[2] , timeMixPost= mixFunction(p_result[2], seg_pred, pars=parMix )
|
||
#print(' -----------line112: ',p_result[2] )
|
||
p_result.append(seg_pred)
|
||
|
||
else:
|
||
timeMixPost = ':0 ms'
|
||
# print('#### line121: segstr:%s timeMixPost:%s timeOut:%s'%( segstr.strip(), timeMixPost,timeOut ))
|
||
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)
|
||
if fiterList:
|
||
p_result[2] = filter_byClass(p_result[2], fiterList)
|
||
|
||
if score_byClass:
|
||
p_result[2] = score_filter_byClass(p_result[2], score_byClass)
|
||
|
||
print('-' * 10, p_result[2])
|
||
return p_result, time_info
|
||
|
||
|
||
def default_mix(predlist, par):
|
||
return predlist[0], ''
|
||
|
||
|
||
def AI_process_N(im0s, modelList, postProcess,score_byClass=None,fiterList=[]):
|
||
# 输入参数
|
||
## im0s---原始图像列表
|
||
## modelList--所有的模型
|
||
# postProcess--字典{},包括后处理函数，及其参数
|
||
#输出参数
|
||
##ret[0]--检测结果；
|
||
##ret[1]--时间信息
|
||
|
||
# modelList包括模型，每个模型是一个类，里面的eval函数可以输出该模型的推理结果
|
||
modelRets = [model.eval(im0s[0]) for model in modelList]
|
||
|
||
timeInfos = [x[1] for x in modelRets]
|
||
timeInfos = ''.join(timeInfos)
|
||
timeInfos = timeInfos
|
||
|
||
# postProcess['function']--后处理函数，输入的就是所有模型输出结果
|
||
mixFunction = postProcess['function']
|
||
predsList = [modelRet[0] for modelRet in modelRets]
|
||
H, W = im0s[0].shape[0:2]
|
||
postProcess['pars']['imgSize'] = (W, H)
|
||
|
||
# ret就是混合处理后的结果
|
||
ret = mixFunction(predsList, postProcess['pars'])
|
||
det = ret[0]
|
||
if fiterList:
|
||
det = filter_byClass(det, fiterList)
|
||
if score_byClass:
|
||
det = score_filter_byClass(det, score_byClass)
|
||
|
||
return det, timeInfos + ret[1]
|
||
def getMaxScoreWords(detRets0):
|
||
maxScore=-1;maxId=0
|
||
for i,detRet in enumerate(detRets0):
|
||
if detRet[4]>maxScore:
|
||
maxId=i
|
||
maxScore = detRet[4]
|
||
return maxId
|
||
|
||
def AI_process_C(im0s, modelList, postProcess,score_byClass,fiterList):
|
||
# 函数定制的原因：
|
||
## 之前模型处理流是
|
||
## 图片---> 模型1-->result1；图片---> 模型2->result2；[result1,result2]--->后处理函数
|
||
## 本函数的处理流程是
|
||
## 图片---> 模型1-->result1；[图片,result1]---> 模型2->result2；[result1,result2]--->后处理函数
|
||
## 模型2的输入，是有模型1的输出决定的。如模型2是ocr模型，需要将模型1检测出来的船名抠图出来输入到模型2.
|
||
## 之前的模型流都是模型2是分割模型，输入就是原始图片，与模型1的输出无关。
|
||
#输入参数
|
||
## im0s---原始图像列表
|
||
## modelList--所有的模型
|
||
# postProcess--字典{},包括后处理函数，及其参数
|
||
#输出参数
|
||
##ret[0]--检测结果；
|
||
##ret[1]--时间信息
|
||
|
||
#modelList包括模型，每个模型是一个类，里面的eval函数可以输出该模型的推理结果
|
||
|
||
t0=time.time()
|
||
detRets0 = modelList[0].eval(im0s[0])
|
||
|
||
#detRets0=[[12, 46, 1127, 1544, 0.2340087890625, 2.0], [1884, 1248, 2992, 1485, 0.64208984375, 1.0]]
|
||
detRets0 = detRets0[0]
|
||
parsIn=postProcess['pars']
|
||
|
||
_detRets0_obj = list(filter(lambda x: x[5] in parsIn['objs'], detRets0 ))
|
||
_detRets0_others = list(filter(lambda x: x[5] not in parsIn['objs'], detRets0 ))
|
||
_detRets0 = []
|
||
if postProcess['name']=='channel2':
|
||
if len(_detRets0_obj)>0:
|
||
maxId=getMaxScoreWords(_detRets0_obj)
|
||
_detRets0 = _detRets0_obj[maxId:maxId+1]
|
||
else: _detRets0 = detRets0
|
||
|
||
|
||
t1=time.time()
|
||
imagePatches = [ im0s[0][int(x[1]):int(x[3] ) ,int(x[0]):int(x[2])] for x in _detRets0 ]
|
||
detRets1 = [modelList[1].eval(patch) for patch in imagePatches]
|
||
print('###line240:',detRets1)
|
||
if postProcess['name']=='crackMeasurement':
|
||
detRets1 = [x[0]*255 for x in detRets1]
|
||
t2=time.time()
|
||
mixFunction =postProcess['function']
|
||
crackInfos = [mixFunction(patchMask,par=parsIn) for patchMask in detRets1]
|
||
|
||
rets = [detRets0[i]+ crackInfos[i] for i in range(len(imagePatches)) ]
|
||
t3=time.time()
|
||
outInfos='total:%.1f (det:%.1f %d次segs:%.1f mixProcess:%.1f) '%( (t3-t0)*1000, (t1-t0)*1000, len(detRets1),(t2-t1)*1000, (t3-t2)*1000 )
|
||
elif postProcess['name']=='channel2':
|
||
H,W = im0s[0].shape[0:2];parsIn['imgSize'] = (W,H)
|
||
mixFunction =postProcess['function']
|
||
_detRets0_others = mixFunction([_detRets0_others], parsIn)
|
||
ocrInfo='no ocr'
|
||
if len(_detRets0_obj)>0:
|
||
res_real = detRets1[0][0]
|
||
res_real="".join( list(filter(lambda x:(ord(x) >19968 and ord(x)<63865 ) or (ord(x) >47 and ord(x)<58 ),res_real)))
|
||
|
||
# detRets1[0][0]="".join( list(filter(lambda x:(ord(x) >19968 and ord(x)<63865 ) or (ord(x) >47 and ord(x)<58 ),detRets1[0][0])))
|
||
_detRets0_obj[maxId].append(res_real)
|
||
_detRets0_obj = [_detRets0_obj[maxId]] ##只输出有OCR的那个船名结果
|
||
ocrInfo = detRets1[0][1]
|
||
print(' _detRets0_obj:{} _detRets0_others:{} '.format(_detRets0_obj, _detRets0_others))
|
||
rets = _detRets0_obj + _detRets0_others
|
||
|
||
if fiterList:
|
||
rets = filter_byClass(rets, fiterList)
|
||
if score_byClass:
|
||
rets = score_filter_byClass(rets, score_byClass)
|
||
t3 = time.time()
|
||
outInfos = 'total:%.1f ,where det:%.1f, ocr:%s' % ((t3 - t0) * 1000, (t1 - t0) * 1000, ocrInfo)
|
||
|
||
#print('###line233:',detRets1,detRets0 )
|
||
|
||
return rets,outInfos
|
||
|
||
def AI_process_forest(im0s, model, segmodel, names, label_arraylist, rainbows, half=True, device=' cuda:0',
|
||
conf_thres=0.25, iou_thres=0.45,
|
||
font={'line_thickness': None, 'fontSize': None, 'boxLine_thickness': None,
|
||
'waterLineColor': (0, 255, 255), 'waterLineWidth': 3}, trtFlag_det=False,
|
||
SecNms=None,ksize=None,score_byClass=None,fiterList=[]):
|
||
# 输入参数
|
||
# 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
|
||
time0 = time.time()
|
||
if trtFlag_det:
|
||
img, padInfos = img_pad(im0s[0], size=(640,640,3)) ;img = [img]
|
||
else:
|
||
img = [letterbox(x, 640, auto=True, stride=32)[0] for x in im0s];padInfos=None
|
||
#img = [letterbox(x, 640, auto=True, stride=32)[0] for x in im0s]
|
||
# 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(device)
|
||
img = img.half() if half else img.float() # uint8 to fp16/32
|
||
|
||
img /= 255.0 # 0 - 255 to 0.0 - 1.0
|
||
|
||
time1 = time.time()
|
||
pred = yolov5Trtforward(model, img) if trtFlag_det else model(img, augment=False)[0]
|
||
|
||
p_result, timeOut = post_process_det(pred,padInfos,img,im0s,conf_thres,iou_thres,label_arraylist,rainbows,font,score_byClass,fiterList)
|
||
if segmodel and len(p_result[2])>0:
|
||
segmodel.set_image(im0s[0])
|
||
p_result, timeOut = post_process_seg(im0s,segmodel,p_result[2],ksize)
|
||
|
||
time2 = time.time()
|
||
time_info = 'letterbox:%.1f, infer:%.1f, ' % ((time1 - time0) * 1000, (time2 - time1) * 1000)
|
||
return p_result, time_info + timeOut
|
||
|
||
|
||
def AI_det_track(im0s_in, modelPar, processPar, sort_tracker, segPar=None):
|
||
im0s, iframe = im0s_in[0], im0s_in[1]
|
||
model = modelPar['det_Model']
|
||
segmodel = modelPar['seg_Model']
|
||
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
|
||
|
||
iou2nd = processPar['iou2nd']
|
||
time0=time.time()
|
||
|
||
if trtFlag_det:
|
||
img, padInfos = img_pad(im0s[0], size=(640,640,3)) ;img = [img]
|
||
else:
|
||
img = [letterbox(x, 640, auto=True, stride=32)[0] for x in im0s];padInfos=None
|
||
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(device)
|
||
img = img.half() if half else img.float() # uint8 to fp16/32
|
||
|
||
img /= 255.0 # 0 - 255 to 0.0 - 1.0
|
||
|
||
seg_pred = None;segFlag=False
|
||
time1=time.time()
|
||
pred = yolov5Trtforward(model,img) if trtFlag_det else model(img,augment=False)[0]
|
||
|
||
time2=time.time()
|
||
|
||
#p_result,timeOut = getDetections(datas,conf_thres, iou_thres,names,label_arraylist,rainbows,10,ObjectPar=ObjectPar,font=font,padInfos=padInfos)
|
||
p_result, timeOut = getDetectionsFromPreds(pred,img,im0s[0],conf_thres=conf_thres,iou_thres=iou_thres,ovlap_thres=iou2nd,padInfos=padInfos)
|
||
if score_byClass:
|
||
p_result[2] = score_filter_byClass(p_result[2],score_byClass)
|
||
if segmodel:
|
||
seg_pred,segstr = segmodel.eval(im0s[0] )
|
||
segFlag=True
|
||
else:
|
||
seg_pred = None;segFlag=False;segstr='No segmodel'
|
||
|
||
|
||
if segPar and segPar['mixFunction']['function']:
|
||
mixFunction = segPar['mixFunction']['function']
|
||
|
||
H,W = im0s[0].shape[0:2]
|
||
parMix = segPar['mixFunction']['pars'];#print('###line117:',parMix,p_result[2])
|
||
parMix['imgSize'] = (W,H)
|
||
|
||
|
||
p_result[2],timeInfos_post = mixFunction(p_result[2], seg_pred, pars=parMix )
|
||
timeInfos_seg_post = 'segInfer:%s ,postMixProcess:%s'%( segstr, timeInfos_post )
|
||
else:
|
||
timeInfos_seg_post = ' '
|
||
'''
|
||
if segmodel:
|
||
timeS1=time.time()
|
||
#seg_pred,segstr = segtrtEval(segmodel,im0s[0],par=segPar) if segPar['trtFlag_seg'] else segmodel.eval(im0s[0] )
|
||
seg_pred,segstr = segmodel.eval(im0s[0] )
|
||
timeS2=time.time()
|
||
mixFunction = segPar['mixFunction']['function']
|
||
|
||
p_result[2],timeInfos_post = mixFunction(p_result[2], seg_pred, pars=segPar['mixFunction']['pars'] )
|
||
|
||
timeInfos_seg_post = 'segInfer:%.1f ,postProcess:%s'%( (timeS2-timeS1)*1000, timeInfos_post )
|
||
|
||
else:
|
||
timeInfos_seg_post = ' '
|
||
#print('######line341:',seg_pred.shape,np.max(seg_pred),np.min(seg_pred) , len(p_result[2]) )
|
||
'''
|
||
time_info = 'letterbox:%.1f, detinfer:%.1f, '%( (time1-time0)*1000,(time2-time1)*1000 )
|
||
|
||
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,time_info+timeOut+timeInfos_seg_post
|
||
def AI_det_track_batch(imgarray_list, iframe_list ,modelPar,processPar,sort_tracker,trackPar,segPar=None):
|
||
'''
|
||
输入：
|
||
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( [ [imgarray_list[0]] ,iframe_list[0] ],modelPar,processPar,None,segPar )
|
||
##下面4行内容只是为了保持格式一致
|
||
detArray = np.array(p_result[2])
|
||
#print('##line371:',detArray)
|
||
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( [ [imgarray_list[index_frame]] ,iframe_list[index_frame] ],modelPar,processPar,sort_tracker,segPar )
|
||
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 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 get_tracker_cls(boxes,scId=4,clsId=5):
|
||
#正常来说一各跟踪链上是一个类别，但是有时目标框检测错误，导致有的跟踪链上有多个类别
|
||
#为此，根据跟踪链上每一个类别对应的所有框的置信度之和，作为这个跟踪链上目标的类别
|
||
#输入boxes--跟踪是保留的box_history,[[xc,yc,width,height,score,class,iframe],[...],[...]]
|
||
## scId=4,score所在的序号; clsId=5;类别所在的序号
|
||
#输出类别
|
||
##这个跟踪链上目标的类别
|
||
ids = list(set(boxes[:,clsId].tolist()))
|
||
scores = [np.sum( boxes[:,scId] [ boxes[:,clsId]==x ] ) for x in ids]
|
||
maxScoreId = scores.index(np.max(scores))
|
||
return int(ids[maxScoreId])
|
||
|
||
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).copy()
|
||
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;
|
||
|
||
#2023.11.17添加的。目的是修正跟踪链上所有的框的类别一样
|
||
chainClsId = get_tracker_cls(bbox_history,scId=4,clsId=5)
|
||
bbox_history[:,5] = chainClsId
|
||
|
||
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,'----------',interpolation_x0s.shape,track_det_result.shape,bbox_history ,'-----')
|
||
|
||
##将[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):
|
||
|
||
img_patch,engine,context,converter,AlignCollate_normal,device=pars[0:6]
|
||
time1 = time.time()
|
||
img_tensor = AlignCollate_normal([ Image.fromarray(img_patch,'L') ])
|
||
img_input = img_tensor.to('cuda:0')
|
||
time2 = time.time()
|
||
|
||
preds,trtstr=OcrTrtForward(engine,[img_input],context)
|
||
time3 = time.time()
|
||
|
||
batch_size = preds.size(0)
|
||
preds_size = torch.IntTensor([preds.size(1)] * batch_size)
|
||
|
||
######## filter ignore_char, rebalance
|
||
preds_prob = F.softmax(preds, dim=2)
|
||
preds_prob = preds_prob.cpu().detach().numpy()
|
||
pred_norm = preds_prob.sum(axis=2)
|
||
preds_prob = preds_prob/np.expand_dims(pred_norm, axis=-1)
|
||
preds_prob = torch.from_numpy(preds_prob).float().to(device)
|
||
_, preds_index = preds_prob.max(2)
|
||
preds_index = preds_index.view(-1)
|
||
time4 = time.time()
|
||
preds_str = converter.decode_greedy(preds_index.data.cpu().detach().numpy(), preds_size.data)
|
||
time5 = time.time()
|
||
|
||
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 ) )
|
||
return preds_str,info_str
|
||
|
||
def AI_process_Ocr(im0s, modelList, device, detpar):
|
||
timeMixPost = ':0 ms'
|
||
new_device = torch.device(device)
|
||
time0 = time.time()
|
||
|
||
img, padInfos = pre_process(im0s[0], new_device)
|
||
ocrModel = modelList[1]
|
||
time1 = time.time()
|
||
if not detpar['trtFlag_det']:
|
||
preds, timeOut = modelList[0].eval(img)
|
||
boxes = post_process(preds, padInfos, device, conf_thres=detpar['conf_thres'], iou_thres=detpar['iou_thres'],
|
||
nc=detpar['nc']) # 后处理
|
||
else:
|
||
boxes, timeOut = modelList[0].eval(im0s[0])
|
||
time2 = time.time()
|
||
|
||
imagePatches = [im0s[0][int(x[1]):int(x[3]), int(x[0]):int(x[2])] for x in boxes]
|
||
|
||
detRets1 = [ocrModel.eval(patch) for patch in imagePatches]
|
||
time3 = time.time()
|
||
dets = []
|
||
for i, (box, ocr) in enumerate(zip(boxes, detRets1)):
|
||
label = plat_format(ocr)
|
||
if label:
|
||
xyxy = box[0:4]
|
||
dets.append([label, xyxy])
|
||
|
||
time_info = 'pre_process:%.1f, det:%.1f , ocr:%.1f ,timeMixPost:%s ' % (
|
||
(time1 - time0) * 1000, (time2 - time1) * 1000, (time3 - time2) * 1000, timeMixPost)
|
||
|
||
return [im0s[0], im0s[0], dets, 0], time_info
|
||
|
||
|
||
def AI_process_Crowd(im0s,model,device,postPar):
|
||
timeMixPost = ':0 ms'
|
||
time0 = time.time()
|
||
h, w = im0s[0].shape[:2]
|
||
Wrate = w // 128 * 128 / w
|
||
Hrate = h // 128 * 128 / h
|
||
detM = np.array([])
|
||
if len(model) == 1:
|
||
detP = model[0].eval(im0s[0])
|
||
p_postPar = postPar[0]
|
||
|
||
else:
|
||
detP = model[1].eval(im0s[0])
|
||
p_postPar = postPar[1]
|
||
detM = np.array(model[0].eval(im0s[0])[0])
|
||
# d_postPar = postPar[0]
|
||
time1 = time.time()
|
||
outputs_scores = torch.nn.functional.softmax(detP['pred_logits'], -1)[:, :, 1][0]
|
||
|
||
outputs_points = detP['pred_points'][0]
|
||
|
||
detP = outputs_points[outputs_scores > p_postPar['conf']].detach().cpu().numpy()
|
||
detP[:, 0] = detP[:, 0] / Wrate
|
||
detP[:, 1] = detP[:, 1] / Hrate
|
||
|
||
dets = []
|
||
if not detM.size==0:
|
||
P = detP.shape[0] # 第二类边界框数量
|
||
if P == 0:
|
||
time_info = 'No Valid object find'
|
||
return [im0s[0], im0s[0], dets, 0], time_info
|
||
x1 = detP[:, 0] - p_postPar['expend']
|
||
y1 = detP[:, 1] - p_postPar['expend']
|
||
x2 = detP[:, 0] + p_postPar['expend']
|
||
y2 = detP[:, 1] + p_postPar['expend']
|
||
detP = np.column_stack((x1, y1, x2, y2))
|
||
detM = detM[:,:4]
|
||
a_x1, a_y1, a_x2, a_y2 = detM[:, 0], detM[:, 1], detM[:, 2], detM[:, 3]
|
||
b_x1, b_y1, b_x2, b_y2 = detP[:, 0], detP[:, 1], detP[:, 2], detP[:, 3]
|
||
|
||
# 扩展维度以便广播 [N,1]
|
||
a_x1 = a_x1[:, np.newaxis]
|
||
a_y1 = a_y1[:, np.newaxis]
|
||
a_x2 = a_x2[:, np.newaxis]
|
||
a_y2 = a_y2[:, np.newaxis]
|
||
|
||
# 计算相交区域坐标 [N,M]
|
||
inter_x1 = np.maximum(a_x1, b_x1)
|
||
inter_y1 = np.maximum(a_y1, b_y1)
|
||
inter_x2 = np.minimum(a_x2, b_x2)
|
||
inter_y2 = np.minimum(a_y2, b_y2)
|
||
# 计算相交区域面积 [N,M]
|
||
inter_area = np.maximum(0, inter_x2 - inter_x1) * np.maximum(0, inter_y2 - inter_y1)
|
||
|
||
# 计算每个A框与多少个B框相交（相交面积>0）
|
||
counts = np.sum(inter_area > 0, axis=1).astype(int)
|
||
|
||
# 构建结果数组：前4列是A框坐标，第5列是相交计数
|
||
detM = np.hstack((detM, counts.reshape(-1, 1)))
|
||
detM = detM[detM[:, -1] != 0]
|
||
# 找出所有与A框相交的B框索引
|
||
p_preds_index = np.any(inter_area > 0, axis=0)
|
||
p_indices = np.where(p_preds_index)[0]
|
||
# 提取与第一类相交的第二类边界框
|
||
if len(p_indices) > 0:
|
||
detP = detP[p_indices]
|
||
detP = detP[:, :2]
|
||
detP[:, 0] = detP[:, 0] + p_postPar['expend']
|
||
detP[:, 1] = detP[:, 1] + p_postPar['expend']
|
||
dets= [[detM.tolist(), detP.tolist(),p_postPar['psize']]]
|
||
else:
|
||
dets = [[detP.tolist(),p_postPar['psize']]]
|
||
|
||
# for b in detM:
|
||
# b_label = '该建筑物下行人及数量：%d'%(int(b[4]))
|
||
|
||
time2 = time.time()
|
||
|
||
time_info = 'det:%.1f , post:%.1f ,timeMixPost:%s ' % (
|
||
(time1 - time0) * 1000, (time2 - time1) * 1000, timeMixPost)
|
||
|
||
return [im0s[0], im0s[0], dets, 0], time_info
|
||
|
||
|
||
def main():
|
||
##预先设置的参数
|
||
device_='1' ##选定模型，可选 cpu,'0','1'
|
||
|
||
##以下参数目前不可改
|
||
Detweights = "weights/yolov5/class5/best_5classes.pt"
|
||
seg_nclass = 2
|
||
Segweights = "weights/BiSeNet/checkpoint.pth"
|
||
conf_thres,iou_thres,classes= 0.25,0.45,5
|
||
labelnames = "weights/yolov5/class5/labelnames.json"
|
||
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]]
|
||
allowedList=[0,1,2,3]
|
||
|
||
|
||
##加载模型，准备好显示字符
|
||
device = select_device(device_)
|
||
names=get_labelnames(labelnames)
|
||
label_arraylist = get_label_arrays(names,rainbows,outfontsize=40,fontpath="conf/platech.ttf")
|
||
half = device.type != 'cpu' # half precision only supported on CUDA
|
||
model = attempt_load(Detweights, map_location=device) # load FP32 model
|
||
if half: model.half()
|
||
segmodel = SegModel(nclass=seg_nclass,weights=Segweights,device=device)
|
||
|
||
|
||
##图像测试
|
||
#url='images/examples/20220624_响水河_12300_1621.jpg'
|
||
impth = 'images/examples/'
|
||
outpth = 'images/results/'
|
||
folders = os.listdir(impth)
|
||
for i in range(len(folders)):
|
||
imgpath = os.path.join(impth, folders[i])
|
||
im0s=[cv2.imread(imgpath)]
|
||
time00 = time.time()
|
||
p_result,timeOut = AI_process(im0s,model,segmodel,names,label_arraylist,rainbows,half,device,conf_thres, iou_thres,allowedList,fontSize=1.0)
|
||
time11 = time.time()
|
||
image_array = p_result[1]
|
||
cv2.imwrite( os.path.join( outpth,folders[i] ) ,image_array )
|
||
#print('----process:%s'%(folders[i]), (time.time() - time11) * 1000)
|
||
|
||
|
||
|
||
|
||
|
||
if __name__=="__main__":
|
||
main()
|