gitadmin
/
tuoheng_alg_move_20240715
ウォッチ 1
スター 0
フォーク 0
コード 課題 0 プルリクエスト 0 リリース 11 Wiki アクティビティ
選択できるのは25トピックまでです。 トピックは、先頭が英数字で、英数字とダッシュ('-')を使用した35文字以内のものにしてください。
137 コミット
8 ブランチ
ブランチ: master
ブランチ タグ
${ item.name }
ブランチ ${ searchTerm } を作成
'master' から
${ noResults }
tuoheng_alg_move_20240715/util/QueRiver.py

347 行
15KB
Raw パーマリンク Blame 履歴 

  1. from kafka import KafkaProducer, KafkaConsumer

  2. from kafka.errors import kafka_errors

  3. import traceback

  4. import json, base64,os

  5. import numpy as np

  6. from multiprocessing import Process,Queue

  7. import time,cv2,string,random

  8. import subprocess as sp

  9. 

  10. import matplotlib.pyplot as plt

  11. from utils.datasets import LoadStreams, LoadImages

  12. from models.experimental import attempt_load

  13. 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

  14. import torch,sys

  15. #from segutils.segmodel import SegModel,get_largest_contours

  16. #sys.path.extend(['../yolov5/segutils'])

  17. 

  18. from segutils.segWaterBuilding import  SegModel,get_largest_contours,illBuildings

  19. 

  20. #from segutils.core.models.bisenet import BiSeNet

  21. from segutils.core.models.bisenet import BiSeNet_MultiOutput

  22. 

  23. from utils.plots import plot_one_box,plot_one_box_PIL,draw_painting_joint,get_label_arrays,get_websource

  24. from collections import  Counter

  25. #import matplotlib

  26. 

  27. import matplotlib.pyplot as plt

  28. # get_labelnames,get_label_arrays,post_process_,save_problem_images,time_str

  29. #FP_DEBUG=open('debut.txt','w')

  30. def bsJpgCode(image_ori):            

  31.     jpgCode = cv2.imencode('.jpg',image_ori)[-1]###np.array,(4502009,1)

  32.     bsCode = str(base64.b64encode(jpgCode))[2:-1] ###str,长6002680

  33.     return bsCode

  34. def bsJpgDecode(bsCode):

  35.     bsDecode = base64.b64decode(bsCode)###types,长4502009

  36.     npString = np.frombuffer(bsDecode,np.uint8)###np.array,(长4502009,)

  37.     jpgDecode = cv2.imdecode(npString,cv2.IMREAD_COLOR)###np.array,(3000,4000,3)

  38.     return jpgDecode

  39. def get_ms(time0,time1):

  40.     str_time ='%.2f ms'%((time1-time0)*1000) 

  41.     return str_time  

  42. rainbows=[

  43.            (0,0,255),(0,255,0),(255,0,0),(255,0,255),(255,255,0),(255,127,0),(255,0,127),

  44.            (127,255,0),(0,255,127),(0,127,255),(127,0,255),(255,127,255),(255,255,127),

  45.            (127,255,255),(0,255,255),(255,127,255),(127,255,255),

  46.            (0,127,0),(0,0,127),(0,255,255)

  47.          ]

  48. 

  49. 

  50. def get_labelnames(labelnames):   

  51.     with open(labelnames,'r') as fp:

  52.         namesjson=json.load(fp)             

  53.         names_fromfile=namesjson['labelnames']                                    

  54.         names = names_fromfile

  55.     return names 

  56. 

  57. def check_stream(stream):

  58.     cap = cv2.VideoCapture(stream)        

  59.     if cap.isOpened():

  60.         return True

  61.     else:

  62.        return False 

  63. #####  

  64. def drawWater(pred,image_array0,river={'color':(0,255,255),'line_width':3,'segRegionCnt':2}):####pred是模型的输出，只有水分割的任务

  65.     ##画出水体区域

  66.     contours, hierarchy = cv2.findContours(pred,cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE)

  67.     water = pred.copy(); water[:,:] = 0 

  68.     

  69.     if len(contours)==0:        

  70.         return image_array0,water

  71.     max_ids = get_largest_contours(contours,river['segRegionCnt']);     

  72.     for max_id in max_ids:    

  73.         cv2.fillPoly(water, [contours[max_id][:,0,:]], 1)     

  74.         cv2.drawContours(image_array0,contours,max_id,river['color'],river['line_width'] )

  75.     return     image_array0,water

  76. 

  77. 

  78. def scale_back(boxes,padInfos):

  79.     top, left,r = padInfos[0:3]

  80.   

  81.     boxes[:,0] = (boxes[:,0] - left) * r

  82.     

  83.     boxes[:,2] = (boxes[:,2] - left) * r

  84.     boxes[:,1] = (boxes[:,1] - top) * r

  85.     boxes[:,3] = (boxes[:,3] - top) * r

  86.     return boxes

  87. 

  88. def img_pad(img, size, pad_value=[114,114,114]):

  89.     ###填充成固定尺寸

  90.     H,W,_ = img.shape

  91.     r = max(H/size[0], W/size[1])

  92.     img_r = cv2.resize(img, (int(W/r), int(H/r)))

  93.     tb = size[0] - img_r.shape[0]

  94.     lr = size[1] - img_r.shape[1]

  95.     top = int(tb/2)

  96.     bottom = tb - top

  97.     left = int(lr/2)

  98.     right = lr - left

  99.     pad_image = cv2.copyMakeBorder(img_r, top, bottom, left, right, cv2.BORDER_CONSTANT,value=pad_value)

  100.     return pad_image,(top, left,r)      

  101. 

  102. 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 ):

  103.     object_config,slopeIndex,segmodel,segRegionCnt=ObjectPar['object_config'],ObjectPar['slopeIndex'],ObjectPar['segmodel'],ObjectPar['segRegionCnt']

  104.     ##输入dataset genereate 生成的数据，model预测的结果pred，nms参数

  105.     ##主要操作NMS ---> 坐标转换 ---> 画图

  106.     ##输出原图、AI处理后的图、检测结果

  107.     time0=time.time()

  108.     path, img, im0s, vid_cap ,pred,seg_pred= datas[0:6];

  109.     #segmodel=True

  110.     pred = non_max_suppression(pred, conf_thres, iou_thres, classes=None, agnostic=False)   

  111.     time1=time.time()

  112.     i=0;det=pred[0]###一次检测一张图片

  113.     time1_1 = time.time()

  114.     #p, s, im0 = path[i], '%g: ' % i, im0s[i].copy()

  115.     p, s, im0 = path[i], '%g: ' % i, im0s[i]

  116.     time1_2 = time.time()

  117.     gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]  # normalization gain whwh

  118.     time1_3 = time.time()

  119.     det_xywh=[];

  120.     #im0_brg=cv2.cvtColor(im0,cv2.COLOR_RGB2BGR);

  121.     if segmodel:

  122.         if len(seg_pred)==2:

  123.             im0,water = illBuildings(seg_pred,im0) 

  124.         else:

  125.             river={ 'color':font['waterLineColor'],'line_width':font['waterLineWidth'],'segRegionCnt':segRegionCnt }

  126.             im0,water = drawWater(seg_pred,im0,river)      

  127.     time2=time.time()

  128.     #plt.imshow(im0);plt.show()

  129.     if len(det)>0:

  130.         # Rescale boxes from img_size to im0 size

  131.         if not padInfos:

  132.             det[:, :4] = scale_coords(img.shape[2:], det[:, :4],im0.shape).round()     

  133.         else:

  134.             #print('####line131:',det[:, :])

  135.             det[:, :4] = scale_back( det[:, :4],padInfos).round()

  136.             #print('####line133:',det[:, :])

  137.         #用seg模型，确定有效检测匡及河道轮廓线

  138.         if segmodel:               

  139.             cls_indexs =    det[:, 5].clone().cpu().numpy().astype(np.int32)

  140.             ##判断哪些目标属于岸坡的

  141.             slope_flag = np.array([x in slopeIndex for x in cls_indexs ]  )

  142.             

  143.             det_c = det.clone(); det_c=det_c.cpu().numpy()  

  144.             try:

  145.                 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]  )  

  146.             except:

  147.                 print('*****************************line143: error:',det_c)

  148.             water_flag = np.array(area_factors>0.1)            

  149.             det = det[water_flag|slope_flag]##如果是水上目标，则需要与水的iou超过0.1；如果是岸坡目标，则直接保留。

  150.         #对检测匡绘图

  151.         for *xyxy, conf, cls in reversed(det):

  152.             xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # normalized xywh

  153.             cls_c = cls.cpu().numpy()

  154.             if int(cls_c) not in object_config: ###如果不是所需要的目标，则不显示

  155.                 continue

  156.             conf_c = conf.cpu().numpy()

  157.             line = [float(cls_c), *xywh, float(conf_c)] # label format      

  158.             det_xywh.append(line)

  159.             label = f'{names[int(cls)]} {conf:.2f}'			    

  160.            

  161.             im0 = draw_painting_joint(xyxy,im0,label_arraylist[int(cls)],score=conf,color=rainbows[int(cls)%20],font=font)

  162.     time3=time.time()

  163.     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)   ) 

  164.     return [im0s[0],im0,det_xywh,iframe],strout

  165.    

  166. 

  167.           

  168. def preprocess(par):

  169.     print('#####process:',par['name'])

  170.     ##负责读取视频，生成原图及供检测的使用图，numpy格式   

  171.     #source='rtmp://liveplay.yunhengzhizao.cn/live/demo_HD5M'

  172.     #img_size=640;    stride=32

  173.     while True:

  174.         cap = cv2.VideoCapture(par['source'])

  175.         iframe = 0

  176.         if cap.isOpened():

  177.             print( '#### read %s  success!'%(par['source']))

  178.             try:

  179.                 dataset = LoadStreams(par['source'], img_size=640, stride=32)           

  180.                 for path, img, im0s, vid_cap in dataset:

  181.                     datas=[path, img, im0s, vid_cap,iframe]

  182.                     par['queOut'].put(datas)

  183.                     iframe +=1

  184.             except Exception as e:

  185.                 print('###read error:%s '%(par['source']))

  186.                 time.sleep(10)

  187.                 iframe = 0

  188.                 

  189.         else:

  190.             print('###read error:%s  '%(par['source'] ))

  191.             time.sleep(10)

  192.             iframe = 0

  193. 

  194. def gpu_process(par):

  195.     print('#####process:',par['name'])

  196.     half=True

  197.     ##gpu运算,检测模型    

  198.     weights = par['weights']

  199.     device = par['device'] 

  200.     print('###line127:',par['device'])

  201.     model = attempt_load(par['weights'], map_location=par['device'])  # load FP32 model

  202.     if half:

  203.         model.half()

  204. 

  205.     ##gpu运算，分割模型  

  206.     seg_nclass = par['seg_nclass']

  207.     seg_weights = par['seg_weights']

  208. 

  209.     #segmodel = SegModel(nclass=seg_nclass,weights=seg_weights,device=device)

  210.     

  211. 

  212.     nclass = [2,2]

  213.     Segmodel =  BiSeNet_MultiOutput(nclass)

  214.     weights='weights/segmentation/WaterBuilding.pth'

  215.     segmodel = SegModel(model=Segmodel,nclass=nclass,weights=weights,device='cuda:0',multiOutput=True)

  216.     while True:

  217.         if not par['queIn'].empty():

  218.             time0=time.time()

  219.             datas = par['queIn'].get() 

  220.             path, img, im0s, vid_cap,iframe = datas[0:5]

  221.             time1=time.time()

  222.             img = torch.from_numpy(img).to(device)

  223.             img = img.half() if half else img.float()  # uint8 to fp16/32

  224.             img /= 255.0  # 0 - 255 to 0.0 - 1.0

  225.             time2 = time.time()

  226.             pred = model(img,augment=False)[0]

  227.             time3 = time.time()

  228.             seg_pred = segmodel.eval(im0s[0],outsize=None,smooth_kernel=20)

  229.             time4 = time.time()

  230.             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()   )  

  231.             #FP_DEBUG.write( fpStr+'\n' )

  232.             datasOut = [path, img, im0s, vid_cap,pred,seg_pred,iframe]

  233.             par['queOut'].put(datasOut)

  234.             if par['debug']:

  235.                 print('#####process:',par['name'],'  line107')

  236.         else:

  237.             time.sleep(1/300)  

  238. def get_cls(array):

  239.     dcs = Counter(array)

  240.     keys = list(dcs.keys())

  241.     values = list(dcs.values())

  242.     max_index = values.index(max(values))

  243.     cls = int(keys[max_index])

  244.     return cls

  245. def save_problem_images(post_results,iimage_cnt,names,streamName='live-THSAHD5M',outImaDir='problems/images_tmp',imageTxtFile=False):

  246.     ## [cls, x,y,w,h, conf]

  247.     problem_image=[[] for i in range(6)]

  248. 

  249.        

  250.     dets_list = [x[2] for x in post_results]

  251. 

  252.     mean_scores=[  np.array(x)[:,5].mean()  for x in dets_list  ] ###mean conf

  253. 

  254.     best_index = mean_scores.index(max(mean_scores)) ##获取该批图片里，问题图片的index 

  255.     best_frame = post_results[ best_index][3] ##获取绝对帧号 

  256.     img_send = post_results[best_index][1]##AI处理后的图

  257.     img_bak = post_results[best_index][0]##原图    

  258.     cls_max = get_cls( x[0] for x in    dets_list[best_index] )

  259. 

  260. 

  261.     time_str = time.strftime("%Y-%m-%d-%H-%M-%S", time.localtime()) 

  262.     uid=''.join(random.sample(string.ascii_letters + string.digits, 16))

  263.     #ori_name = '2022-01-20-15-57-36_frame-368-720_type-漂浮物_qVh4zI08ZlwJN9on_s-live-THSAHD5M_OR.jpg'

  264.     #2022-01-13-15-07-57_frame-9999-9999_type-结束_9999999999999999_s-off-XJRW20220110115904_AI.jpg

  265.     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)

  266.     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)

  267.     

  268.     cv2.imwrite(outnameOR,img_send)

  269.     cv2.imwrite(outnameAR,img_bak)

  270.     if imageTxtFile:

  271.         outnameOR_txt = outnameOR.replace('.jpg','.txt')

  272.         fp=open(outnameOR_txt,'w');fp.write(outnameOR+'\n');fp.close()

  273.         outnameAI_txt = outnameAR.replace('.jpg','.txt')

  274.         fp=open(outnameAI_txt,'w');fp.write(outnameAR+'\n');fp.close()

  275.        

  276.     parOut = {}; parOut['imgOR'] =  img_send; parOut['imgAR'] =  img_send; parOut['uid']=uid

  277.     parOut['imgORname']=os.path.basename(outnameOR);parOut['imgARname']=os.path.basename(outnameAR);

  278.     parOut['time_str'] = time_str;parOut['type'] = names[cls_max]

  279.     return parOut

  280.        

  281. 

  282.      

  283.     

  284. def post_process(par):

  285.     

  286.     print('#####process:',par['name'])

  287.     ###post-process参数

  288.     conf_thres,iou_thres,classes=par['conf_thres'],par['iou_thres'],par['classes']

  289.     labelnames=par['labelnames']    

  290.     rainbows=par['rainbows']

  291.     fpsample = par['fpsample']

  292.     names=get_labelnames(labelnames)   

  293.     label_arraylist = get_label_arrays(names,rainbows,outfontsize=40)

  294.     iimage_cnt = 0

  295.     post_results=[]

  296.     while True:

  297.         if not par['queIn'].empty():

  298.             time0=time.time()

  299.             datas = par['queIn'].get() 

  300.             iframe = datas[6]

  301.             if par['debug']:

  302.                 print('#####process:',par['name'],'  line129')

  303.             p_result,timeOut = post_process_(datas,conf_thres, iou_thres,names,label_arraylist,rainbows,iframe)

  304.             par['queOut'].put(p_result)   

  305.             ##输出结果

  306.             

  307.     

  308. 

  309.             ##每隔 fpsample帧处理一次，如果有问题就保存图片       

  310.             if (iframe % fpsample == 0) and (len(post_results)>0) :

  311.                #print('####line204:',iframe,post_results) 

  312.                save_problem_images(post_results,iframe,names)

  313.                post_results=[]   

  314. 

  315.             if len(p_result[2] )>0:  ##                

  316.                 #post_list = p_result.append(iframe)

  317.                 post_results.append(p_result)

  318.                 #print('####line201:',type(p_result)) 

  319.                 

  320.             time1=time.time()

  321.             outstr='process:%s ,iframe:%d,%s , time:%s, queLen:%d  '%( par['name'],iframe,timeOut,get_ms(time0,time1) ,par['queIn'].qsize()   ) 

  322.             #FP_DEBUG.write(outstr +'\n')

  323.             #print( 'process:%s ,iframe:%d,%s , time:%s, queLen:%d  '%( par['name'],iframe,timeOut,get_ms(time0,time1) ,par['queIn'].qsize()   )  )        

  324.         else:            

  325.             time.sleep(1/300)

  326. 

  327.      

  328. def save_logfile(name,txt):

  329.     if os.path.exists(name):

  330.         fp=open(name,'r+')

  331.     else:

  332.         fp=open(name,'w')

  333.  

  334.     fp.write('%s %s \n'%(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()),txt))

  335.     fp.close()

  336. def time_str():     

  337.     return time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())

  338.     

  339.   

  340. 

  341. if __name__=='__main__':

  342.     jsonfile='config/queRiver.json'

  343.     #image_encode_decode()

  344.     work_stream(jsonfile)

  345.     #par={'name':'preprocess'}

  346.     #preprocess(par)