59 lines
2.2 KiB
Python
59 lines
2.2 KiB
Python
|
import sys
|
||
|
|
sys.path.extend(['..','../AIlib' ])
|
||
|
|
|
||
|
|
from AI import AI_process
|
||
|
|
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
|
||
|
|
|
||
|
|
def main():
|
||
|
|
##预先设置的参数
|
||
|
|
device_='1' ##选定模型，可选 cpu,'0','1'
|
||
|
|
|
||
|
|
##以下参数目前不可改
|
||
|
|
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]]
|
||
|
|
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/examples/'
|
||
|
|
outpth = 'images/results/'
|
||
|
|
folders = os.listdir(impth)
|
||
|
|
for i in range(len(folders)):
|
||
|
|
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('----process:%s'%(folders[i]), (time.time() - time11) * 1000)
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
if __name__=="__main__":
|
||
|
|
main()
|