trt accerate

demo.py

@@ -1,18 +1,77 @@
 import sys
 sys.path.extend(['..','../AIlib' ])
 
-from AI import AI_process,AI_process_forest,get_postProcess_para
+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
 from utilsK.masterUtils import get_needed_objectsIndex
-def river_demo():
+def river_demo_v2():
     ##预先设置的参数
     device_='0'  ##选定模型，可选 cpu,'0','1'
+    ###注意TRT模型生成时，就需要对应cuda device，下面的trt文件是cuda:0生成的，device只能是0  
+    ##以下参数目前不可改
+    
+    Detweights = "../AIlib/weights/yolov5/class8/bestcao_fp16.engine" ##升级后的检测模型
+    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.engine' ##升级的分割模型
+    
+    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
+    
+    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/'
+    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" 
@@ -54,6 +113,7 @@ def river_demo():
     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)
     
     
@@ -74,7 +134,6 @@ def river_demo():
         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']
@@ -187,7 +246,7 @@ def jkm_demo():
 
         
 if __name__=="__main__":
-    #river_demo()
+    river_demo_v2()
     #road_forest_demo('forest'  )
-    jkm_demo()
+    #jkm_demo()
     

readme.md

@@ -11,4 +11,10 @@
 增加道路模型，测试图像在test/images/road 
 
 2022.11.21
-增加健康码（健康吗、行程卡）检测、车牌检测，测试图像在test/images/jkm、test/images/plate  
+增加健康码（健康吗、行程卡）检测、车牌检测，测试图像在test/images/jkm、test/images/plate 
+
+2022.12.06
+将河道测试的代码换成trt的。
+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