diff --git a/demo.py b/demo.py index fed7e39..4c03b79 100644 --- a/demo.py +++ b/demo.py @@ -1,36 +1,81 @@ import sys +from concurrent.futures import ThreadPoolExecutor + sys.path.extend(['..','../AIlib' ]) -from AI import AI_process,AI_process_v2,AI_process_forest,get_postProcess_para +from AI import AI_process,AI_process_forest,get_postProcess_para import cv2,os,time from segutils.segmodel import SegModel -from segutils.segmodel_trt import SegModel_STDC_trt -from segutils.trtUtils import DetectMultiBackend from models.experimental import attempt_load from utils.torch_utils import select_device from utilsK.queRiver import get_labelnames,get_label_arrays import numpy as np import torch +import tensorrt as trt from utilsK.masterUtils import get_needed_objectsIndex -def river_demo_v2(): + +''' + 多线程 +''' + +def process_v1(frame): + #try: + + 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],frame[6],frame[7],frame[8], frame[9],frame[10],font=frame[11],trtFlag_det=frame[13],trtFlag_seg=frame[14],segPar=frame[15]) + time11 = time.time() + image_array = p_result[1] + cv2.imwrite(os.path.join('images/results/',frame[12] ) ,image_array) + time22 = time.time() + print('%s,%d*%d,AI-process: %.1f,image save:%.1f , %s'%(frame[12],H,W, (time11 - time00) * 1000.0, (time22-time11)*1000.0,timeOut)) + return 'success' + #except Exception as e: + # return 'failed:'+str(e) + +def river_demo_v3(): ##预先设置的参数 device_='0' ##选定模型，可选 cpu,'0','1' - ###注意TRT模型生成时，就需要对应cuda device，下面的trt文件是cuda:0生成的，device只能是0 + ###注意TRT模型生成时，就需要对应cuda device，下面的trt文件是cuda:0生成的，device只能是0 ##以下参数目前不可改 - - gpuname='3090'###3090,2080Ti. 只能是其中之一 - Detweights = "../AIlib/weights/yolov5/class8/bestcao_%s_fp16.engine"%(gpuname) ##升级后的检测模型 labelnames = "../AIlib/weights/yolov5/class8/labelnames.json" ##对应类别表 + gpuname='3090'; + max_workers=1; + trtFlag_det=True;trtFlag_seg=True + device = select_device(device_) + names=get_labelnames(labelnames) + + half = device.type != 'cpu' # half precision only supported on CUDA + if trtFlag_det: + Detweights = "../AIlib/weights/yolov5/class8/bestcao_%s_fp16.engine"%(gpuname) ##升级后的检测模型 + 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 trt success#######') + else: + Detweights = "../AIlib/weights/yolov5/class8/bestcao.pt" + model = attempt_load(Detweights, map_location=device) # load FP32 model + print('############locad det model pth success#######') + if half: model.half() + + + seg_nclass = 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} + if trtFlag_seg: + Segweights = '../AIlib/weights/STDC/model_maxmIOU75_1720_0.946_360640_%s_fp16.engine'%(gpuname) + 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 = '../AIlib/weights/STDC/model_maxmIOU75_1720_0.946_360640.pth' + segmodel = SegModel(nclass=seg_nclass,weights=Segweights,device=device) + print('############locad seg model pth success#######') - - seg_nclass = 2 - #Segweights = "../AIlib/weights/BiSeNet/checkpoint.pth" - Segweights = '../AIlib/weights/STDC/model_maxmIOU75_1720_0.946_360640_%s_fp16.engine'%(gpuname) ##升级的分割模型 - postFile= '../AIlib/conf/para.json' - digitFont= { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'waterLineWidth':3} + digitFont= { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'waterLineWidth':3} conf_thres,iou_thres,classes,rainbows=get_postProcess_para(postFile) - + ####模型选择参数用如下： mode_paras=[ {"id":"0","config":{"k1":"v1","k2":"v2"}}, @@ -42,34 +87,40 @@ def river_demo_v2(): {"id":"6","config":{"k1":"v1","k2":"v2"}}, {"id":"7","config":{"k1":"v1","k2":"v2"}}, ] - allowedList,allowedList_string=get_needed_objectsIndex(mode_paras) #allowedList=[0,1,2,3] ##加载模型，准备好显示字符 - device = select_device(device_) - names=get_labelnames(labelnames) - label_arraylist = get_label_arrays(names,rainbows,outfontsize=40,fontpath="../AIlib/conf/platech.ttf") - half = device.type != 'cpu' # half precision only supported on CUDA - - segmodel = SegModel_STDC_trt(weights=Segweights,modelsize=(640,360),std=(0.229, 0.224, 0.225),mean=(0.485, 0.456, 0.406),device=device) - model = DetectMultiBackend(Detweights, device=device) + ##图像测试 - #url='images/examples/20220624_响水河_12300_1621.jpg' - impth = 'images/slope/' + #impth = 'images/slope/' + impth = '../../../data/无人机起飞测试图像/' outpth = 'images/results/' folders = os.listdir(impth) + frames=[] + for i in range(len(folders)): imgpath = os.path.join(impth, folders[i]) im0s=[cv2.imread(imgpath)] - H,W,C = im0s[0].shape - time00 = time.time() - p_result,timeOut = AI_process_v2(im0s,model,segmodel,names,label_arraylist,rainbows,half,device,conf_thres, iou_thres,allowedList,font=digitFont) - time11 = time.time() - image_array = p_result[1] - cv2.imwrite( os.path.join( outpth,folders[i] ) ,image_array ) - print('%s,%d*%d,AI-process: %.1f, %s'%(folders[i],H,W, (time11 - time00) * 1000,timeOut)) + frame=(im0s,model,segmodel,names,label_arraylist,rainbows,half,device,conf_thres, iou_thres,allowedList,digitFont,folders[i],trtFlag_det,trtFlag_seg,segPar) + frames.append(frame) + t1=time.time() + if max_workers==1: + for i in range(len(folders)): + t5=time.time() + process_v1(frames[i]) + t6=time.time() + print('#######%s, ms:%.1f , accumetate time:%.1f, avage:%1.f '%(folders[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() + print('All %d images time:%.1f ms ,each:%.1f ms, with %d threads'%(len(folders),(t2-t1)*1000, (t2-t1)*1000.0/len(folders) , max_workers) ) + def river_demo(): ##预先设置的参数 @@ -140,58 +191,50 @@ def road_forest_demo(business ): ##使用森林，道路模型,business 控制['forest','road'] ##预先设置的参数 - device_='1' ##选定模型，可选 cpu,'0','1' - - + gpuname='3090'#如果用trt就需要此参数，只能是"3090" "2080Ti" + device_='0' ##选定模型，可选 cpu,'0','1' + device = select_device(device_) + half = device.type != 'cpu' # half precision only supported on CUDA + + trtFlag_det=False ###是否采用TRT模型加速 ##以下参数目前不可改 #business='forest';imageW=4916 ####森林模型 #business='road'; imageW=1536 ####道路模型 - digitFont= { 'line_thickness':2, 'fontSize':1.0} ###数字显示的线宽度，大小; 如果都是None，则采用默认大小 - Detweights="../AIlib/weights/%s/best.pt"%(business) + digitFont= { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'waterLineWidth':3 } + + + if trtFlag_det: + Detweights="../AIlib/weights/%s/best_%s_fp16.engine"%(business,gpuname) + 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="../AIlib/weights/%s/best.pt"%(business) + model = attempt_load(Detweights, map_location=device) # load FP32 model + if half: model.half() + labelnames = "../AIlib/weights/%s/labelnames.json"%(business) - - postFile= '../AIlib/conf/para.json' + postFile= '../AIlib/weights/%s/para.json'%(business) print( Detweights,labelnames ) conf_thres,iou_thres,classes,rainbows=get_postProcess_para(postFile) ####模型选择参数用如下： - mode_paras=[ - { - "id":"0", - "config":{ - "k1":"v1", - "k2":"v2" - } - }, - { - "id":"1", - "config":{ - "k1":"v1", - "k2":"v2" - } - } - ] - - allowedList,allowedList_string=get_needed_objectsIndex(mode_paras) - #allowedList=[0,1,2,3] - - print('####line108###') - + if business == 'road': + mode_paras=[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,3,4,5,6] ]###类别2为“修补”，不输出 + else: + mode_paras=[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,4,5,6] ] + allowedList,allowedList_string=get_needed_objectsIndex(mode_paras) ##只加载检测模型，准备好显示字符 - device = select_device(device_) + names=get_labelnames(labelnames) #imageW=4915;###默认是1920，在森林巡检的高清图像中是4920 outfontsize=int(imageW/1920*40);### label_arraylist = get_label_arrays(names,rainbows,outfontsize=outfontsize,fontpath="../AIlib/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 = None - - - ##图像测试 #url='images/examples/20220624_响水河_12300_1621.jpg' impth = 'images/%s/'%(business) @@ -204,7 +247,7 @@ def road_forest_demo(business ): 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) + 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,folders[i] ) ,image_array ) @@ -246,9 +289,9 @@ def jkm_demo(): if isinstance(dataBack[key],list): cv2.imwrite( 'images/results/%s_%s.jpg'%( imgp.replace('.jpg','').replace('.png',''),key),dataBack[key][0] ) ###返回值： dataBack - + 1 if __name__=="__main__": - river_demo_v2() - #road_forest_demo('forest' ) + #river_demo_v3() + road_forest_demo('road' ) #jkm_demo() diff --git a/demo_cyk.py b/demo_cyk.py deleted file mode 100644 index 89037af..0000000 --- a/demo_cyk.py +++ /dev/null @@ -1,381 +0,0 @@ -import sys -from concurrent.futures import ThreadPoolExecutor - -sys.path.extend(['..','../AIlib' ]) - -from AI import AI_process,AI_process_v2,AI_process_forest,get_postProcess_para -import cv2,os,time -from segutils.segmodel import SegModel -from segutils.segmodel_trt import SegModel_STDC_trt -from segutils.trtUtils import DetectMultiBackend -from models.experimental import attempt_load -from utils.torch_utils import select_device -from utilsK.queRiver import get_labelnames,get_label_arrays -import numpy as np -import torch -import tensorrt as trt -from utilsK.masterUtils import get_needed_objectsIndex - -''' - 多线程 -''' - -def process_v1(frame): - #try: - - 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],frame[6],frame[7],frame[8], frame[9],frame[10],font=frame[11],trtFlag_det=frame[13],trtFlag_seg=frame[14],segPar=frame[15]) - time11 = time.time() - image_array = p_result[1] - #cv2.imwrite(os.path.join('images/results/',frame[12] ) ,image_array) - time22 = time.time() - print('%s,%d*%d,AI-process: %.1f,image save:%.1f , %s'%(frame[12],H,W, (time11 - time00) * 1000.0, (time22-time11)*1000.0,timeOut)) - return 'success' - #except Exception as e: - # return 'failed:'+str(e) - -def river_demo_v3(): - ##预先设置的参数 - device_='0' ##选定模型，可选 cpu,'0','1' - ###注意TRT模型生成时，就需要对应cuda device，下面的trt文件是cuda:0生成的，device只能是0 - ##以下参数目前不可改 - labelnames = "../AIlib/weights/yolov5/class8/labelnames.json" ##对应类别表 - gpuname='3090'; - max_workers=6; - trtFlag_det=True;trtFlag_seg=True - device = select_device(device_) - names=get_labelnames(labelnames) - - half = device.type != 'cpu' # half precision only supported on CUDA - if trtFlag_det: - Detweights = "../AIlib/weights/yolov5/class8/bestcao_%s_fp16.engine"%(gpuname) ##升级后的检测模型 - - logger = trt.Logger(trt.Logger.INFO) - with open(Detweights, "rb") as f, trt.Runtime(logger) as runtime: - model=runtime.deserialize_cuda_engine(f.read())# 输入trt本地文件，返回ICudaEngine对象 - print('############locad det model trt success#######') - else: - Detweights = "../AIlib/weights/yolov5/class8/bestcao.pt" - model = attempt_load(Detweights, map_location=device) # load FP32 model - print('############locad det model pth success#######') - if half: model.half() - - - seg_nclass = 2 - #Segweights = "../AIlib/weights/BiSeNet/checkpoint.pth" - segPar={'modelSize':(640,360),'mean':(0.485, 0.456, 0.406),'std' :(0.229, 0.224, 0.225),'numpy':False, 'RGB_convert_first':True} - if trtFlag_seg: - Segweights = '../AIlib/weights/STDC/model_maxmIOU75_1720_0.946_360640_%s_fp16.engine'%(gpuname) ##升级的分割模型 - segmodel = SegModel_STDC_trt(weights=Segweights,modelsize=(640,360),std=(0.229, 0.224, 0.225),mean=(0.485, 0.456, 0.406),device=device) - - logger = trt.Logger(trt.Logger.INFO) - 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 = '../AIlib/weights/STDC/model_maxmIOU75_1720_0.946_360640.pth' - segmodel = SegModel(nclass=seg_nclass,weights=Segweights,device=device) - print('############locad seg model pth success#######') - - postFile= '../AIlib/conf/para.json' - digitFont= { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'waterLineWidth':3} - conf_thres,iou_thres,classes,rainbows=get_postProcess_para(postFile) - - ####模型选择参数用如下： - mode_paras=[ - {"id":"0","config":{"k1":"v1","k2":"v2"}}, - {"id":"1","config":{"k1":"v1","k2":"v2"}}, - {"id":"2","config":{"k1":"v1","k2":"v2"}}, - {"id":"3","config":{"k1":"v1","k2":"v2"}}, - {"id":"4","config":{"k1":"v1","k2":"v2"}}, - {"id":"5","config":{"k1":"v1","k2":"v2"}}, - {"id":"6","config":{"k1":"v1","k2":"v2"}}, - {"id":"7","config":{"k1":"v1","k2":"v2"}}, - ] - allowedList,allowedList_string=get_needed_objectsIndex(mode_paras) - #allowedList=[0,1,2,3] - ##加载模型，准备好显示字符 - label_arraylist = get_label_arrays(names,rainbows,outfontsize=40,fontpath="../AIlib/conf/platech.ttf") - - - ##图像测试 - impth = 'images/slope/' - outpth = 'images/results/' - folders = os.listdir(impth) - frames=[] - - for i in range(len(folders)): - imgpath = os.path.join(impth, folders[i]) - im0s=[cv2.imread(imgpath)] - frame=(im0s,model,segmodel,names,label_arraylist,rainbows,half,device,conf_thres, iou_thres,allowedList,digitFont,folders[i],trtFlag_det,trtFlag_seg,segPar) - frames.append(frame) - t1=time.time() - if max_workers==1: - for i in range(len(folders)): - t5=time.time() - process_v1(frames[i]) - t6=time.time() - print('#######%s, ms:%.1f , accumetate time:%.1f, avage:%1.f '%(folders[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() - print('All %d images time:%.1f ms ,each:%.1f ms, with %d threads'%(len(folders),(t2-t1)*1000, (t2-t1)*1000.0/len(folders) , max_workers) ) - -def river_demo_v2(): - ##预先设置的参数 - device_='0' ##选定模型，可选 cpu,'0','1' - ###注意TRT模型生成时，就需要对应cuda device，下面的trt文件是cuda:0生成的，device只能是0 - ##以下参数目前不可改 - - trtFlag=False - gpuname='3090' - if trtFlag: - Detweights = "../AIlib/weights/yolov5/class8/bestcao_%s_fp16.engine"%(gpuname) ##升级后的检测模型 - else: - Detweights = "../AIlib/weights/yolov5/class8/bestcao.pt" ##升级后的检测模型 - - labelnames = "../AIlib/weights/yolov5/class8/labelnames.json" ##对应类别表 - - - seg_nclass = 2 - #Segweights = "../AIlib/weights/BiSeNet/checkpoint.pth" - if trtFlag: - Segweights = '../AIlib/weights/STDC/model_maxmIOU75_1720_0.946_360640_%s_fp16.engine'%(gpuname) ##升级的分割模型 - else: - Segweights = '../AIlib/weights/STDC/model_maxmIOU75_1720_0.946_360640.pth' - - postFile= '../AIlib/conf/para.json' - digitFont= { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'waterLineWidth':3} - conf_thres,iou_thres,classes,rainbows=get_postProcess_para(postFile) - - ####模型选择参数用如下： - mode_paras=[ - {"id":"0","config":{"k1":"v1","k2":"v2"}}, - {"id":"1","config":{"k1":"v1","k2":"v2"}}, - {"id":"2","config":{"k1":"v1","k2":"v2"}}, - {"id":"3","config":{"k1":"v1","k2":"v2"}}, - {"id":"4","config":{"k1":"v1","k2":"v2"}}, - {"id":"5","config":{"k1":"v1","k2":"v2"}}, - {"id":"6","config":{"k1":"v1","k2":"v2"}}, - {"id":"7","config":{"k1":"v1","k2":"v2"}}, - ] - - allowedList,allowedList_string=get_needed_objectsIndex(mode_paras) - #allowedList=[0,1,2,3] - ##加载模型，准备好显示字符 - device = select_device(device_) - names=get_labelnames(labelnames) - - label_arraylist = get_label_arrays(names,rainbows,outfontsize=40,fontpath="../AIlib/conf/platech.ttf") - half = device.type != 'cpu' # half precision only supported on CUDA - - if trtFlag: - segmodel = SegModel_STDC_trt(weights=Segweights,modelsize=(640,360),std=(0.229, 0.224, 0.225),mean=(0.485, 0.456, 0.406),device=device) - model = DetectMultiBackend(Detweights, device=device) - else: - 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/slope/' - outpth = 'images/results/' - folders = os.listdir(impth) - for i in range(len(folders)): - imgpath = os.path.join(impth, folders[i]) - im0s=[cv2.imread(imgpath)] - H,W,C = im0s[0].shape - time00 = time.time() - p_result,timeOut = AI_process_v2(im0s,model,segmodel,names,label_arraylist,rainbows,half,device,conf_thres, iou_thres,allowedList,font=digitFont) - time11 = time.time() - image_array = p_result[1] - cv2.imwrite( os.path.join( outpth,folders[i] ) ,image_array ) - print('%s,%d*%d,AI-process: %.1f, %s'%(folders[i],H,W, (time11 - time00) * 1000,timeOut)) - -def river_demo(): - ##预先设置的参数 - device_='1' ##选定模型，可选 cpu,'0','1' - - ##以下参数目前不可改 - #Detweights = "../AIlib/weights/yolov5/class5/best_5classes.pt" - #labelnames = "../AIlib/weights/yolov5/class5/labelnames.json" - - Detweights = "../AIlib/weights/yolov5/class8/bestcao.pt" ##升级后的检测模型 - labelnames = "../AIlib/weights/yolov5/class8/labelnames.json" ##对应类别表 - - seg_nclass = 2 - #Segweights = "../AIlib/weights/BiSeNet/checkpoint.pth" - Segweights = '../AIlib/weights/STDC/model_maxmIOU75_1720_0.946_360640.pth' ##升级的分割模型 - - postFile= '../AIlib/conf/para.json' - digitFont= { 'line_thickness':2,'boxLine_thickness':1, 'fontSize':1.0,'waterLineColor':(0,255,255),'waterLineWidth':3} - conf_thres,iou_thres,classes,rainbows=get_postProcess_para(postFile) - - ####模型选择参数用如下： - mode_paras=[ - {"id":"0","config":{"k1":"v1","k2":"v2"}}, - {"id":"1","config":{"k1":"v1","k2":"v2"}}, - {"id":"2","config":{"k1":"v1","k2":"v2"}}, - {"id":"3","config":{"k1":"v1","k2":"v2"}}, - {"id":"4","config":{"k1":"v1","k2":"v2"}}, - {"id":"5","config":{"k1":"v1","k2":"v2"}}, - {"id":"6","config":{"k1":"v1","k2":"v2"}}, - {"id":"7","config":{"k1":"v1","k2":"v2"}}, - ] - - allowedList,allowedList_string=get_needed_objectsIndex(mode_paras) - #allowedList=[0,1,2,3] - - - - ##加载模型，准备好显示字符 - device = select_device(device_) - names=get_labelnames(labelnames) - - label_arraylist = get_label_arrays(names,rainbows,outfontsize=40,fontpath="../AIlib/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/slope/' - outpth = 'images/results/' - folders = os.listdir(impth) - for i in range(len(folders)): - imgpath = os.path.join(impth, folders[i]) - im0s=[cv2.imread(imgpath)] - H,W,C = im0s[0].shape - time00 = time.time() - p_result,timeOut = AI_process(im0s,model,segmodel,names,label_arraylist,rainbows,half,device,conf_thres, iou_thres,allowedList,font=digitFont) - time11 = time.time() - image_array = p_result[1] - cv2.imwrite( os.path.join( outpth,folders[i] ) ,image_array ) - print('%s,%d*%d,AI-process: %.1f, %s'%(folders[i],H,W, (time11 - time00) * 1000,timeOut)) - -def road_forest_demo(business ): - - ##使用森林，道路模型,business 控制['forest','road'] - ##预先设置的参数 - device_='1' ##选定模型，可选 cpu,'0','1' - - - - ##以下参数目前不可改 - #business='forest';imageW=4916 ####森林模型 - #business='road'; - imageW=1536 ####道路模型 - digitFont= { 'line_thickness':2, 'fontSize':1.0} ###数字显示的线宽度，大小; 如果都是None，则采用默认大小 - Detweights="../AIlib/weights/%s/best.pt"%(business) - labelnames = "../AIlib/weights/%s/labelnames.json"%(business) - - postFile= '../AIlib/conf/para.json' - print( Detweights,labelnames ) - conf_thres,iou_thres,classes,rainbows=get_postProcess_para(postFile) - - ####模型选择参数用如下： - mode_paras=[ - { - "id":"0", - "config":{ - "k1":"v1", - "k2":"v2" - } - }, - { - "id":"1", - "config":{ - "k1":"v1", - "k2":"v2" - } - } - ] - - allowedList,allowedList_string=get_needed_objectsIndex(mode_paras) - #allowedList=[0,1,2,3] - - print('####line108###') - - ##只加载检测模型，准备好显示字符 - device = select_device(device_) - names=get_labelnames(labelnames) - #imageW=4915;###默认是1920，在森林巡检的高清图像中是4920 - outfontsize=int(imageW/1920*40);### - label_arraylist = get_label_arrays(names,rainbows,outfontsize=outfontsize,fontpath="../AIlib/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 = None - - - - ##图像测试 - #url='images/examples/20220624_响水河_12300_1621.jpg' - impth = 'images/%s/'%(business) - outpth = 'images/results/' - folders = os.listdir(impth) - folders.sort() - for i in range(len(folders)): - #for i in range(2): - imgpath = os.path.join(impth, folders[i]) - 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) - time11 = time.time() - image_array = p_result[1] - cv2.imwrite( os.path.join( outpth,folders[i] ) ,image_array ) - - print('----image:%s, process:%s ,save:%s, %s'%(folders[i],(time11-time00) * 1000, (time.time() - time11) * 1000,timeOut) ) - - -def jkm_demo(): - from utilsK.jkmUtils import pre_process,post_process,get_return_data - img_type = 'plate' ## code,plate - par={'code':{'weights':'../AIlib/weights/jkm/health_yolov5s_v3.jit','img_type':'code','nc':10 }, - 'plate':{'weights':'../AIlib/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 - - 1 -if __name__=="__main__": - river_demo_v3() - #road_forest_demo('forest' ) - #jkm_demo() - diff --git a/images/forest/DJI_20220531130118_0298_2_c1.JPG b/images/forest/DJI_20220531130118_0298_2_c1.JPG new file mode 100755 index 0000000..a97a93d Binary files /dev/null and b/images/forest/DJI_20220531130118_0298_2_c1.JPG differ diff --git a/images/forest/DJI_20220531130118_0298_2_c2.JPG b/images/forest/DJI_20220531130118_0298_2_c2.JPG new file mode 100755 index 0000000..a97a93d Binary files /dev/null and b/images/forest/DJI_20220531130118_0298_2_c2.JPG differ diff --git a/images/forest/DJI_20220531130118_0298_2_c3.JPG b/images/forest/DJI_20220531130118_0298_2_c3.JPG new file mode 100755 index 0000000..a97a93d Binary files /dev/null and b/images/forest/DJI_20220531130118_0298_2_c3.JPG differ diff --git a/images/forest/DJI_20220531130118_0298_2_c4.JPG b/images/forest/DJI_20220531130118_0298_2_c4.JPG new file mode 100755 index 0000000..a97a93d Binary files /dev/null and b/images/forest/DJI_20220531130118_0298_2_c4.JPG differ diff --git a/images/forest/DJI_20220531130118_0298_2_c5.JPG b/images/forest/DJI_20220531130118_0298_2_c5.JPG new file mode 100755 index 0000000..a97a93d Binary files /dev/null and b/images/forest/DJI_20220531130118_0298_2_c5.JPG differ diff --git a/images/road/c1.jpg b/images/road/c1.jpg deleted file mode 100644 index 60e40d0..0000000 Binary files a/images/road/c1.jpg and /dev/null differ diff --git a/images/road/c2.jpg b/images/road/c2.jpg deleted file mode 100644 index d6a0d6e..0000000 Binary files a/images/road/c2.jpg and /dev/null differ diff --git a/images/road/c3.jpg b/images/road/c3.jpg deleted file mode 100644 index 7f7439e..0000000 Binary files a/images/road/c3.jpg and /dev/null differ diff --git a/images/road/c4.jpg b/images/road/c4.jpg deleted file mode 100644 index 61009b8..0000000 Binary files a/images/road/c4.jpg and /dev/null differ diff --git a/images/road/pro2th_101.jpg b/images/road/pro2th_101.jpg new file mode 100644 index 0000000..f470468 Binary files /dev/null and b/images/road/pro2th_101.jpg differ diff --git a/images/road/pro2th_11.jpg b/images/road/pro2th_11.jpg new file mode 100644 index 0000000..1f7b902 Binary files /dev/null and b/images/road/pro2th_11.jpg differ diff --git a/images/road/pro2th_17.jpg b/images/road/pro2th_17.jpg new file mode 100644 index 0000000..0238a4f Binary files /dev/null and b/images/road/pro2th_17.jpg differ diff --git a/images/road/pro2th_21.jpg b/images/road/pro2th_21.jpg new file mode 100644 index 0000000..949a39c Binary files /dev/null and b/images/road/pro2th_21.jpg differ diff --git a/images/road/pro2th_57.jpg b/images/road/pro2th_57.jpg new file mode 100644 index 0000000..149e8d6 Binary files /dev/null and b/images/road/pro2th_57.jpg differ diff --git a/images/road/pro2th_60.jpg b/images/road/pro2th_60.jpg new file mode 100644 index 0000000..48287fc Binary files /dev/null and b/images/road/pro2th_60.jpg differ diff --git a/images/road/pro2th_64.jpg b/images/road/pro2th_64.jpg new file mode 100644 index 0000000..a1cbe28 Binary files /dev/null and b/images/road/pro2th_64.jpg differ diff --git a/images/road/pro2th_79.jpg b/images/road/pro2th_79.jpg new file mode 100644 index 0000000..a3e529e Binary files /dev/null and b/images/road/pro2th_79.jpg differ diff --git a/images/road/pro2th_9.jpg b/images/road/pro2th_9.jpg new file mode 100644 index 0000000..ef703a9 Binary files /dev/null and b/images/road/pro2th_9.jpg differ diff --git a/images/road/pro2th_99.jpg b/images/road/pro2th_99.jpg new file mode 100644 index 0000000..d5eb5cf Binary files /dev/null and b/images/road/pro2th_99.jpg differ diff --git a/readme.md b/readme.md index b2a30ac..d663ddf 100644 --- a/readme.md +++ b/readme.md @@ -18,3 +18,9 @@ fp16,trt:yolov5,torch1.8 :4.1 ms STDC:1.2 ms fp16,pth:yolovt,torch1.8 :9.6 ms STDC:4.6 ms fp16,pth:yolov5,torch1.7 :6.0ms, STDC:7.2 ms + +2022.12.13 +1.修正trt调用方式 +2.道路模型换成新的 +3.road,forest换成TRT +