diff --git a/demo.py b/demo.py index 57426e6..b503536 100644 --- a/demo.py +++ b/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 main(): +def river_demo(): ##预先设置的参数 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() + diff --git a/images/forest/DJI_20220531130118_0298_2.JPG b/images/forest/DJI_20220531130118_0298_2.JPG new file mode 100755 index 0000000..a97a93d Binary files /dev/null and b/images/forest/DJI_20220531130118_0298_2.JPG differ