347 lines
15 KiB
Python
347 lines
15 KiB
Python
from kafka import KafkaProducer, KafkaConsumer
|
||
from kafka.errors import kafka_errors
|
||
import traceback
|
||
import json, base64,os
|
||
import numpy as np
|
||
from multiprocessing import Process,Queue
|
||
import time,cv2,string,random
|
||
import subprocess as sp
|
||
|
||
import matplotlib.pyplot as plt
|
||
from utils.datasets import LoadStreams, LoadImages
|
||
from models.experimental import attempt_load
|
||
from utils.general import check_img_size, check_requirements, check_imshow, non_max_suppression, apply_classifier, scale_coords, xyxy2xywh, strip_optimizer, set_logging, increment_path
|
||
import torch,sys
|
||
#from segutils.segmodel import SegModel,get_largest_contours
|
||
#sys.path.extend(['../yolov5/segutils'])
|
||
|
||
from segutils.segWaterBuilding import SegModel,get_largest_contours,illBuildings
|
||
|
||
#from segutils.core.models.bisenet import BiSeNet
|
||
from segutils.core.models.bisenet import BiSeNet_MultiOutput
|
||
|
||
from utils.plots import plot_one_box,plot_one_box_PIL,draw_painting_joint,get_label_arrays,get_websource
|
||
from collections import Counter
|
||
#import matplotlib
|
||
|
||
import matplotlib.pyplot as plt
|
||
# get_labelnames,get_label_arrays,post_process_,save_problem_images,time_str
|
||
#FP_DEBUG=open('debut.txt','w')
|
||
def bsJpgCode(image_ori):
|
||
jpgCode = cv2.imencode('.jpg',image_ori)[-1]###np.array,(4502009,1)
|
||
bsCode = str(base64.b64encode(jpgCode))[2:-1] ###str,长6002680
|
||
return bsCode
|
||
def bsJpgDecode(bsCode):
|
||
bsDecode = base64.b64decode(bsCode)###types,长4502009
|
||
npString = np.frombuffer(bsDecode,np.uint8)###np.array,(长4502009,)
|
||
jpgDecode = cv2.imdecode(npString,cv2.IMREAD_COLOR)###np.array,(3000,4000,3)
|
||
return jpgDecode
|
||
def get_ms(time0,time1):
|
||
str_time ='%.2f ms'%((time1-time0)*1000)
|
||
return str_time
|
||
rainbows=[
|
||
(0,0,255),(0,255,0),(255,0,0),(255,0,255),(255,255,0),(255,127,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)
|
||
]
|
||
|
||
|
||
def get_labelnames(labelnames):
|
||
with open(labelnames,'r') as fp:
|
||
namesjson=json.load(fp)
|
||
names_fromfile=namesjson['labelnames']
|
||
names = names_fromfile
|
||
return names
|
||
|
||
def check_stream(stream):
|
||
cap = cv2.VideoCapture(stream)
|
||
if cap.isOpened():
|
||
return True
|
||
else:
|
||
return False
|
||
#####
|
||
def drawWater(pred,image_array0,river={'color':(0,255,255),'line_width':3,'segRegionCnt':2}):####pred是模型的输出，只有水分割的任务
|
||
##画出水体区域
|
||
contours, hierarchy = cv2.findContours(pred,cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE)
|
||
water = pred.copy(); water[:,:] = 0
|
||
|
||
if len(contours)==0:
|
||
return image_array0,water
|
||
max_ids = get_largest_contours(contours,river['segRegionCnt']);
|
||
for max_id in max_ids:
|
||
cv2.fillPoly(water, [contours[max_id][:,0,:]], 1)
|
||
cv2.drawContours(image_array0,contours,max_id,river['color'],river['line_width'] )
|
||
return image_array0,water
|
||
|
||
|
||
def scale_back(boxes,padInfos):
|
||
top, left,r = padInfos[0:3]
|
||
|
||
boxes[:,0] = (boxes[:,0] - left) * r
|
||
|
||
boxes[:,2] = (boxes[:,2] - left) * r
|
||
boxes[:,1] = (boxes[:,1] - top) * r
|
||
boxes[:,3] = (boxes[:,3] - top) * r
|
||
return boxes
|
||
|
||
def img_pad(img, size, pad_value=[114,114,114]):
|
||
###填充成固定尺寸
|
||
H,W,_ = img.shape
|
||
r = max(H/size[0], W/size[1])
|
||
img_r = cv2.resize(img, (int(W/r), int(H/r)))
|
||
tb = size[0] - img_r.shape[0]
|
||
lr = size[1] - img_r.shape[1]
|
||
top = int(tb/2)
|
||
bottom = tb - top
|
||
left = int(lr/2)
|
||
right = lr - left
|
||
pad_image = cv2.copyMakeBorder(img_r, top, bottom, left, right, cv2.BORDER_CONSTANT,value=pad_value)
|
||
return pad_image,(top, left,r)
|
||
|
||
def post_process_(datas,conf_thres, iou_thres,names,label_arraylist,rainbows,iframe,ObjectPar={ 'object_config':[0,1,2,3,4], 'slopeIndex':[5,6,7] ,'segmodel':True,'segRegionCnt':1 },font={ 'line_thickness':None, 'fontSize':None,'boxLine_thickness':None,'waterLineColor':(0,255,255),'waterLineWidth':3},padInfos=None ):
|
||
object_config,slopeIndex,segmodel,segRegionCnt=ObjectPar['object_config'],ObjectPar['slopeIndex'],ObjectPar['segmodel'],ObjectPar['segRegionCnt']
|
||
##输入dataset genereate 生成的数据，model预测的结果pred，nms参数
|
||
##主要操作NMS ---> 坐标转换 ---> 画图
|
||
##输出原图、AI处理后的图、检测结果
|
||
time0=time.time()
|
||
path, img, im0s, vid_cap ,pred,seg_pred= datas[0:6];
|
||
#segmodel=True
|
||
pred = non_max_suppression(pred, conf_thres, iou_thres, classes=None, agnostic=False)
|
||
time1=time.time()
|
||
i=0;det=pred[0]###一次检测一张图片
|
||
time1_1 = time.time()
|
||
#p, s, im0 = path[i], '%g: ' % i, im0s[i].copy()
|
||
p, s, im0 = path[i], '%g: ' % i, im0s[i]
|
||
time1_2 = time.time()
|
||
gn = torch.tensor(im0.shape)[[1, 0, 1, 0]] # normalization gain whwh
|
||
time1_3 = time.time()
|
||
det_xywh=[];
|
||
#im0_brg=cv2.cvtColor(im0,cv2.COLOR_RGB2BGR);
|
||
if segmodel:
|
||
if len(seg_pred)==2:
|
||
im0,water = illBuildings(seg_pred,im0)
|
||
else:
|
||
river={ 'color':font['waterLineColor'],'line_width':font['waterLineWidth'],'segRegionCnt':segRegionCnt }
|
||
im0,water = drawWater(seg_pred,im0,river)
|
||
time2=time.time()
|
||
#plt.imshow(im0);plt.show()
|
||
if len(det)>0:
|
||
# Rescale boxes from img_size to im0 size
|
||
if not padInfos:
|
||
det[:, :4] = scale_coords(img.shape[2:], det[:, :4],im0.shape).round()
|
||
else:
|
||
#print('####line131:',det[:, :])
|
||
det[:, :4] = scale_back( det[:, :4],padInfos).round()
|
||
#print('####line133:',det[:, :])
|
||
#用seg模型，确定有效检测匡及河道轮廓线
|
||
if segmodel:
|
||
cls_indexs = det[:, 5].clone().cpu().numpy().astype(np.int32)
|
||
##判断哪些目标属于岸坡的
|
||
slope_flag = np.array([x in slopeIndex for x in cls_indexs ] )
|
||
|
||
det_c = det.clone(); det_c=det_c.cpu().numpy()
|
||
try:
|
||
area_factors = np.array([np.sum(water[int(x[1]):int(x[3]), int(x[0]):int(x[2])] )*1.0/(1.0*(x[2]-x[0])*(x[3]-x[1])+0.00001) for x in det_c] )
|
||
except:
|
||
print('*****************************line143: error:',det_c)
|
||
water_flag = np.array(area_factors>0.1)
|
||
det = det[water_flag|slope_flag]##如果是水上目标，则需要与水的iou超过0.1；如果是岸坡目标，则直接保留。
|
||
#对检测匡绘图
|
||
for *xyxy, conf, cls in reversed(det):
|
||
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh
|
||
cls_c = cls.cpu().numpy()
|
||
if int(cls_c) not in object_config: ###如果不是所需要的目标，则不显示
|
||
continue
|
||
conf_c = conf.cpu().numpy()
|
||
line = [float(cls_c), *xywh, float(conf_c)] # label format
|
||
det_xywh.append(line)
|
||
label = f'{names[int(cls)]} {conf:.2f}'
|
||
|
||
im0 = draw_painting_joint(xyxy,im0,label_arraylist[int(cls)],score=conf,color=rainbows[int(cls)%20],font=font)
|
||
time3=time.time()
|
||
strout='nms:%s drawWater:%s,copy:%s,toTensor:%s,detDraw:%s '%(get_ms(time0,time1),get_ms(time1,time2),get_ms(time1_1,time1_2),get_ms(time1_2,time1_3), get_ms(time2,time3) )
|
||
return [im0s[0],im0,det_xywh,iframe],strout
|
||
|
||
|
||
|
||
def preprocess(par):
|
||
print('#####process:',par['name'])
|
||
##负责读取视频，生成原图及供检测的使用图，numpy格式
|
||
#source='rtmp://liveplay.yunhengzhizao.cn/live/demo_HD5M'
|
||
#img_size=640; stride=32
|
||
while True:
|
||
cap = cv2.VideoCapture(par['source'])
|
||
iframe = 0
|
||
if cap.isOpened():
|
||
print( '#### read %s success!'%(par['source']))
|
||
try:
|
||
dataset = LoadStreams(par['source'], img_size=640, stride=32)
|
||
for path, img, im0s, vid_cap in dataset:
|
||
datas=[path, img, im0s, vid_cap,iframe]
|
||
par['queOut'].put(datas)
|
||
iframe +=1
|
||
except Exception as e:
|
||
print('###read error:%s '%(par['source']))
|
||
time.sleep(10)
|
||
iframe = 0
|
||
|
||
else:
|
||
print('###read error:%s '%(par['source'] ))
|
||
time.sleep(10)
|
||
iframe = 0
|
||
|
||
def gpu_process(par):
|
||
print('#####process:',par['name'])
|
||
half=True
|
||
##gpu运算,检测模型
|
||
weights = par['weights']
|
||
device = par['device']
|
||
print('###line127:',par['device'])
|
||
model = attempt_load(par['weights'], map_location=par['device']) # load FP32 model
|
||
if half:
|
||
model.half()
|
||
|
||
##gpu运算，分割模型
|
||
seg_nclass = par['seg_nclass']
|
||
seg_weights = par['seg_weights']
|
||
|
||
#segmodel = SegModel(nclass=seg_nclass,weights=seg_weights,device=device)
|
||
|
||
|
||
nclass = [2,2]
|
||
Segmodel = BiSeNet_MultiOutput(nclass)
|
||
weights='weights/segmentation/WaterBuilding.pth'
|
||
segmodel = SegModel(model=Segmodel,nclass=nclass,weights=weights,device='cuda:0',multiOutput=True)
|
||
while True:
|
||
if not par['queIn'].empty():
|
||
time0=time.time()
|
||
datas = par['queIn'].get()
|
||
path, img, im0s, vid_cap,iframe = datas[0:5]
|
||
time1=time.time()
|
||
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
|
||
time2 = time.time()
|
||
pred = model(img,augment=False)[0]
|
||
time3 = time.time()
|
||
seg_pred = segmodel.eval(im0s[0],outsize=None,smooth_kernel=20)
|
||
time4 = time.time()
|
||
fpStr= 'process:%s ,iframe:%d,getdata:%s,copygpu:%s,dettime:%s,segtime:%s , time:%s, queLen:%d '%( par['name'],iframe,get_ms(time0,time1) ,get_ms(time1,time2) ,get_ms(time2,time3) ,get_ms(time3,time4),get_ms(time0,time4) ,par['queIn'].qsize() )
|
||
#FP_DEBUG.write( fpStr+'\n' )
|
||
datasOut = [path, img, im0s, vid_cap,pred,seg_pred,iframe]
|
||
par['queOut'].put(datasOut)
|
||
if par['debug']:
|
||
print('#####process:',par['name'],' line107')
|
||
else:
|
||
time.sleep(1/300)
|
||
def get_cls(array):
|
||
dcs = Counter(array)
|
||
keys = list(dcs.keys())
|
||
values = list(dcs.values())
|
||
max_index = values.index(max(values))
|
||
cls = int(keys[max_index])
|
||
return cls
|
||
def save_problem_images(post_results,iimage_cnt,names,streamName='live-THSAHD5M',outImaDir='problems/images_tmp',imageTxtFile=False):
|
||
## [cls, x,y,w,h, conf]
|
||
problem_image=[[] for i in range(6)]
|
||
|
||
|
||
dets_list = [x[2] for x in post_results]
|
||
|
||
mean_scores=[ np.array(x)[:,5].mean() for x in dets_list ] ###mean conf
|
||
|
||
best_index = mean_scores.index(max(mean_scores)) ##获取该批图片里，问题图片的index
|
||
best_frame = post_results[ best_index][3] ##获取绝对帧号
|
||
img_send = post_results[best_index][1]##AI处理后的图
|
||
img_bak = post_results[best_index][0]##原图
|
||
cls_max = get_cls( x[0] for x in dets_list[best_index] )
|
||
|
||
|
||
time_str = time.strftime("%Y-%m-%d-%H-%M-%S", time.localtime())
|
||
uid=''.join(random.sample(string.ascii_letters + string.digits, 16))
|
||
#ori_name = '2022-01-20-15-57-36_frame-368-720_type-漂浮物_qVh4zI08ZlwJN9on_s-live-THSAHD5M_OR.jpg'
|
||
#2022-01-13-15-07-57_frame-9999-9999_type-结束_9999999999999999_s-off-XJRW20220110115904_AI.jpg
|
||
outnameOR= '%s/%s_frame-%d-%d_type-%s_%s_s-%s_AI.jpg'%(outImaDir,time_str,best_frame,iimage_cnt,names[cls_max],uid,streamName)
|
||
outnameAR= '%s/%s_frame-%d-%d_type-%s_%s_s-%s_OR.jpg'%(outImaDir,time_str,best_frame,iimage_cnt,names[cls_max],uid,streamName)
|
||
|
||
cv2.imwrite(outnameOR,img_send)
|
||
cv2.imwrite(outnameAR,img_bak)
|
||
if imageTxtFile:
|
||
outnameOR_txt = outnameOR.replace('.jpg','.txt')
|
||
fp=open(outnameOR_txt,'w');fp.write(outnameOR+'\n');fp.close()
|
||
outnameAI_txt = outnameAR.replace('.jpg','.txt')
|
||
fp=open(outnameAI_txt,'w');fp.write(outnameAR+'\n');fp.close()
|
||
|
||
parOut = {}; parOut['imgOR'] = img_send; parOut['imgAR'] = img_send; parOut['uid']=uid
|
||
parOut['imgORname']=os.path.basename(outnameOR);parOut['imgARname']=os.path.basename(outnameAR);
|
||
parOut['time_str'] = time_str;parOut['type'] = names[cls_max]
|
||
return parOut
|
||
|
||
|
||
|
||
|
||
def post_process(par):
|
||
|
||
print('#####process:',par['name'])
|
||
###post-process参数
|
||
conf_thres,iou_thres,classes=par['conf_thres'],par['iou_thres'],par['classes']
|
||
labelnames=par['labelnames']
|
||
rainbows=par['rainbows']
|
||
fpsample = par['fpsample']
|
||
names=get_labelnames(labelnames)
|
||
label_arraylist = get_label_arrays(names,rainbows,outfontsize=40)
|
||
iimage_cnt = 0
|
||
post_results=[]
|
||
while True:
|
||
if not par['queIn'].empty():
|
||
time0=time.time()
|
||
datas = par['queIn'].get()
|
||
iframe = datas[6]
|
||
if par['debug']:
|
||
print('#####process:',par['name'],' line129')
|
||
p_result,timeOut = post_process_(datas,conf_thres, iou_thres,names,label_arraylist,rainbows,iframe)
|
||
par['queOut'].put(p_result)
|
||
##输出结果
|
||
|
||
|
||
|
||
##每隔 fpsample帧处理一次，如果有问题就保存图片
|
||
if (iframe % fpsample == 0) and (len(post_results)>0) :
|
||
#print('####line204:',iframe,post_results)
|
||
save_problem_images(post_results,iframe,names)
|
||
post_results=[]
|
||
|
||
if len(p_result[2] )>0: ##
|
||
#post_list = p_result.append(iframe)
|
||
post_results.append(p_result)
|
||
#print('####line201:',type(p_result))
|
||
|
||
time1=time.time()
|
||
outstr='process:%s ,iframe:%d,%s , time:%s, queLen:%d '%( par['name'],iframe,timeOut,get_ms(time0,time1) ,par['queIn'].qsize() )
|
||
#FP_DEBUG.write(outstr +'\n')
|
||
#print( 'process:%s ,iframe:%d,%s , time:%s, queLen:%d '%( par['name'],iframe,timeOut,get_ms(time0,time1) ,par['queIn'].qsize() ) )
|
||
else:
|
||
time.sleep(1/300)
|
||
|
||
|
||
def save_logfile(name,txt):
|
||
if os.path.exists(name):
|
||
fp=open(name,'r+')
|
||
else:
|
||
fp=open(name,'w')
|
||
|
||
fp.write('%s %s \n'%(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()),txt))
|
||
fp.close()
|
||
def time_str():
|
||
return time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
|
||
|
||
|
||
|
||
if __name__=='__main__':
|
||
jsonfile='config/queRiver.json'
|
||
#image_encode_decode()
|
||
work_stream(jsonfile)
|
||
#par={'name':'preprocess'}
|
||
#preprocess(par)
|