diff --git a/demo.py b/demo.py index 5d27d2d..fed7e39 100644 --- a/demo.py +++ b/demo.py @@ -18,7 +18,7 @@ def river_demo_v2(): ###注意TRT模型生成时,就需要对应cuda device,下面的trt文件是cuda:0生成的,device只能是0 ##以下参数目前不可改 - gpuname='3090' + gpuname='3090'###3090,2080Ti. 只能是其中之一 Detweights = "../AIlib/weights/yolov5/class8/bestcao_%s_fp16.engine"%(gpuname) ##升级后的检测模型 labelnames = "../AIlib/weights/yolov5/class8/labelnames.json" ##对应类别表 diff --git a/demo_cyk.py b/demo_cyk.py new file mode 100644 index 0000000..89037af --- /dev/null +++ b/demo_cyk.py @@ -0,0 +1,381 @@ +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() +