1)新增M029 火焰面积 2）算法支持按类过滤 3）算法支持按置信度过滤 4）其他优化
This commit is contained in:
parent
796fb6bcad
commit
464923ee37
283
AI.py
283
AI.py
|
|
@ -5,11 +5,13 @@ 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,get_label_arrays,post_process_,img_pad,draw_painting_joint,detectDraw,getDetections,getDetectionsFromPreds
|
||||
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
|
||||
|
|
@ -18,6 +20,7 @@ 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=[];###获取文件里所有的图像
|
||||
|
|
@ -33,9 +36,9 @@ def get_images_videos(impth, imageFixs=['.jpg','.JPG','.PNG','.png'],videoFixs=[
|
|||
if postfix in videoFixs: videopaths = [impth ]
|
||||
|
||||
print('%s: test Images:%d , test videos:%d '%(impth, len(imgpaths), len(videopaths)))
|
||||
return imgpaths,videopaths
|
||||
return imgpaths,videopaths
|
||||
|
||||
def xywh2xyxy(box,iW=None,iH=None):
|
||||
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)
|
||||
|
|
@ -73,13 +76,15 @@ def score_filter_byClass(pdetections,score_para_2nd):
|
|||
ret.append(det)
|
||||
return ret
|
||||
# 按类过滤
|
||||
def filter_byClass(pdetections,allowedList):
|
||||
ret=[]
|
||||
def filter_byClass(pdetections, fiterList):
|
||||
ret = []
|
||||
for det in pdetections:
|
||||
score,cls = det[4],det[5]
|
||||
if int(cls) in allowedList:
|
||||
ret.append(det)
|
||||
elif str(int(cls)) in allowedList:
|
||||
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
|
||||
|
|
@ -111,9 +116,90 @@ def plat_format(ocr):
|
|||
|
||||
return label.upper()
|
||||
|
||||
def AI_process(im0s,model,segmodel,names,label_arraylist,rainbows,objectPar={ 'half':True,'device':'cuda:0' ,'conf_thres':0.25,'iou_thres':0.45,'allowedList':[0,1,2,3],'segRegionCnt':1, 'trtFlag_det':False,'trtFlag_seg':False,'score_byClass':{x:0.1 for x in range(30)} }, 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):
|
||||
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）
|
||||
#
|
||||
|
|
@ -126,16 +212,13 @@ def AI_process(im0s,model,segmodel,names,label_arraylist,rainbows,objectPar={ 'h
|
|||
# #strout---统计AI处理个环节的时间
|
||||
# Letterbox
|
||||
|
||||
half,device,conf_thres,iou_thres,allowedList = objectPar['half'],objectPar['device'],objectPar['conf_thres'],objectPar['iou_thres'],objectPar['allowedList']
|
||||
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']
|
||||
if 'ovlap_thres_crossCategory' in objectPar.keys(): ovlap_thres = objectPar['ovlap_thres_crossCategory']
|
||||
else: ovlap_thres = None
|
||||
trtFlag_det, trtFlag_seg, segRegionCnt = objectPar['trtFlag_det'], objectPar['trtFlag_seg'], objectPar[
|
||||
'segRegionCnt']
|
||||
|
||||
if 'score_byClass' in objectPar.keys(): score_byClass = objectPar['score_byClass']
|
||||
else: score_byClass = None
|
||||
|
||||
time0=time.time()
|
||||
time0 = time.time()
|
||||
if trtFlag_det:
|
||||
img, padInfos = img_pad(im0s[0], size=(640,640,3)) ;img = [img]
|
||||
else:
|
||||
|
|
@ -151,10 +234,10 @@ def AI_process(im0s,model,segmodel,names,label_arraylist,rainbows,objectPar={ 'h
|
|||
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()
|
||||
time01 = time.time()
|
||||
|
||||
if segmodel:
|
||||
seg_pred,segstr = segmodel.eval(im0s[0] )
|
||||
if segmodel:
|
||||
seg_pred,segstr = segmodel.eval(im0s[0])
|
||||
segFlag=True
|
||||
else:
|
||||
seg_pred = None;segFlag=False;segstr='Not implemented'
|
||||
|
|
@ -170,12 +253,11 @@ def AI_process(im0s,model,segmodel,names,label_arraylist,rainbows,objectPar={ 'h
|
|||
time2=time.time()
|
||||
|
||||
|
||||
p_result, timeOut = getDetectionsFromPreds(pred,img,im0s[0],conf_thres=conf_thres,iou_thres=iou_thres,ovlap_thres=ovlap_thres,padInfos=padInfos)
|
||||
if score_byClass:
|
||||
p_result[2] = score_filter_byClass(p_result[2],score_byClass)
|
||||
#if mode=='highWay3.0':
|
||||
#if segmodel:
|
||||
if segPar and segPar['mixFunction']['function']:
|
||||
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])
|
||||
|
|
@ -186,19 +268,26 @@ def AI_process(im0s,model,segmodel,names,label_arraylist,rainbows,objectPar={ 'h
|
|||
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 allowedList:
|
||||
p_result[2] = filter_byClass(p_result[2],allowedList)
|
||||
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)
|
||||
|
||||
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):
|
||||
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--字典{},包括后处理函数，及其参数
|
||||
|
|
@ -206,22 +295,28 @@ def AI_process_N(im0s,modelList,postProcess):
|
|||
##ret[0]--检测结果；
|
||||
##ret[1]--时间信息
|
||||
|
||||
#modelList包括模型，每个模型是一个类，里面的eval函数可以输出该模型的推理结果
|
||||
modelRets=[ model.eval(im0s[0]) for model in modelList]
|
||||
# modelList包括模型，每个模型是一个类，里面的eval函数可以输出该模型的推理结果
|
||||
modelRets = [model.eval(im0s[0]) for model in modelList]
|
||||
|
||||
timeInfos = [ x[1] for x in modelRets]
|
||||
timeInfos=''.join(timeInfos)
|
||||
timeInfos=timeInfos
|
||||
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)
|
||||
# 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'])
|
||||
return ret[0],timeInfos+ret[1]
|
||||
# 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):
|
||||
|
|
@ -230,8 +325,8 @@ def getMaxScoreWords(detRets0):
|
|||
maxScore = detRet[4]
|
||||
return maxId
|
||||
|
||||
def AI_process_C(im0s,modelList,postProcess):
|
||||
#函数定制的原因：
|
||||
def AI_process_C(im0s, modelList, postProcess,score_byClass,fiterList):
|
||||
# 函数定制的原因：
|
||||
## 之前模型处理流是
|
||||
## 图片---> 模型1-->result1；图片---> 模型2->result2；[result1,result2]--->后处理函数
|
||||
## 本函数的处理流程是
|
||||
|
|
@ -287,21 +382,30 @@ def AI_process_C(im0s,modelList,postProcess):
|
|||
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
|
||||
t3=time.time()
|
||||
outInfos='total:%.1f ,where det:%.1f, ocr:%s'%( (t3-t0)*1000, (t1-t0)*1000, ocrInfo)
|
||||
# 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,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):
|
||||
#输入参数
|
||||
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
|
||||
|
|
@ -313,7 +417,7 @@ def AI_process_forest(im0s,model,segmodel,names,label_arraylist,rainbows,half=Tr
|
|||
# #strout---统计AI处理个环节的时间
|
||||
|
||||
# Letterbox
|
||||
time0=time.time()
|
||||
time0 = time.time()
|
||||
if trtFlag_det:
|
||||
img, padInfos = img_pad(im0s[0], size=(640,640,3)) ;img = [img]
|
||||
else:
|
||||
|
|
@ -329,26 +433,22 @@ def AI_process_forest(im0s,model,segmodel,names,label_arraylist,rainbows,half=Tr
|
|||
img = img.half() if half else img.float() # uint8 to fp16/32
|
||||
|
||||
img /= 255.0 # 0 - 255 to 0.0 - 1.0
|
||||
if segmodel:
|
||||
seg_pred,segstr = segmodel.eval(im0s[0] )
|
||||
segFlag=True
|
||||
else:
|
||||
seg_pred = None;segFlag=False
|
||||
time1=time.time()
|
||||
pred = yolov5Trtforward(model,img) if trtFlag_det else model(img,augment=False)[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
|
||||
|
||||
|
||||
time2=time.time()
|
||||
datas = [[''], img, im0s, None,pred,seg_pred,10]
|
||||
|
||||
ObjectPar={ 'object_config':allowedList, 'slopeIndex':[] ,'segmodel':segFlag,'segRegionCnt':0 }
|
||||
p_result,timeOut = post_process_(datas,conf_thres, iou_thres,names,label_arraylist,rainbows,10,ObjectPar=ObjectPar,font=font,padInfos=padInfos,ovlap_thres=SecNms)
|
||||
#print('###line274:',p_result[2])
|
||||
#p_result,timeOut = post_process_(datas,conf_thres, iou_thres,names,label_arraylist,rainbows,10,object_config=allowedList,segmodel=segFlag,font=font,padInfos=padInfos)
|
||||
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]
|
||||
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']
|
||||
|
|
@ -705,17 +805,22 @@ def ocr_process(pars):
|
|||
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):
|
||||
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()
|
||||
preds,timeOut = modelList[0].eval(img)
|
||||
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()
|
||||
boxes = post_process(preds, padInfos, device, conf_thres=detpar['conf_thres'], iou_thres=detpar['iou_thres'],
|
||||
nc=detpar['nc']) # 后处理
|
||||
|
||||
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]
|
||||
|
|
@ -728,9 +833,9 @@ def AI_process_Ocr(im0s,modelList,device,detpar):
|
|||
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)
|
||||
(time1 - time0) * 1000, (time2 - time1) * 1000, (time3 - time2) * 1000, timeMixPost)
|
||||
|
||||
return [im0s[0],im0s[0],dets,0],time_info
|
||||
return [im0s[0], im0s[0], dets, 0], time_info
|
||||
|
||||
|
||||
def AI_process_Crowd(im0s,model,device,postPar):
|
||||
|
|
|
|||
Loading…
Reference in New Issue