161 lines
5.9 KiB
Python
161 lines
5.9 KiB
Python
import sys
|
||
sys.path.extend(['..','../AIlib' ])
|
||
|
||
from AI import AI_process,AI_process_forest,get_postProcess_para
|
||
import cv2,os,time
|
||
from segutils.segmodel import SegModel
|
||
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():
|
||
##预先设置的参数
|
||
device_='0' ##选定模型，可选 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) )
|
||
|
||
|
||
|
||
|
||
|
||
if __name__=="__main__":
|
||
river_demo()
|
||
#road_forest_demo('forest' )
|
||
|