add forest model

demo.py

@@ -1,7 +1,7 @@
 import sys
 sys.path.extend(['..','../AIlib' ])
 
-from AI import AI_process
+from AI import AI_process,get_postProcess_para
 import cv2,os,time
 from segutils.segmodel import SegModel
 from models.experimental import attempt_load
@@ -10,8 +10,7 @@ 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 main():
     ##预先设置的参数
     device_='1'  ##选定模型，可选 cpu,'0','1'
     
@@ -19,9 +18,11 @@ def main():
     Detweights = "../AIlib/weights/yolov5/class5/best_5classes.pt"
     seg_nclass = 2 
     Segweights = "../AIlib/weights/BiSeNet/checkpoint.pth"
-    conf_thres,iou_thres,classes= 0.25,0.45,5
     labelnames = "../AIlib/weights/yolov5/class5/labelnames.json"
-    rainbows = [ [0,0,255],[0,255,0],[255,0,0],[255,0,255],[255,255,0],[255,129,0],[255,0,127],[127,255,0],[0,255,127],[0,127,255],[127,0,255],[255,127,255],[255,255,127],[127,255,255],[0,255,255],[255,127,255],[127,255,255],  [0,127,0],[0,0,127],[0,255,255]]
+    postFile= '../AIlib/conf/para.json'
+    
+    conf_thres,iou_thres,classes,rainbows=get_postProcess_para(postFile)
+   
     ####模型选择参数用如下：
     mode_paras=[
         {
@@ -48,6 +49,7 @@ def main():
     ##加载模型，准备好显示字符
     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
@@ -69,11 +71,77 @@ def main():
         time11 = time.time()
         image_array = p_result[1]
         cv2.imwrite(  os.path.join( outpth,folders[i] ) ,image_array )
-        print('----process:%s'%(folders[i]), (time.time() - time11) * 1000)
+        print('----process:%s'%(folders[i]), (time.time() - time11) * 1000,timeOut)
+        
+def forest_demo():
+    ##预先设置的参数
+    device_='1'  ##选定模型，可选 cpu,'0','1'
+    
+    ##以下参数目前不可改
+    Detweights = "../AIlib/weights/forest/best.pt"
+    labelnames = "../AIlib/weights/forest/labelnames.json"
+    postFile= '../AIlib/conf/para.json'
+    
+    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/forest/'
+    outpth = 'images/results/'
+    folders = os.listdir(impth)
+    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(im0s,model,segmodel,names,label_arraylist,rainbows,half,device,conf_thres, iou_thres,allowedList)
+        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) )
 
     
     
     
         
 if __name__=="__main__":
-    main()
+    #river_demo()
+    forest_demo()
+