|
|
|
@@ -0,0 +1,922 @@ |
|
|
|
import sys, yaml
|
|
|
|
from easydict import EasyDict as edict
|
|
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
|
|
|
|
|
|
sys.path.extend(['..','../AIlib2' ])
|
|
|
|
|
|
|
|
from AI import AI_process,AI_process_forest,get_postProcess_para,AI_Seg_process,ocr_process
|
|
|
|
import cv2,os,time
|
|
|
|
from segutils.segmodel import SegModel
|
|
|
|
from models.experimental import attempt_load
|
|
|
|
from utils.torch_utils import select_device
|
|
|
|
from utilsK.queRiver import get_labelnames,get_label_arrays,save_problem_images
|
|
|
|
from ocrUtils.ocrUtils import CTCLabelConverter,AlignCollate
|
|
|
|
from obbUtils.shipUtils import OBB_infer
|
|
|
|
from obbUtils.load_obb_model import load_model_decoder_OBB
|
|
|
|
import numpy as np
|
|
|
|
import torch,glob
|
|
|
|
import tensorrt as trt
|
|
|
|
from utilsK.masterUtils import get_needed_objectsIndex
|
|
|
|
#import warnings
|
|
|
|
#warnings.filterwarnings("error")
|
|
|
|
|
|
|
|
def view_bar(num, total,time1,prefix='prefix'):
|
|
|
|
rate = num / total
|
|
|
|
time_n=time.time()
|
|
|
|
rate_num = int(rate * 30)
|
|
|
|
rate_nums = np.round(rate * 100)
|
|
|
|
r = '\r %s %d / %d [%s%s] %.2f s'%(prefix,num,total, ">" * rate_num, " " * (30 - rate_num), time_n-time1 )
|
|
|
|
sys.stdout.write(r)
|
|
|
|
sys.stdout.flush()
|
|
|
|
|
|
|
|
|
|
|
|
'''
|
|
|
|
多线程
|
|
|
|
'''
|
|
|
|
|
|
|
|
def process_v1(frame):
|
|
|
|
#try:
|
|
|
|
print('demo.py beging to :',frame[8])
|
|
|
|
time00 = time.time()
|
|
|
|
H,W,C = frame[0][0].shape
|
|
|
|
|
|
|
|
p_result,timeOut = AI_process(frame[0],frame[1],frame[2],frame[3],frame[4],frame[5],objectPar=frame[6],font=frame[7],segPar=frame[9],mode=frame[10],postPar=frame[11])
|
|
|
|
|
|
|
|
time11 = time.time()
|
|
|
|
image_array = p_result[1]
|
|
|
|
|
|
|
|
cv2.imwrite(os.path.join('images/results/',frame[8] ) ,image_array)
|
|
|
|
bname = frame[8].split('.')[0]
|
|
|
|
if len(p_result)==5:
|
|
|
|
image_mask = p_result[4]
|
|
|
|
cv2.imwrite(os.path.join('images/results/',bname+'_mask.png' ) , (image_mask*50).astype(np.uint8))
|
|
|
|
|
|
|
|
boxes=p_result[2]
|
|
|
|
with open( os.path.join('images/results/',bname+'.txt' ),'w' ) as fp:
|
|
|
|
for box in boxes:
|
|
|
|
box_str=[str(x) for x in box]
|
|
|
|
out_str=','.join(box_str)+'\n'
|
|
|
|
fp.write(out_str)
|
|
|
|
time22 = time.time()
|
|
|
|
print('%s,%d*%d,AI-process: %.1f,image save:%.1f , %s'%(frame[8],H,W, (time11 - time00) * 1000.0, (time22-time11)*1000.0,timeOut), boxes)
|
|
|
|
return 'success'
|
|
|
|
#except Exception as e:
|
|
|
|
# return 'failed:'+str(e)
|
|
|
|
def process_video(video,par0,mode='detSeg'):
|
|
|
|
cap=cv2.VideoCapture(video)
|
|
|
|
if not cap.isOpened():
|
|
|
|
print('#####error url:',video)
|
|
|
|
return False
|
|
|
|
bname=os.path.basename(video).split('.')[0]
|
|
|
|
fps = int(cap.get(cv2.CAP_PROP_FPS)+0.5)
|
|
|
|
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH )+0.5)
|
|
|
|
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)+0.5)
|
|
|
|
framecnt=int(cap.get(7)+0.5)
|
|
|
|
save_path_AI = os.path.join(par0['outpth'],os.path.basename(video))
|
|
|
|
problem_image_dir= os.path.join( par0['outpth'], 'probleImages' )
|
|
|
|
os.makedirs(problem_image_dir,exist_ok=True)
|
|
|
|
vid_writer_AI = cv2.VideoWriter(save_path_AI, cv2.VideoWriter_fourcc(*'mp4v'), fps, (width,height))
|
|
|
|
num=0
|
|
|
|
iframe=0;post_results=[];fpsample=30*10
|
|
|
|
while cap.isOpened():
|
|
|
|
ret, imgarray = cap.read() #读取摄像头画面
|
|
|
|
if not ret:break
|
|
|
|
if mode=='detSeg':
|
|
|
|
p_result,timeOut = AI_process([imgarray],par0['model'],par0['segmodel'],par0['names'],par0['label_arraylist'],par0['rainbows'],objectPar=par0['objectPar'],font=par0['digitFont'],segPar=par0['segPar'])
|
|
|
|
else:
|
|
|
|
p_result,timeOut = AI_process_forest([imgarray],par0['model'],par0['segmodel'],par0['names'],par0['label_arraylist'],par0['rainbows'],par0['half'],par0['device'],par0['conf_thres'], par0['iou_thres'],par0['allowedList'],font=par0['digitFont'],trtFlag_det=par0['trtFlag_det'])
|
|
|
|
|
|
|
|
|
|
|
|
image_array = p_result[1];num+=1
|
|
|
|
ret = vid_writer_AI.write(image_array)
|
|
|
|
view_bar(num, framecnt,time.time(),prefix=os.path.basename(video))
|
|
|
|
##每隔 fpsample帧处理一次，如果有问题就保存图片
|
|
|
|
if (iframe % fpsample == 0) and (len(post_results)>0) :
|
|
|
|
parImage=save_problem_images(post_results,iframe,par0['names'],streamName=bname,outImaDir=problem_image_dir,imageTxtFile=False)
|
|
|
|
post_results=[]
|
|
|
|
|
|
|
|
if len(p_result[2] )>0: ##
|
|
|
|
post_results.append(p_result)
|
|
|
|
iframe+=1
|
|
|
|
vid_writer_AI.release();
|
|
|
|
def detSeg_demo(opt):
|
|
|
|
|
|
|
|
|
|
|
|
###河道巡检的参数####
|
|
|
|
if opt['business'] == 'river':
|
|
|
|
par={
|
|
|
|
'device':'0', ###显卡号,如果用TRT模型，只支持0（单显卡）
|
|
|
|
'labelnames':"../AIlib2/weights/conf/river/labelnames.json", ###检测类别对照表
|
|
|
|
'gpuname':'3090',###显卡名称
|
|
|
|
'max_workers':1, ###并行线程数
|
|
|
|
'trtFlag_det':True,###检测模型是否采用TRT
|
|
|
|
'trtFlag_seg':True,###分割模型是否采用TRT
|
|
|
|
'Detweights':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###检测模型路径
|
|
|
|
|
|
|
|
'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,4,5,6,7] ],###控制哪些检测类别显示，输出
|
|
|
|
#'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [] ],###控制哪些检测类别显示、输出
|
|
|
|
'slopeIndex':[5,6,7],###岸坡类别（或者其它业务里的类别），不与河道（分割的前景区域）计算交并比，即不论是否在河道内都显示。
|
|
|
|
'seg_nclass':2,###分割模型类别数目，默认2类
|
|
|
|
'segRegionCnt':1,###分割模型结果需要保留的等值线数目
|
|
|
|
'segPar':{'modelSize':(640,360),'mean':(0.485, 0.456, 0.406),'std' :(0.229, 0.224, 0.225),'numpy':False, 'RGB_convert_first':True},###分割模型预处理参数
|
|
|
|
'Segweights' : "../weights/%s/AIlib2/%s/stdc_360X640_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###分割模型权重位置
|
|
|
|
'postFile': '../AIlib2/weights/conf/river/para.json',###后处理参数文件
|
|
|
|
'txtFontSize':40,###文本字符的大小
|
|
|
|
'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'segLineShow':True,'waterLineWidth':3},###显示框、线、数字设置
|
|
|
|
'testImgPath':'images/tt/',
|
|
|
|
#'testImgPath':'/home/thsw2/WJ/data/XunHe/huData/Huzhou/original/',
|
|
|
|
#'images/river2/',
|
|
|
|
#'../../../data/无人机起飞测试图像/',###测试图像的位置
|
|
|
|
'testOutPath':'images/results/',###输出测试图像位置
|
|
|
|
#'testOutPath':'/home/thsw2/WJ/data/XunHe/huData/Huzhou/AIProcess',
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
####公路巡检的参数
|
|
|
|
if opt['business'] == 'road':
|
|
|
|
par={
|
|
|
|
'device':'0', ###显卡号,如果用TRT模型，只支持0（单显卡）
|
|
|
|
'labelnames':"../AIlib2/weights/conf/road/labelnames.json", ###检测类别对照表
|
|
|
|
'gpuname':'3090',###显卡名称
|
|
|
|
'max_workers':1, ###并行线程数
|
|
|
|
'trtFlag_det':True,###检测模型是否采用TRT
|
|
|
|
'trtFlag_seg':True,###分割模型是否采用TRT
|
|
|
|
'Detweights':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###检测模型路径
|
|
|
|
'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,3,4,5,6] ],###控制哪些检测类别显示、输出
|
|
|
|
'slopeIndex':[5,6,7],###岸坡类别（或者其它业务里的类别），不与河道（分割的前景区域）计算交并比，即不论是否在河道内都显示。
|
|
|
|
'seg_nclass':2,###分割模型类别数目，默认2类
|
|
|
|
'segRegionCnt':2,###分割模型结果需要保留的等值线数目
|
|
|
|
'segPar':{'modelSize':(640,360),'mean':(0.485, 0.456, 0.406),'std' :(0.229, 0.224, 0.225),'numpy':False, 'RGB_convert_first':True},###分割模型预处理参数
|
|
|
|
'Segweights' : "../weights/%s/AIlib2/%s/stdc_360X640_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###分割模型权重位置
|
|
|
|
'postFile': '../AIlib2/weights/conf/road/para.json',###后处理参数文件
|
|
|
|
'txtFontSize':20,###文本字符的大小
|
|
|
|
'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'waterLineWidth':2},###显示框、线设置
|
|
|
|
'testImgPath':'../AIdemo2/images/road/',###测试图像的位置
|
|
|
|
'testOutPath':'images/results/',###输出测试图像位置
|
|
|
|
}
|
|
|
|
if opt['business'] == 'highWay2':
|
|
|
|
par={
|
|
|
|
'device':'0', ###显卡号,如果用TRT模型，只支持0（单显卡）
|
|
|
|
'labelnames':"../AIlib2/weights/conf/%s/labelnames.json"%(opt['business']), ###检测类别对照表
|
|
|
|
'gpuname':'3090',###显卡名称
|
|
|
|
'max_workers':1, ###并行线程数
|
|
|
|
'trtFlag_det':True,###检测模型是否采用TRT
|
|
|
|
'trtFlag_seg':True,###分割模型是否采用TRT
|
|
|
|
'Detweights':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###检测模型路径
|
|
|
|
#'Detweights':"../AIlib2/weights/conf/highWay2/yolov5.pt",
|
|
|
|
'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,5,6,7,8,9] ],###控制哪些检测类别显示、输出
|
|
|
|
#'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [] ],
|
|
|
|
#'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [9] ],###控制哪些检测类别显示、输出
|
|
|
|
|
|
|
|
'slopeIndex':[],###岸坡类别（或者其它业务里的类别），不与河道（分割的前景区域）计算交并比，即不论是否在河道内都显示。
|
|
|
|
'seg_nclass':3,###分割模型类别数目，默认2类
|
|
|
|
'segRegionCnt':2,###分割模型结果需要保留的等值线数目
|
|
|
|
'segPar':{'modelSize':(640,360),'mean':(0.485, 0.456, 0.406),'std' :(0.229, 0.224, 0.225),'predResize':True,'numpy':False, 'RGB_convert_first':True},###分割模型预处理参数
|
|
|
|
|
|
|
|
'postPar': {'label_csv': '../AIlib2/weights/conf/trafficAccident/class_dict.csv', 'speedRoadArea': 16000, 'vehicleArea': 10, 'speedRoadVehicleAngleMin': 15, 'speedRoadVehicleAngleMax': 75, 'roundness': 0.7, 'cls': 9, 'vehicleFactor': 0.1,'cls':9},
|
|
|
|
|
|
|
|
'mode':'highWay3.0',
|
|
|
|
'Segweights' : "../weights/%s/AIlib2/%s/stdc_360X640_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###分割模型权重位置
|
|
|
|
'postFile': '../AIlib2/weights/conf/%s/para.json'%(opt['business'] ),###后处理参数文件
|
|
|
|
'txtFontSize':20,###文本字符的大小
|
|
|
|
'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'segLineShow':True,'waterLineWidth':2},###显示框、线设置
|
|
|
|
'testImgPath':'../AIdemo2/images/trafficAccident/',###测试图像的位置
|
|
|
|
'testOutPath':'images/results/',###输出测试图像位置
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
mode = par['mode'] if 'mode' in par.keys() else 'others'
|
|
|
|
postPar = par['postPar'] if 'postPar' in par.keys() else None
|
|
|
|
|
|
|
|
device_=par['device']
|
|
|
|
labelnames = par['labelnames'] ##对应类别表
|
|
|
|
gpuname= par['gpuname']
|
|
|
|
max_workers=par['max_workers'];
|
|
|
|
trtFlag_det=par['trtFlag_det'];trtFlag_seg=par['trtFlag_seg'];segRegionCnt=par['segRegionCnt']
|
|
|
|
device = select_device(device_)
|
|
|
|
names=get_labelnames(labelnames)
|
|
|
|
|
|
|
|
half = device.type != 'cpu' # half precision only supported on CUDA
|
|
|
|
if trtFlag_det:
|
|
|
|
Detweights = par['Detweights']##升级后的检测模型
|
|
|
|
logger = trt.Logger(trt.Logger.ERROR)
|
|
|
|
with open(Detweights, "rb") as f, trt.Runtime(logger) as runtime:
|
|
|
|
model=runtime.deserialize_cuda_engine(f.read())# 输入trt本地文件，返回ICudaEngine对象
|
|
|
|
print('############locad det model trtsuccess:',Detweights)
|
|
|
|
else:
|
|
|
|
Detweights = par['Detweights']
|
|
|
|
model = attempt_load(Detweights, map_location=device) # load FP32 model
|
|
|
|
print('############locad det model pth success:',Detweights)
|
|
|
|
if half: model.half()
|
|
|
|
|
|
|
|
|
|
|
|
seg_nclass = par['seg_nclass']
|
|
|
|
segPar=par['segPar']
|
|
|
|
if trtFlag_seg:
|
|
|
|
Segweights = par['Segweights']
|
|
|
|
logger = trt.Logger(trt.Logger.ERROR)
|
|
|
|
with open(Segweights, "rb") as f, trt.Runtime(logger) as runtime:
|
|
|
|
segmodel=runtime.deserialize_cuda_engine(f.read())# 输入trt本地文件，返回ICudaEngine对象
|
|
|
|
print('############locad seg model trt success: ',Segweights)
|
|
|
|
else:
|
|
|
|
Segweights = par['Segweights']
|
|
|
|
segmodel = SegModel(nclass=seg_nclass,weights=Segweights,device=device)
|
|
|
|
print('############locad seg model pth success:',Segweights)
|
|
|
|
|
|
|
|
postFile= par['postFile']
|
|
|
|
digitFont= par['digitFont']
|
|
|
|
conf_thres,iou_thres,classes,rainbows=get_postProcess_para(postFile)
|
|
|
|
|
|
|
|
#print('#####line225:',par['mode'])
|
|
|
|
####模型选择参数用如下：
|
|
|
|
mode_paras=par['detModelpara'];slopeIndex=par['slopeIndex']
|
|
|
|
allowedList,allowedList_string=get_needed_objectsIndex(mode_paras)
|
|
|
|
#allowedList=[0,1,2,3]
|
|
|
|
##加载模型，准备好显示字符
|
|
|
|
label_arraylist = get_label_arrays(names,rainbows,outfontsize=par['txtFontSize'],fontpath="../AIlib2/conf/platech.ttf")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
##图像测试
|
|
|
|
#impth = 'images/slope/'
|
|
|
|
impth = par['testImgPath']
|
|
|
|
outpth = par['testOutPath']
|
|
|
|
imgpaths=[]###获取文件里所有的图像
|
|
|
|
for postfix in ['.jpg','.JPG','.PNG','.png']:
|
|
|
|
imgpaths.extend(glob.glob('%s/*%s'%(impth,postfix )) )
|
|
|
|
videopaths=[]###获取文件里所有的视频
|
|
|
|
for postfix in ['.MP4','.mp4','.avi']:
|
|
|
|
videopaths.extend(glob.glob('%s/*%s'%(impth,postfix )) )
|
|
|
|
|
|
|
|
|
|
|
|
###先处理图像
|
|
|
|
frames=[]
|
|
|
|
for imgpath in imgpaths:
|
|
|
|
im0s=[cv2.imread(imgpath)]
|
|
|
|
objectPar={ 'half':half,'device':device,'conf_thres':conf_thres,'iou_thres':iou_thres,'allowedList':allowedList,'slopeIndex':slopeIndex,'segRegionCnt':segRegionCnt, 'trtFlag_det':trtFlag_det,'trtFlag_seg':trtFlag_seg }
|
|
|
|
|
|
|
|
frame=(im0s,model,segmodel,names,label_arraylist,rainbows,objectPar,digitFont,os.path.basename(imgpath),segPar,mode,postPar)
|
|
|
|
frames.append(frame)
|
|
|
|
t1=time.time()
|
|
|
|
if max_workers==1:
|
|
|
|
for i in range(len(imgpaths)):
|
|
|
|
t5=time.time()
|
|
|
|
process_v1(frames[i])
|
|
|
|
t6=time.time()
|
|
|
|
print('#######%s, ms:%.1f , accumetate time:%.1f, avage:%1.f '%(os.path.basename(imgpaths[i]), (t6-t5)*1000.0,(t6-t1)*1000.0, (t6-t1)*1000.0/(i+1)))
|
|
|
|
else:
|
|
|
|
with ThreadPoolExecutor(max_workers=max_workers) as t:
|
|
|
|
for result in t.map(process_v1, frames):
|
|
|
|
#print(result)
|
|
|
|
t=result
|
|
|
|
|
|
|
|
t2=time.time()
|
|
|
|
if len(imgpaths)>0:
|
|
|
|
print('All %d images time:%.1f ms ,each:%.1f ms, with %d threads'%(len(imgpaths),(t2-t1)*1000, (t2-t1)*1000.0/len(imgpaths) , max_workers) )
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
objectPar={ 'half':half,'device':device,'conf_thres':conf_thres,'iou_thres':iou_thres,'allowedList':allowedList,'slopeIndex':slopeIndex,'segRegionCnt':segRegionCnt, 'trtFlag_det':trtFlag_det,'trtFlag_seg':trtFlag_seg }
|
|
|
|
|
|
|
|
par0={ 'model':model,'segmodel':segmodel,'names':names,'label_arraylist':label_arraylist,'rainbows':rainbows,
|
|
|
|
'objectPar':objectPar,'digitFont':digitFont,'segPar':segPar,'outpth':outpth
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
###如果是视频文件
|
|
|
|
for video in videopaths:
|
|
|
|
process_video(video,par0)
|
|
|
|
print(' ')
|
|
|
|
|
|
|
|
|
|
|
|
def det_demo(business ):
|
|
|
|
|
|
|
|
####森林巡检的参数
|
|
|
|
if opt['business'] == 'forest':
|
|
|
|
par={
|
|
|
|
'device':'0', ###显卡号,如果用TRT模型，只支持0（单显卡）
|
|
|
|
'labelnames':"../AIlib2/weights/conf/forest/labelnames.json", ###检测类别对照表
|
|
|
|
'gpuname':'3090',###显卡名称
|
|
|
|
'max_workers':1, ###并行线程数
|
|
|
|
'trtFlag_det':True,###检测模型是否采用TRT
|
|
|
|
'trtFlag_seg':False,###分割模型是否采用TRT
|
|
|
|
'Detweights':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###检测模型路径
|
|
|
|
'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,4,5,6] ],###控制哪些检测类别显示、输出
|
|
|
|
'slopeIndex':[],###岸坡类别（或者其它业务里的类别），不与河道（分割的前景区域）计算交并比，即不论是否在河道内都显示。
|
|
|
|
'seg_nclass':2,###分割模型类别数目，默认2类
|
|
|
|
'segRegionCnt':0,###分割模型结果需要保留的等值线数目
|
|
|
|
'segPar':None,###分割模型预处理参数
|
|
|
|
'Segweights' : None,###分割模型权重位置
|
|
|
|
'postFile': '../AIlib2/weights/conf/forest/para.json',###后处理参数文件
|
|
|
|
'txtFontSize':80,###文本字符的大小
|
|
|
|
'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'waterLineWidth':3},###显示框、线设置
|
|
|
|
'testImgPath':'../AIdemo2/images/forest/',###测试图像的位置
|
|
|
|
'testOutPath':'images/results/',###输出测试图像位置
|
|
|
|
}
|
|
|
|
#
|
|
|
|
if opt['business'] == 'forest2':
|
|
|
|
par={
|
|
|
|
'device':'0', ###显卡号,如果用TRT模型，只支持0（单显卡）
|
|
|
|
'labelnames':"../AIlib2/weights/conf/forest2/labelnames.json", ###检测类别对照表
|
|
|
|
'gpuname':opt['gpu'],###显卡名称
|
|
|
|
'max_workers':1, ###并行线程数
|
|
|
|
'trtFlag_det':True,###检测模型是否采用TRT
|
|
|
|
'trtFlag_seg':False,###分割模型是否采用TRT
|
|
|
|
'Detweights':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###检测模型路径
|
|
|
|
#'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,4,5,6] ],###控制哪些检测类别显示、输出
|
|
|
|
'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [] ],
|
|
|
|
|
|
|
|
'slopeIndex':[],###岸坡类别（或者其它业务里的类别），不与河道（分割的前景区域）计算交并比，即不论是否在河道内都显示。
|
|
|
|
'seg_nclass':2,###分割模型类别数目，默认2类
|
|
|
|
'segRegionCnt':0,###分割模型结果需要保留的等值线数目
|
|
|
|
'segPar':None,###分割模型预处理参数
|
|
|
|
'Segweights' : None,###分割模型权重位置
|
|
|
|
'postFile': '../AIlib2/weights/conf/forest/para.json',###后处理参数文件
|
|
|
|
'txtFontSize':80,###文本字符的大小
|
|
|
|
'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'waterLineWidth':3},###显示框、线设置
|
|
|
|
'testImgPath':'../AIdemo2/images/forest2/',###测试图像的位置
|
|
|
|
'testOutPath':'images/results/',###输出测试图像位置
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
###车辆巡检参数
|
|
|
|
if opt['business'] == 'vehicle':
|
|
|
|
par={
|
|
|
|
'device':'0', ###显卡号,如果用TRT模型，只支持0（单显卡）
|
|
|
|
'labelnames':"../AIlib2/weights/conf/vehicle/labelnames.json", ###检测类别对照表
|
|
|
|
'gpuname':'2080T',###显卡名称
|
|
|
|
'max_workers':1, ###并行线程数
|
|
|
|
'trtFlag_det':True,###检测模型是否采用TRT
|
|
|
|
'trtFlag_seg':False,###分割模型是否采用TRT
|
|
|
|
'Detweights':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###检测模型路径
|
|
|
|
'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,4,5,6] ],###控制哪些检测类别显示、输出
|
|
|
|
'slopeIndex':[],###岸坡类别（或者其它业务里的类别），不与河道（分割的前景区域）计算交并比，即不论是否在河道内都显示。
|
|
|
|
'seg_nclass':2,###分割模型类别数目，默认2类
|
|
|
|
'segRegionCnt':0,###分割模型结果需要保留的等值线数目
|
|
|
|
'segPar':None,###分割模型预处理参数
|
|
|
|
'Segweights' : None,###分割模型权重位置
|
|
|
|
'postFile': '../AIlib2/weights/conf/vehicle/para.json',###后处理参数文件
|
|
|
|
'txtFontSize':40,###文本字符的大小
|
|
|
|
'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'waterLineWidth':3},###显示框、线设置
|
|
|
|
'testImgPath':'../AIdemo2/images/vehicle/',###测试图像的位置
|
|
|
|
'testOutPath':'images/results/',###输出测试图像位置
|
|
|
|
}
|
|
|
|
|
|
|
|
###行人检测模型
|
|
|
|
if opt['business'] == 'pedestrian':
|
|
|
|
par={
|
|
|
|
'device':'0', ###显卡号,如果用TRT模型，只支持0（单显卡）
|
|
|
|
'labelnames':"../AIlib2/weights/conf/pedestrian/labelnames.json", ###检测类别对照表
|
|
|
|
'gpuname':'2080T',###显卡名称
|
|
|
|
'max_workers':1, ###并行线程数
|
|
|
|
'trtFlag_det':True,###检测模型是否采用TRT
|
|
|
|
'trtFlag_seg':False,###分割模型是否采用TRT
|
|
|
|
'Detweights':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###检测模型路径
|
|
|
|
'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,4,5,6] ],###控制哪些检测类别显示、输出
|
|
|
|
'slopeIndex':[],###岸坡类别（或者其它业务里的类别），不与河道（分割的前景区域）计算交并比，即不论是否在河道内都显示。
|
|
|
|
'seg_nclass':2,###分割模型类别数目，默认2类
|
|
|
|
'segRegionCnt':0,###分割模型结果需要保留的等值线数目
|
|
|
|
'segPar':None,###分割模型预处理参数
|
|
|
|
'Segweights' : None,###分割模型权重位置
|
|
|
|
'postFile': '../AIlib2/weights/conf/pedestrian/para.json',###后处理参数文件
|
|
|
|
'txtFontSize':40,###文本字符的大小
|
|
|
|
'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'waterLineWidth':3},###显示框、线设置
|
|
|
|
'testImgPath':'../AIdemo2/images/pedestrian/',###测试图像的位置
|
|
|
|
'testOutPath':'images/results/',###输出测试图像位置
|
|
|
|
}
|
|
|
|
|
|
|
|
###烟雾火焰检测模型
|
|
|
|
if opt['business'] == 'smogfire':
|
|
|
|
par={
|
|
|
|
'device':'0', ###显卡号,如果用TRT模型，只支持0（单显卡）
|
|
|
|
'labelnames':"../AIlib2/weights/conf/smogfire/labelnames.json", ###检测类别对照表
|
|
|
|
'gpuname':'2080T',###显卡名称
|
|
|
|
'max_workers':1, ###并行线程数
|
|
|
|
'trtFlag_det':True,###检测模型是否采用TRT
|
|
|
|
'trtFlag_seg':False,###分割模型是否采用TRT
|
|
|
|
'Detweights':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###检测模型路径
|
|
|
|
'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,4,5,6] ],###控制哪些检测类别显示、输出
|
|
|
|
'slopeIndex':[],###岸坡类别（或者其它业务里的类别），不与河道（分割的前景区域）计算交并比，即不论是否在河道内都显示。
|
|
|
|
'seg_nclass':2,###没有分割模型，此处不用
|
|
|
|
'segRegionCnt':0,###没有分割模型，此处不用
|
|
|
|
'segPar':None,###分割模型预处理参数
|
|
|
|
'Segweights' : None,###分割模型权重位置
|
|
|
|
'postFile': '../AIlib2/weights/conf/smogfire/para.json',###后处理参数文件
|
|
|
|
'txtFontSize':40,###文本字符的大小
|
|
|
|
'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'waterLineWidth':3},###显示框、线设置
|
|
|
|
'testImgPath':'../AIdemo2/images/smogfire/',###测试图像的位置
|
|
|
|
'testOutPath':'images/results/',###输出测试图像位置
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
###钓鱼游泳检测
|
|
|
|
if opt['business'] == 'AnglerSwimmer':
|
|
|
|
par={
|
|
|
|
'device':'0', ###显卡号,如果用TRT模型，只支持0（单显卡）
|
|
|
|
'labelnames':"../AIlib2/weights/conf/AnglerSwimmer/labelnames.json", ###检测类别对照表
|
|
|
|
'gpuname':'2080T',###显卡名称
|
|
|
|
'max_workers':1, ###并行线程数
|
|
|
|
'trtFlag_det':True,###检测模型是否采用TRT
|
|
|
|
'trtFlag_seg':False,###分割模型是否采用TRT
|
|
|
|
'Detweights':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###检测模型路径
|
|
|
|
'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,4,5,6] ],###控制哪些检测类别显示、输出
|
|
|
|
'slopeIndex':[],###岸坡类别（或者其它业务里的类别），不与河道（分割的前景区域）计算交并比，即不论是否在河道内都显示。
|
|
|
|
'seg_nclass':2,###没有分割模型，此处不用
|
|
|
|
'segRegionCnt':0,###没有分割模型，此处不用
|
|
|
|
'segPar':None,###分割模型预处理参数
|
|
|
|
'Segweights' : None,###分割模型权重位置
|
|
|
|
'postFile': '../AIlib2/weights/conf/AnglerSwimmer/para.json',###后处理参数文件
|
|
|
|
'txtFontSize':40,###文本字符的大小
|
|
|
|
'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'waterLineWidth':3},###显示框、线设置
|
|
|
|
'testImgPath':'../AIdemo2/images/AnglerSwimmer/',###测试图像的位置
|
|
|
|
'testOutPath':'images/results/',###输出测试图像位置
|
|
|
|
}
|
|
|
|
###航道应急，做落水人员检测, channelEmergency
|
|
|
|
if opt['business'] == 'channelEmergency':
|
|
|
|
par={
|
|
|
|
'device':'0', ###显卡号,如果用TRT模型，只支持0（单显卡）
|
|
|
|
'labelnames':"../AIlib2/weights/conf/channelEmergency/labelnames.json", ###检测类别对照表
|
|
|
|
'gpuname':'2080T',###显卡名称
|
|
|
|
'max_workers':1, ###并行线程数
|
|
|
|
'trtFlag_det':True,###检测模型是否采用TRT
|
|
|
|
'trtFlag_seg':False,###分割模型是否采用TRT
|
|
|
|
'Detweights':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###检测模型路径
|
|
|
|
#'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,4,5,6] ],###控制哪些检测类别显示、输出
|
|
|
|
'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [] ],###控制哪些检测类别显示、输出
|
|
|
|
'slopeIndex':[],###岸坡类别（或者其它业务里的类别），不与河道（分割的前景区域）计算交并比，即不论是否在河道内都显示。
|
|
|
|
'seg_nclass':2,###没有分割模型，此处不用
|
|
|
|
'segRegionCnt':0,###没有分割模型，此处不用
|
|
|
|
'segPar':None,###分割模型预处理参数
|
|
|
|
'Segweights' : None,###分割模型权重位置
|
|
|
|
'postFile': '../AIlib2/weights/conf/channelEmergency/para.json',###后处理参数文件
|
|
|
|
'txtFontSize':40,###文本字符的大小
|
|
|
|
'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'waterLineWidth':3},###显示框、线设置
|
|
|
|
'testImgPath':'../AIdemo2/images/channelEmergency/',###测试图像的位置
|
|
|
|
'testOutPath':'images/results/',###输出测试图像位置
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
###乡村路违法种植
|
|
|
|
if opt['business'] == 'countryRoad':
|
|
|
|
par={
|
|
|
|
'device':'0', ###显卡号,如果用TRT模型，只支持0（单显卡）
|
|
|
|
'labelnames':"../AIlib2/weights/conf/countryRoad/labelnames.json", ###检测类别对照表
|
|
|
|
'gpuname':'2080T',###显卡名称
|
|
|
|
'max_workers':1, ###并行线程数
|
|
|
|
'trtFlag_det':True,###检测模型是否采用TRT
|
|
|
|
'trtFlag_seg':False,###分割模型是否采用TRT
|
|
|
|
'Detweights':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###检测模型路径
|
|
|
|
'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,4,5,6] ],###控制哪些检测类别显示、输出
|
|
|
|
'slopeIndex':[],###岸坡类别（或者其它业务里的类别），不与河道（分割的前景区域）计算交并比，即不论是否在河道内都显示。
|
|
|
|
'seg_nclass':2,###没有分割模型，此处不用
|
|
|
|
'segRegionCnt':0,###没有分割模型，此处不用
|
|
|
|
'segPar':None,###分割模型预处理参数
|
|
|
|
'Segweights' : None,###分割模型权重位置
|
|
|
|
'postFile': '../AIlib2/weights/conf/countryRoad/para.json',###后处理参数文件
|
|
|
|
'txtFontSize':40,###文本字符的大小
|
|
|
|
'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'waterLineWidth':3},###显示框、线设置
|
|
|
|
'testImgPath':'../AIdemo2/images/countryRoad/',###测试图像的位置
|
|
|
|
'testOutPath':'images/results/',###输出测试图像位置
|
|
|
|
}
|
|
|
|
###河道上大型船只
|
|
|
|
if opt['business'] == 'ship':
|
|
|
|
par={
|
|
|
|
'device':'0', ###显卡号,如果用TRT模型，只支持0（单显卡）
|
|
|
|
'labelnames':"../AIlib2/weights/conf/%s/labelnames.json"%(opt['business']), ###检测类别对照表
|
|
|
|
'gpuname':'2080T',###显卡名称
|
|
|
|
'max_workers':1, ###并行线程数
|
|
|
|
'trtFlag_det':True,###检测模型是否采用TRT
|
|
|
|
'trtFlag_seg':False,###分割模型是否采用TRT
|
|
|
|
'Detweights':"../weights/%s/AIlib2/%s/yolov5_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###检测模型路径
|
|
|
|
'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,4,5,6] ],###控制哪些检测类别显示、输出
|
|
|
|
'slopeIndex':[],###岸坡类别（或者其它业务里的类别），不与河道（分割的前景区域）计算交并比，即不论是否在河道内都显示。
|
|
|
|
'seg_nclass':2,###没有分割模型，此处不用
|
|
|
|
'segRegionCnt':0,###没有分割模型，此处不用
|
|
|
|
'segPar':None,###分割模型预处理参数
|
|
|
|
'Segweights' : None,###分割模型权重位置
|
|
|
|
'postFile': '../AIlib2/weights/conf/%s/para.json'%(opt['business']),###后处理参数文件
|
|
|
|
'txtFontSize':40,###文本字符的大小
|
|
|
|
'digitFont': { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'waterLineWidth':3},###显示框、线设置
|
|
|
|
'testImgPath':'../../../data/XunHe/shipData/',###测试图像的位置
|
|
|
|
'testOutPath':'images/results/',###输出测试图像位置
|
|
|
|
}
|
|
|
|
|
|
|
|
segRegionCnt=par['segRegionCnt']
|
|
|
|
trtFlag_seg = par['trtFlag_seg'];segPar=par['segPar']
|
|
|
|
##使用森林，道路模型,business 控制['forest','road']
|
|
|
|
##预先设置的参数
|
|
|
|
gpuname=par['gpuname']#如果用trt就需要此参数，只能是"3090" "2080Ti"
|
|
|
|
device_=par['device'] ##选定模型，可选 cpu,'0','1'
|
|
|
|
|
|
|
|
|
|
|
|
device = select_device(device_)
|
|
|
|
half = device.type != 'cpu' # half precision only supported on CUDA
|
|
|
|
trtFlag_det=par['trtFlag_det'] ###是否采用TRT模型加速
|
|
|
|
|
|
|
|
##以下参数目前不可改
|
|
|
|
imageW=1536 ####道路模型
|
|
|
|
digitFont= par['digitFont']
|
|
|
|
|
|
|
|
|
|
|
|
if trtFlag_det:
|
|
|
|
Detweights=par['Detweights']
|
|
|
|
logger = trt.Logger(trt.Logger.ERROR)
|
|
|
|
with open(Detweights, "rb") as f, trt.Runtime(logger) as runtime:
|
|
|
|
model=runtime.deserialize_cuda_engine(f.read())# 输入trt本地文件，返回ICudaEngine对象
|
|
|
|
print('####load TRT model :%s'%(Detweights))
|
|
|
|
else:
|
|
|
|
Detweights=par['Detweights']
|
|
|
|
model = attempt_load(Detweights, map_location=device) # load FP32 model
|
|
|
|
if half: model.half()
|
|
|
|
|
|
|
|
labelnames = par['labelnames']
|
|
|
|
postFile= par['postFile']
|
|
|
|
print( Detweights,labelnames )
|
|
|
|
conf_thres,iou_thres,classes,rainbows=get_postProcess_para(postFile)
|
|
|
|
|
|
|
|
####模型选择参数用如下：
|
|
|
|
mode_paras=par['detModelpara']
|
|
|
|
|
|
|
|
allowedList,allowedList_string=get_needed_objectsIndex(mode_paras)
|
|
|
|
slopeIndex = par['slopeIndex']
|
|
|
|
##只加载检测模型，准备好显示字符
|
|
|
|
|
|
|
|
names=get_labelnames(labelnames)
|
|
|
|
#imageW=4915;###默认是1920，在森林巡检的高清图像中是4920
|
|
|
|
outfontsize=int(imageW/1920*40);###
|
|
|
|
label_arraylist = get_label_arrays(names,rainbows,outfontsize=par['txtFontSize'],fontpath="../AIlib2/conf/platech.ttf")
|
|
|
|
|
|
|
|
segmodel = None
|
|
|
|
##图像测试
|
|
|
|
#url='images/examples/20220624_响水河_12300_1621.jpg'
|
|
|
|
impth = par['testImgPath']
|
|
|
|
outpth = par['testOutPath']
|
|
|
|
|
|
|
|
|
|
|
|
imgpaths=[]###获取文件里所有的图像
|
|
|
|
for postfix in ['.jpg','.JPG','.PNG','.png']:
|
|
|
|
imgpaths.extend(glob.glob('%s/*%s'%(impth,postfix )) )
|
|
|
|
videopaths=[]###获取文件里所有的视频
|
|
|
|
for postfix in ['.MP4','.mp4','.avi']:
|
|
|
|
videopaths.extend(glob.glob('%s/*%s'%(impth,postfix )) )
|
|
|
|
|
|
|
|
|
|
|
|
imgpaths.sort()
|
|
|
|
|
|
|
|
for i in range(len(imgpaths)):
|
|
|
|
#for i in range(2):
|
|
|
|
#imgpath = os.path.join(impth, folders[i])
|
|
|
|
imgpath = imgpaths[i]
|
|
|
|
bname = os.path.basename(imgpath )
|
|
|
|
im0s=[cv2.imread(imgpath)]
|
|
|
|
time00 = time.time()
|
|
|
|
#使用不同的函数。每一个领域采用一个函数
|
|
|
|
p_result,timeOut = AI_process_forest(im0s,model,segmodel,names,label_arraylist,rainbows,half,device,conf_thres, iou_thres,allowedList,font=digitFont,trtFlag_det=trtFlag_det)
|
|
|
|
|
|
|
|
time11 = time.time()
|
|
|
|
image_array = p_result[1]
|
|
|
|
cv2.imwrite( os.path.join( outpth,bname ) ,image_array )
|
|
|
|
|
|
|
|
print('----image:%s, process:%s ,save:%s, %s'%(bname,(time11-time00) * 1000, (time.time() - time11) * 1000,timeOut) , p_result[2] )
|
|
|
|
|
|
|
|
##process video
|
|
|
|
|
|
|
|
print('##begin to process videos, total %d videos'%( len(videopaths)))
|
|
|
|
for i,video in enumerate(videopaths):
|
|
|
|
print('process video%d :%s '%(i,video))
|
|
|
|
par0={'model':model,'segmodel':segmodel, 'names':names,'label_arraylist':label_arraylist,'rainbows':rainbows,'outpth':par['testOutPath'],
|
|
|
|
'half':half,'device':device,'conf_thres':conf_thres, 'iou_thres':iou_thres,'allowedList':allowedList,'digitFont':digitFont,'trtFlag_det': trtFlag_det
|
|
|
|
}
|
|
|
|
process_video(video,par0,mode='det')
|
|
|
|
|
|
|
|
|
|
|
|
def Seg_demo(opt):
|
|
|
|
###河道巡检的参数####
|
|
|
|
if opt['business'] == 'trafficAccident':
|
|
|
|
par={
|
|
|
|
'device':'0', ###显卡号,如果用TRT模型，只支持0（单显卡）
|
|
|
|
'labelnames':"../AIlib2/weights/conf/trafficAccident/labelnames.json", ###检测类别对照表
|
|
|
|
'gpuname':'3090',###显卡名称,只能是"3090" "2080Ti",'4090'
|
|
|
|
'max_workers':1, ###并行线程数
|
|
|
|
'trtFlag_seg':True,###分割模型是否采用TRT
|
|
|
|
#'seg_nclass':2,###分割模型类别数目，默认2类
|
|
|
|
'segPar':{'modelSize':(1440,810),'mean':(0.485, 0.456, 0.406),'std' :(0.229, 0.224, 0.225),'nclass':3,'predResize':False,'numpy':False, 'RGB_convert_first':True},###分割模型预处理参数
|
|
|
|
'postPar': {'label_csv': '../AIlib2/weights/conf/trafficAccident/class_dict.csv', 'speedRoadArea': 5100, 'vehicleArea': 400, 'speedRoadVehicleAngleMin': 15, 'speedRoadVehicleAngleMax': 75, 'vehicleLengthWidthThreshold': 12, 'vehicleSafeDistance': 7},
|
|
|
|
##'vehicleArea': 400， 'speedRoadArea': 5100---都是相对于1920*1080的大小。如果是其它分辨率要相应放缩
|
|
|
|
#'Segweights' : "../weights/%s/AIlib2/%s/stdc_810X1440_%s_fp16.engine"%(opt['gpu'], opt['business'] ,opt['gpu'] ),###分割模型权重位置
|
|
|
|
'Segweights' : "../weights/2080Ti/AIlib2/trafficAccident/stdc_1440X810_fp16_2080Ti.engine",
|
|
|
|
'postFile': '../AIlib2/weights/conf/trafficAccident/para.json',###后处理参数文件
|
|
|
|
'digitWordFont': { 'line_thickness':2,'boxLine_thickness':1,'wordSize':40, 'fontSize':1.0},
|
|
|
|
###line_thickness:类别汉字线宽，'fontSize':汉字相对大小 ，boxLine_thickness：矩形框线宽
|
|
|
|
'testImgPath':'images/trafficAccident2',
|
|
|
|
#'testImgPath':'/home/thsw2/WJ/data/XunHe/huData/Huzhou/original/',
|
|
|
|
#'images/river2/',
|
|
|
|
#'../../../data/无人机起飞测试图像/',###测试图像的位置
|
|
|
|
'testOutPath':'images/results/',###输出测试图像位置
|
|
|
|
#'testOutPath':'/home/thsw2/WJ/data/XunHe/huData/Huzhou/AIProcess',
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
max_workers=par['max_workers'];
|
|
|
|
trtFlag_seg=par['trtFlag_seg'];
|
|
|
|
device = select_device(par['device'])
|
|
|
|
names=get_labelnames(par['labelnames'])
|
|
|
|
|
|
|
|
#加载分割模型
|
|
|
|
seg_nclass = par['segPar']['nclass']
|
|
|
|
segPar=par['segPar']
|
|
|
|
if trtFlag_seg:
|
|
|
|
Segweights = par['Segweights']
|
|
|
|
logger = trt.Logger(trt.Logger.ERROR)
|
|
|
|
with open(Segweights, "rb") as f, trt.Runtime(logger) as runtime:
|
|
|
|
segmodel=runtime.deserialize_cuda_engine(f.read())# 输入trt本地文件，返回ICudaEngine对象
|
|
|
|
print('############locad seg model trt success#######')
|
|
|
|
else:
|
|
|
|
Segweights = par['Segweights']
|
|
|
|
segmodel = SegModel(nclass=seg_nclass,weights=Segweights,device=device)
|
|
|
|
print('############locad seg model pth success#######')
|
|
|
|
|
|
|
|
postFile= par['postFile']
|
|
|
|
digitWordFont= par['digitWordFont']
|
|
|
|
conf_thres,iou_thres,classes,rainbows=get_postProcess_para(postFile)
|
|
|
|
|
|
|
|
##加载模型，准备好显示字符
|
|
|
|
label_arraylist = get_label_arrays(names,rainbows,outfontsize=par['digitWordFont']['wordSize'],fontpath="../AIlib2/conf/platech.ttf")
|
|
|
|
digitWordFont['names']=names;digitWordFont['rainbows']=rainbows;digitWordFont['label_arraylist']=label_arraylist
|
|
|
|
wordFont={ 'fontSize':None,'boxLine_thickness':None,'waterLineColor':(0,255,255),'waterLineWidth':3}
|
|
|
|
|
|
|
|
|
|
|
|
##图像测试
|
|
|
|
#impth = 'images/slope/'
|
|
|
|
impth = par['testImgPath']
|
|
|
|
outpth = par['testOutPath']
|
|
|
|
imgpaths=[]###获取文件里所有的图像
|
|
|
|
for postfix in ['.jpg','.JPG','.PNG','.png']:
|
|
|
|
imgpaths.extend(glob.glob('%s/*%s'%(impth,postfix )) )
|
|
|
|
videopaths=[]###获取文件里所有的视频
|
|
|
|
for postfix in ['.MP4','.mp4','.avi']:
|
|
|
|
videopaths.extend(glob.glob('%s/*%s'%(impth,postfix )) )
|
|
|
|
|
|
|
|
##图像测试
|
|
|
|
#url='images/examples/20220624_响水河_12300_1621.jpg'
|
|
|
|
impth = par['testImgPath']
|
|
|
|
outpth = par['testOutPath']
|
|
|
|
folders = os.listdir(impth)
|
|
|
|
folders.sort()
|
|
|
|
|
|
|
|
for i in range(len(imgpaths)):
|
|
|
|
#for i in range(3,4):
|
|
|
|
imgpath = imgpaths[i]
|
|
|
|
bname = os.path.basename(imgpath)
|
|
|
|
im0s=[cv2.imread(imgpath)]
|
|
|
|
time00 = time.time()
|
|
|
|
seg_pred,img_draw,segstr,list1 = AI_Seg_process(im0s,segmodel,digitWordFont,trtFlag_seg,segPar,par['postPar'])
|
|
|
|
list1 = [ np.array(x) for x in list1]
|
|
|
|
|
|
|
|
cv2.imwrite( os.path.join(outpth, bname), img_draw)
|
|
|
|
ret=cv2.imwrite( os.path.join(outpth, bname.replace('.png','_mask.png')), (seg_pred*50).astype(np.uint8))
|
|
|
|
print(bname,segstr)
|
|
|
|
|
|
|
|
def OCR_demo(opt):
|
|
|
|
if opt['business'] == 'ocr':
|
|
|
|
par={
|
|
|
|
'image_dir':'images/ocr_en',
|
|
|
|
'outtxt':'images/results',
|
|
|
|
'TRTfile':'../weights/2080Ti/AIlib2/ocr_en/english_2080Ti_g2_h64_fp16.engine',
|
|
|
|
'device':'cuda:0',
|
|
|
|
'dict_list':{'en': '../AIlib2/weights/conf/ocr_en/en.txt'},
|
|
|
|
'char_file':'../AIlib2/weights/conf/ocr_en/en_character.csv',
|
|
|
|
'imgH':64,
|
|
|
|
'imgW':896,
|
|
|
|
'workers':1
|
|
|
|
|
|
|
|
|
|
|
|
}
|
|
|
|
image_dir=par['image_dir']
|
|
|
|
outtxt=par['outtxt']
|
|
|
|
workers=par['workers']
|
|
|
|
TRTfile= par['TRTfile']
|
|
|
|
device=par['device']
|
|
|
|
dict_list= par['dict_list']
|
|
|
|
char_file=par['char_file']
|
|
|
|
imgH=par['imgH']
|
|
|
|
imgW=par['imgW']
|
|
|
|
|
|
|
|
img_urls=glob.glob('%s/*.jpg'%( image_dir ))
|
|
|
|
img_urls.extend( glob.glob('%s/*.png'%( image_dir )) )
|
|
|
|
cnt=len(img_urls)
|
|
|
|
print('%s has %d images'%(image_dir ,len(img_urls) ) )
|
|
|
|
logger = trt.Logger(trt.Logger.ERROR)
|
|
|
|
with open(TRTfile, "rb") as f, trt.Runtime(logger) as runtime:
|
|
|
|
engine=runtime.deserialize_cuda_engine(f.read())# 输入trt本地文件，返回ICudaEngine对象
|
|
|
|
print('#####load TRT file:',TRTfile,'success #####')
|
|
|
|
context= engine.create_execution_context()
|
|
|
|
print('###line693:',type(context),context,dir(context))
|
|
|
|
|
|
|
|
with open(char_file,'r') as fp:
|
|
|
|
characters = fp.readlines()[0].strip()
|
|
|
|
converter = CTCLabelConverter(characters, {}, dict_list)
|
|
|
|
|
|
|
|
AlignCollate_normal = AlignCollate(imgH=imgH, imgW=imgW, keep_ratio_with_pad=True)
|
|
|
|
|
|
|
|
|
|
|
|
# 准备数据
|
|
|
|
parList=[]
|
|
|
|
for i in range(cnt):
|
|
|
|
img_patch=cv2.imread( img_urls[i] , cv2.IMREAD_GRAYSCALE)
|
|
|
|
time1 = time.time()
|
|
|
|
par=[img_patch,engine,context,converter,AlignCollate_normal,device]
|
|
|
|
parList.append(par)
|
|
|
|
|
|
|
|
if workers==1:
|
|
|
|
for i in range(len(parList)):
|
|
|
|
preds_str,info_str=ocr_process(parList[i])
|
|
|
|
print('filename:%s preds_str:%s %s'%( os.path.basename( img_urls[i] ), preds_str, info_str))
|
|
|
|
else:
|
|
|
|
with ThreadPoolExecutor(max_workers=workers) as t:
|
|
|
|
for result in t.map(processM, parList):
|
|
|
|
#print(result)
|
|
|
|
t=result
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def OCR_demo2(opt):
|
|
|
|
from ocrUtils2 import crnn_model
|
|
|
|
from ocrUtils2.ocrUtils import get_cfg,recognition_ocr,strLabelConverter
|
|
|
|
|
|
|
|
if opt['business'] == 'ocr2':
|
|
|
|
par={
|
|
|
|
'image_dir':'images/ocr_en',
|
|
|
|
'outtxt':'images/results',
|
|
|
|
'weights':'/mnt/thsw2/DSP2/weights/ocr2/checkpoint_99_acc_0.9737_448X32.engine',
|
|
|
|
#'weights':'/mnt/thsw2/DSP2/weights/ocr2/checkpoint_99_acc_0.9737.pth',
|
|
|
|
'device':'cuda:0',
|
|
|
|
'cfg':'/mnt/thsw2/DSP2/weights/ocr2/360CC_config.yaml',
|
|
|
|
|
|
|
|
'dict_list':{'en': '../AIlib2/weights/conf/ocr_en/en.txt'},
|
|
|
|
'char_file':'/mnt/thsw2/DSP2/weights/ocr2/chars.txt',
|
|
|
|
'imgH':32,
|
|
|
|
'imgW':448,
|
|
|
|
'workers':1
|
|
|
|
|
|
|
|
|
|
|
|
}
|
|
|
|
image_dir=par['image_dir']
|
|
|
|
outtxt=par['outtxt']
|
|
|
|
workers=par['workers']
|
|
|
|
weights= par['weights']
|
|
|
|
device=par['device']
|
|
|
|
dict_list= par['dict_list']
|
|
|
|
char_file=par['char_file']
|
|
|
|
imgH=par['imgH']
|
|
|
|
imgW=par['imgW']
|
|
|
|
cfg = par['cfg']
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
config = get_cfg(cfg, char_file)
|
|
|
|
|
|
|
|
|
|
|
|
par['contextFlag']=False
|
|
|
|
device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')
|
|
|
|
if weights.endswith('.pth'):
|
|
|
|
model = crnn_model.get_crnn(config,weights=weights).to(device)
|
|
|
|
par['model_mode']='pth'
|
|
|
|
else:
|
|
|
|
logger = trt.Logger(trt.Logger.ERROR)
|
|
|
|
with open(weights, "rb") as f, trt.Runtime(logger) as runtime:
|
|
|
|
model = runtime.deserialize_cuda_engine(f.read())# 输入trt本地文件，返回ICudaEngine对象
|
|
|
|
print('#####load TRT file:',weights,'success #####')
|
|
|
|
context = model.create_execution_context()
|
|
|
|
par['model_mode']='trt';par['contextFlag']=context
|
|
|
|
|
|
|
|
converter = strLabelConverter(config.DATASET.ALPHABETS)
|
|
|
|
|
|
|
|
img_urls=glob.glob('%s/*.jpg'%( image_dir ))
|
|
|
|
img_urls.extend( glob.glob('%s/*.png'%( image_dir )) )
|
|
|
|
cnt=len(img_urls)
|
|
|
|
print('%s has %d images'%(image_dir ,len(img_urls) ) )
|
|
|
|
# 准备数据
|
|
|
|
parList=[]
|
|
|
|
for i in range(cnt):
|
|
|
|
img_patch=cv2.imread( img_urls[i] , cv2.IMREAD_GRAYSCALE)
|
|
|
|
started = time.time()
|
|
|
|
img = cv2.imread(img_urls[i])
|
|
|
|
sim_pred = recognition_ocr(config, img, model, converter, device,par=par)
|
|
|
|
finished = time.time()
|
|
|
|
print('{0}: elapsed time: {1} prd:{2} '.format( os.path.basename( img_urls[i] ), finished - started, sim_pred ))
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def OBB_demo(opt):
|
|
|
|
###倾斜框（OBB）的ship目标检测
|
|
|
|
par={
|
|
|
|
'model_size':(608,608), #width,height
|
|
|
|
'K':100, #Maximum of objects'
|
|
|
|
'conf_thresh':0.18,##Confidence threshold, 0.1 for general evaluation
|
|
|
|
'device':"cuda:0",
|
|
|
|
|
|
|
|
'down_ratio':4,'num_classes':15,
|
|
|
|
#'weights':'../AIlib2/weights/conf/ship2/obb_608X608.engine',
|
|
|
|
'weights':'../weights/%s/AIlib2/%s/obb_608X608_%s_fp16.engine'%(opt['gpu'],opt['business'],opt['gpu']),
|
|
|
|
'dataset':'dota',
|
|
|
|
'test_dir': 'images/ship/',
|
|
|
|
'result_dir': 'images/results',
|
|
|
|
'half': False,
|
|
|
|
'mean':(0.5, 0.5, 0.5),
|
|
|
|
'std':(1, 1, 1),
|
|
|
|
'model_size':(608,608),##width,height
|
|
|
|
'heads': {'hm': None,'wh': 10,'reg': 2,'cls_theta': 1},
|
|
|
|
'decoder':None,
|
|
|
|
'test_flag':True,
|
|
|
|
'postFile': '../AIlib2/weights/conf/%s/para.json'%(opt['business'] ),###后处理参数文件
|
|
|
|
'drawBox':False,#####是否画框
|
|
|
|
'digitWordFont': { 'line_thickness':2,'boxLine_thickness':1,'wordSize':40, 'fontSize':1.0,'label_location':'leftTop'},
|
|
|
|
'labelnames':"../AIlib2/weights/conf/%s/labelnames.json"%(opt['business'] ), ###检测类别对照表
|
|
|
|
}
|
|
|
|
|
|
|
|
####加载模型
|
|
|
|
model,decoder2=load_model_decoder_OBB(par)
|
|
|
|
par['decoder']=decoder2
|
|
|
|
|
|
|
|
names=get_labelnames(par['labelnames']);par['labelnames']=names
|
|
|
|
conf_thres,iou_thres,classes,rainbows=get_postProcess_para(par['postFile']);par['rainbows']=rainbows
|
|
|
|
label_arraylist = get_label_arrays(names,rainbows,outfontsize=par['digitWordFont']['wordSize'],fontpath="../AIlib2/conf/platech.ttf")
|
|
|
|
par['label_array']=label_arraylist
|
|
|
|
|
|
|
|
img_urls=glob.glob('%s/*'%( par['test_dir'] ))
|
|
|
|
for img_url in img_urls:
|
|
|
|
#print(img_url)
|
|
|
|
ori_image=cv2.imread(img_url)
|
|
|
|
ori_image_list,infos = OBB_infer(model,ori_image,par)
|
|
|
|
imgName = os.path.basename(img_url)
|
|
|
|
saveFile = os.path.join(par['result_dir'], imgName)
|
|
|
|
ret=cv2.imwrite(saveFile, ori_image_list[1])
|
|
|
|
|
|
|
|
if not ret:
|
|
|
|
print(saveFile, ' not created ')
|
|
|
|
print( os.path.basename(img_url),':',infos,ori_image_list[2])
|
|
|
|
|
|
|
|
|
|
|
|
def jkm_demo():
|
|
|
|
from utilsK.jkmUtils import pre_process,post_process,get_return_data
|
|
|
|
img_type = 'plate' ## code,plate
|
|
|
|
par={'code':{'weights':'../AIlib2/weights/jkm/health_yolov5s_v3.jit','img_type':'code','nc':10 },
|
|
|
|
'plate':{'weights':'../AIlib2/weights/jkm/plate_yolov5s_v3.jit','img_type':'plate','nc':1 },
|
|
|
|
'conf_thres': 0.4,
|
|
|
|
'iou_thres':0.45,
|
|
|
|
'device':'cuda:0',
|
|
|
|
'plate_dilate':(0.5,0.1)
|
|
|
|
}
|
|
|
|
###加载模型
|
|
|
|
device = torch.device(par['device'])
|
|
|
|
jit_weights = par['code']['weights']
|
|
|
|
model = torch.jit.load(jit_weights)
|
|
|
|
|
|
|
|
jit_weights = par['plate']['weights']
|
|
|
|
model_plate = torch.jit.load(jit_weights)
|
|
|
|
|
|
|
|
imgd='images/plate'
|
|
|
|
imgpaths = os.listdir(imgd)
|
|
|
|
for imgp in imgpaths[0:]:
|
|
|
|
#imgp = 'plate_IMG_20221030_100612.jpg'
|
|
|
|
imgpath = os.path.join(imgd,imgp)
|
|
|
|
im0 = cv2.imread(imgpath) #读取数据
|
|
|
|
img ,padInfos = pre_process(im0,device) ##预处理
|
|
|
|
if img_type=='code': pred = model(img) ##模型推理
|
|
|
|
else: pred = model_plate(img)
|
|
|
|
boxes = post_process(pred,padInfos,device,conf_thres= par['conf_thres'], iou_thres= par['iou_thres'],nc=par[img_type]['nc']) #后处理
|
|
|
|
dataBack=get_return_data(im0,boxes,modelType=img_type,plate_dilate=par['plate_dilate'])
|
|
|
|
print(imgp,boxes,dataBack['type'])
|
|
|
|
for key in dataBack.keys():
|
|
|
|
if isinstance(dataBack[key],list):
|
|
|
|
cv2.imwrite( 'images/results/%s_%s.jpg'%( imgp.replace('.jpg','').replace('.png',''),key),dataBack[key][0] ) ###返回值： dataBack
|
|
|
|
|
|
|
|
if __name__=="__main__":
|
|
|
|
|
|
|
|
#jkm_demo()
|
|
|
|
businessAll=['river','road','AnglerSwimmer', 'countryRoad', 'forest','forest2', 'pedestrian' , 'smogfire' , 'vehicle',"ship",'ship2',"highWay2","highWay3","trafficAccident","channelEmergency"]
|
|
|
|
businessAll = ['channelEmergency']
|
|
|
|
|
|
|
|
for busi in businessAll:
|
|
|
|
opt={'gpu':'2080Ti','business':busi}
|
|
|
|
if opt['business'] in ['river','road','highWay2']:
|
|
|
|
detSeg_demo(opt)
|
|
|
|
elif opt['business'] in ['trafficAccident']:
|
|
|
|
Seg_demo(opt)
|
|
|
|
elif opt['business'] in ['ship2']:
|
|
|
|
OBB_demo(opt)
|
|
|
|
elif opt['business'] in ['ocr']:
|
|
|
|
OCR_demo(opt)
|
|
|
|
elif opt['business'] in ['ocr2']:
|
|
|
|
OCR_demo2(opt)
|
|
|
|
else:
|
|
|
|
det_demo( opt )
|
|
|
|
|
|
|
|
|
wangjin0928
1 anno fa
