91 lines
3.7 KiB
Python
91 lines
3.7 KiB
Python
|
||
import cv2,os,time
|
||
from models.experimental import attempt_load
|
||
from segutils.segmodel import SegModel,get_largest_contours
|
||
from utils.torch_utils import select_device
|
||
from utilsK.queRiver import get_labelnames,get_label_arrays,post_process_
|
||
from utils.datasets import letterbox
|
||
import numpy as np
|
||
import torch
|
||
|
||
def AI_process(im0s,model,segmodel,names,label_arraylist,rainbows,half=True,device=' cuda:0',conf_thres=0.25, iou_thres=0.45,allowedList=[0,1,2,3]):
|
||
#输入参数
|
||
# im0s---原始图像列表
|
||
# model---检测模型，segmodel---分割模型
|
||
#输出：两个元素（列表，字符）构成的元组，[im0s[0],im0,det_xywh,iframe],strout
|
||
# [im0s[0],im0,det_xywh,iframe]中，
|
||
# im0s[0]--原始图像，im0--AI处理后的图像,iframe--帧号/暂时不需用到。
|
||
# det_xywh--检测结果，是一个列表。
|
||
# 其中每一个元素表示一个目标构成如：[float(cls_c), xc,yc,w,h, float(conf_c)]
|
||
# #cls_c--类别，如0,1,2,3; xc,yc,w,h--中心点坐标及宽；conf_c--得分, 取值范围在0-1之间
|
||
# #strout---统计AI处理个环节的时间
|
||
|
||
# Letterbox
|
||
img = [letterbox(x, 640, auto=True, stride=32)[0] for x in im0s]
|
||
# Stack
|
||
img = np.stack(img, 0)
|
||
# Convert
|
||
img = img[:, :, :, ::-1].transpose(0, 3, 1, 2) # BGR to RGB, to bsx3x416x416
|
||
img = np.ascontiguousarray(img)
|
||
|
||
|
||
img = torch.from_numpy(img).to(device)
|
||
img = img.half() if half else img.float() # uint8 to fp16/32
|
||
|
||
img /= 255.0 # 0 - 255 to 0.0 - 1.0
|
||
|
||
seg_pred,segstr = segmodel.eval(im0s[0] )
|
||
pred = model(img,augment=False)[0]
|
||
datas = [[''], img, im0s, None,pred,seg_pred,10]
|
||
|
||
p_result,timeOut = post_process_(datas,conf_thres, iou_thres,names,label_arraylist,rainbows,10,object_config=allowedList)
|
||
|
||
return p_result,timeOut
|
||
|
||
def main():
|
||
##预先设置的参数
|
||
device_='1' ##选定模型，可选 cpu,'0','1'
|
||
|
||
##以下参数目前不可改
|
||
Detweights = "weights/yolov5/class5/best_5classes.pt"
|
||
seg_nclass = 2
|
||
Segweights = "weights/BiSeNet/checkpoint.pth"
|
||
conf_thres,iou_thres,classes= 0.25,0.45,5
|
||
labelnames = "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="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()
|