import argparse import time from pathlib import Path import random,string import cv2 import torch import torch.backends.cudnn as cudnn import base64 import requests from models.experimental import attempt_load from utils.datasets import LoadStreams, LoadImages 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 from utils.plots import plot_one_box,plot_one_box_PIL,draw_painting_joint,get_label_arrays,get_websource,smooth_outline_auto from utils.get_offline_url import update_websource_offAndLive,platurlToJsonfile, get_websource_fromTxt from utils.torch_utils import select_device, load_classifier, time_synchronized import cv2 import queue import os,json,sys import numpy as np from threading import Thread import datetime,_thread import subprocess as sp import time from PIL import Image, ImageDraw, ImageFont from segutils.segmodel import SegModel,get_largest_contours sys.path.extend(['/home/thsw2/WJ/src/yolov5/segutils']) #from segutils.segWaterBuilding import SegModel,get_largest_contours,illBuildings from segutils.core.models.bisenet import BiSeNet_MultiOutput from collections import Counter import matplotlib.pyplot as plt platform_query_url='http://47.96.182.154:9051/api/suanfa/getPlatformInfo' offlineFile='mintors/offlines/doneCodes.txt' 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 ##9月3日之后，需要新增语义分割模型。 ##10月22日， source 改成文件输入 ##12月31日，支持从platform读取离线视频 # 使用线程锁，防止线程死锁 mutex = _thread.allocate_lock() # 存图片的队列 frame_queue = queue.Queue() # 推流的地址，前端通过这个地址拉流，主机的IP，2019是ffmpeg在nginx中设置的端口号 #camera_path='rtmp://58.200.131.2:1935/livetv/cctv1' camera_path='/data/WJ/data/THexit/vedio/XiYuDaiHe4.MP4' ###lables colors (BGR)##### 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 detect(save_img=False): rtmpUrl = "rtmp://127.0.0.1:1935/live/test" OutVideoW,OutVideoH,OutVideoFps=int(opt.OutVideoW),int(opt.OutVideoH),int(opt.OutVideoFps) command=['ffmpeg', '-y', #'-re',' ', '-f', 'rawvideo', '-vcodec','rawvideo', '-pix_fmt', 'bgr24', '-s', "{}x{}".format(OutVideoW,OutVideoH),# 图片分辨率 #'-vcodec','libx264', #'-b','2500k', '-r', str(OutVideoFps),# 视频帧率 '-i', '-', '-c:v', 'libx264', '-pix_fmt', 'yuv420p', #'-preset', 'ultrafast', '-f', 'flv', rtmpUrl] source, weights, view_img, save_txt, imgsz = opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size #save_img = not opt.nosave and not source.endswith('.txt') # save inference images save_img = not opt.nosave webcam = source.isnumeric() or source.endswith('.txt') or source.lower().startswith( ('rtsp://', 'rtmp://', 'http://', 'https://')) sourceTxt=source if source.endswith('.txt'): #source_list,port_list,streamName_list = get_websource(source) source_infos = get_websource_fromTxt(source) else: #source_list,port_list,streamName_list = [source],[1935],['demo'] source_infos = [{'url':source,'port':1935,'name':'demo' }] # Directories save_dir = Path(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok)) # increment run (save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True) # make dir # Initialize set_logging() device = select_device(opt.device) half = device.type != 'cpu' # half precision only supported on CUDA # Load model model = attempt_load(weights, map_location=device) # load FP32 model stride = int(model.stride.max()) # model stride imgsz = check_img_size(imgsz, s=stride) # check img_size #print('###'*20,imgsz,stride) ##加载分割模型### seg_nclass = 2 weights = 'weights/segmentation/BiSeNet/checkpoint.pth' segmodel = SegModel(nclass=seg_nclass,weights=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)''' if half: model.half() # to FP16 # Second-stage classifier classify = False if classify: modelc = load_classifier(name='resnet101', n=2) # initialize modelc.load_state_dict(torch.load('weights/resnet101.pt', map_location=device)['model']).to(device).eval() if webcam: # create file pointer fp_out=open('mintors/%s.txt'%(time.strftime("Start-%Y-%m-%d-%H-%M-%S", time.localtime())) ,'w') # Set Dataloader vid_path, vid_writer = None, None stream_id = 0 platurlToJsonfile(platform_query_url) while True: #for isource in range(len(source_list)): for isource in range(len(source_infos)): #source , port ,streamName = source_list[isource],port_list[isource],streamName_list[isource] source , port ,streamName = source_infos[isource]['url'],source_infos[isource]['port'],source_infos[isource]['name'] print('########## detect.py line129 souce informations:',isource, source , port ,streamName,webcam ) Push_Flag = False #for wang in ['test']: try: if webcam: #view_img = check_imshow() cudnn.benchmark = True # set True to speed up constant image size inference print('#########Using web cam#################') dataset = LoadStreams(source, img_size=imgsz, stride=stride) # Get names and colors,fp_log,fp_out都是日志文件 fp_log=open('mintors/detection/stream_%s_%d-%s.txt'%(streamName,stream_id,time.strftime("%Y-%m-%d-%H-%M-%S", time.localtime())) ,'w') fp_out.write(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())+ ' rtmp stream-%s-%d starts \n'%(streamName,stream_id) ) fp_out.flush() problem_image = [[],[],[],[],[]]; else: dataset = LoadImages(source, img_size=imgsz, stride=stride) #### iimage_cnt = 0 names = model.module.names if hasattr(model, 'module') else model.names EngLish_label = True if os.path.exists(opt.labelnames): with open(opt.labelnames,'r') as fp: namesjson=json.load(fp) names_fromfile=namesjson['labelnames'] if len(names_fromfile) == len(names): names = names_fromfile EngLish_label = False else: print('******Warning 文件:%s读取的类别数目与模型中的数目不一致，使用模型的类别********'%(opt.labelnames)) colors = rainbows label_arraylist = get_label_arrays(names,colors,outfontsize=40) # Run inference if device.type != 'cpu': model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters()))) # run once t00 = time.time(); if webcam: while True: if len(command) > 0: rtmpUrl = "rtmp://127.0.0.1:%s/live/test"%(port) command[-1] = rtmpUrl # 管道配置，其中用到管道 print(command) ppipe = sp.Popen(command, stdin=sp.PIPE) Push_Flag = True break time00=time.time() for path, img, im0s, vid_cap in dataset: t0= time_synchronized() 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 timeseg0 = time.time() if segmodel: if webcam: seg_pred,segstr = segmodel.eval(im0s[0] ) #seg_pred = segmodel.eval(im0s[0],outsize=None,smooth_kernel=20) else: seg_pred,segstr = segmodel.eval(im0s ) #seg_pred = segmodel.eval(im0s[0],outsize=None,smooth_kernel=20) timeseg1 = time.time() if img.ndimension() == 3: img = img.unsqueeze(0) # Inference t1 = time_synchronized() pred = model(img, augment=opt.augment)[0] #print('###','line197:',img.shape,opt.augment,opt.conf_thres, opt.iou_thres, opt.classes, opt.agnostic_nms) # Apply NMS pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, classes=opt.classes, agnostic=opt.agnostic_nms) t2 = time_synchronized() # Apply Classifier if classify: pred = apply_classifier(pred, modelc, img, im0s) # Process detections for i, det in enumerate(pred): # detections per image if webcam: # batch_size >= 1 p, s, im0, frame = path[i], '%g: ' % i, im0s[i].copy(), dataset.count im0_bak = im0.copy() else: p, s, im0, frame = path, '', im0s, getattr(dataset, 'frame', 0) iimage_cnt += 1 p = Path(p) # to Path save_path = str(save_dir / p.name) # img.jpg txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}') # img.txt s += '%gx%g ' % img.shape[2:] # print string gn = torch.tensor(im0.shape)[[1, 0, 1, 0]] # normalization gain whwh if segmodel: contours, hierarchy = cv2.findContours(seg_pred,cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE) water = seg_pred.copy(); if len(contours)>0: max_id = get_largest_contours(contours) water[:,:]=0 cv2.fillPoly(water, [contours[max_id][:,0,:]], 1) cv2.drawContours(im0,contours,max_id,(0,255,255),3) else: water[:,:] = 0 #im0,water = illBuildings(seg_pred,im0) if len(det): # Rescale boxes from img_size to im0 size det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round() #check weather the box is inside of water area if segmodel: det_c = det.clone(); det_c=det_c.cpu().numpy() #area_factors = np.array([np.sum(water[int(x[1]):int(x[3]), int(x[0]):int(x[2])] )/((x[2]-x[0])*(x[3]-x[1])) for x in det] ) area_factors = np.array([np.sum(water[int(x[1]):int(x[3]), int(x[0]):int(x[2])] )/((x[2]-x[0])*(x[3]-x[1])) for x in det_c] ) #det = det[area_factors>0.1] det = det[area_factors>0.03] ###联通要求，临时屏蔽掉水生植被 '''if len(det): clss = det[:,5] det = det[clss!=2]''' if len(det)>0: # Print results for c in det[:, -1].unique(): n = (det[:, -1] == c).sum() # detections per class s += f"{n} {names[int(c)]}{'s' * (n > 1)}, " # add to string # Write results for *xyxy, conf, cls in reversed(det): if save_txt: # Write to file xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh line = (cls, *xywh, conf) if opt.save_conf else (cls, *xywh) # label format with open(txt_path + '.txt', 'a') as f: f.write(('%g ' * len(line)).rstrip() % line + '\n') if save_img or view_img: # Add bbox to image label = f'{names[int(cls)]} {conf:.2f}' if EngLish_label: plot_one_box(xyxy, im0, label=label, color=colors[int(cls)%20], line_thickness=3) else: #im0=plot_one_box_PIL(xyxy, im0, label=label, color=colors[int(cls)%20], line_thickness=3) im0 = draw_painting_joint(xyxy,im0,label_arraylist[int(cls)],score=conf,color=rainbows[int(cls)%20],line_thickness=None) ###处理问题图片，每fpsample帧，上报一张图片。以平均得分最大的为最佳图片 if webcam: problem_image[0].append( det[:,4].mean()); problem_image[1].append(det); problem_image[2].append(im0);problem_image[3].append(iimage_cnt);problem_image[4].append(im0_bak) if webcam & (iimage_cnt % opt.fpsample == 0) & (dataset.mode != 'image') & (len(problem_image[0])>0): best_index = problem_image[0].index(max(problem_image[0])) best_frame = problem_image[3][ best_index] img_send = problem_image[2][ best_index] img_bak = problem_image[4][ best_index] dets = problem_image[1][best_index] cls_max = get_cls(dets[:,5].cpu().detach().numpy()) time_str = time.strftime("%Y-%m-%d-%H-%M-%S", time.localtime()) uid=''.join(random.sample(string.ascii_letters + string.digits, 16)) AIUrl='problems/images_tmp/%s_frame-%d-%d_type-%s_%s_s-%s_AI.jpg'%(time_str,best_frame,iimage_cnt,names[cls_max],uid,streamName) ORIUrl='problems/images_tmp/%s_frame-%d-%d_type-%s_%s_s-%s_OR.jpg'%(time_str,best_frame,iimage_cnt,names[cls_max],uid,streamName) cv2.imwrite(AIUrl,img_send); cv2.imwrite(ORIUrl,img_bak) outstr='%s save images to %s \n'%(time_str, AIUrl) fp_log.write(outstr ) problem_image=[[],[],[],[],[]] # Print time (inference + NMS) t3 = time_synchronized() # Save results (image with detections) if save_img: if dataset.mode == 'image': cv2.imwrite(save_path, im0) else: # 'video' or 'stream' if vid_path != save_path: # new video vid_path = save_path if isinstance(vid_writer, cv2.VideoWriter): vid_writer.release() # release previous video writer if vid_cap: # video fps = vid_cap.get(cv2.CAP_PROP_FPS) w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH)) h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) else: # stream fps, w, h = 30, im0.shape[1], im0.shape[0] save_path += '.mp4' vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h)) im0 = cv2.resize(im0,(OutVideoW,OutVideoH)) if dataset.mode == 'stream': ppipe.stdin.write(im0.tostring()) t4 = time_synchronized() #outstr='%s Done.read:%.1f ms, infer:%.1f ms, seginfer:%.1f ms,draw:%.1f ms, save:%.1f ms total:%.1f ms \n'%(s,(t1 - t0)*1000, (t2 - t1)*1000,(timeseg1-timeseg0)*1000, (t3 - t2)*1000,(t4-t3)*1000, (t4-t00)*1000) timestr=time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()) outstr='%s ,%s,iframe:%d,,read:%.1f ms,copy:%.1f, infer:%.1f ms, detinfer:%.1f ms,draw:%.1f ms, save:%.1f ms total:%.1f ms \n'%(s,timestr,iimage_cnt,(t0 - t00)*1000,(timeseg0-t0)*1000, (t1 - timeseg0)*1000,(t2-t1)*1000, (t3 - t2)*1000,(t4-t3)*1000, (t4-t00)*1000) if webcam: if len(det): fp_log.write(outstr ) else: print(outstr) print(segstr) sys.stdout.flush() t00=t4 if save_txt or save_img: s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else '' print(f"Results saved to {save_dir}{s}") print(f'Done. ({time.time() - t0:.3f}s)') if not webcam: break; except Exception as e: print('#######reading souce:%s ,error :%s:'%(source,e )) if Push_Flag and webcam: ####source end 推流or视频结束 ### ppipe.kill() fp_out.write(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()) +' rtmp stream-%s-%d ends \n'%(streamName,stream_id) );fp_out.flush() stream_id += 1 time_str = time.strftime("%Y-%m-%d-%H-%M-%S", time.localtime()) if 'off' in streamName:##只有离线视频时，才会写结束文件。 EndUrl='problems/images_tmp/%s_frame-9999-9999_type-结束_9999999999999999_s-%s_AI.jpg'%(time_str,streamName) img_end=np.zeros((100,100),dtype=np.uint8);cv2.imwrite(EndUrl,img_end) EndUrl='problems/images_tmp/%s_frame-9999-9999_type-结束_9999999999999999_s-%s_OR.jpg'%(time_str,streamName) cv2.imwrite(EndUrl,img_end) time_str = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()) with open( 'mintors/offlines/doneCodes.txt','a+' ) as fp: fp.write('%s %s\n'%(time_str,streamName )) #source_infos=update_websource_offAndLive(platform_query_url,sourceTxt,offlineFile) fp_log.close() if webcam: fp_out.write( time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())+' GPU server sleep 10s \n' ) ;fp_out.flush() time.sleep(10) if not webcam: break; ###update source (online or offline) source_infos=update_websource_offAndLive(platform_query_url,sourceTxt,offlineFile) if len(source_infos)==0: print('######NO valid source sleep 10s#####') time.sleep(10) if webcam: fp_out.close() if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument('--weights', nargs='+', type=str, default='yolov5s.pt', help='model.pt path(s)') parser.add_argument('--source', type=str, default='data/images', help='source') # file/folder, 0 for webcam parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)') parser.add_argument('--conf-thres', type=float, default=0.25, help='object confidence threshold') parser.add_argument('--iou-thres', type=float, default=0.45, help='IOU threshold for NMS') parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') parser.add_argument('--view-img', action='store_true', help='display results') parser.add_argument('--save-txt', action='store_true', help='save results to *.txt') parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels') parser.add_argument('--nosave', action='store_true', help='do not save images/videos') parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --class 0, or --class 0 2 3') parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS') parser.add_argument('--augment', action='store_true', help='augmented inference') parser.add_argument('--update', action='store_true', help='update all models') parser.add_argument('--project', default='runs/detect', help='save results to project/name') parser.add_argument('--name', default='exp', help='save results to project/name') parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment') parser.add_argument('--labelnames', type=str,default=None, help='labes nams') parser.add_argument('--fpsample', type=int, default=240, help='fpsample') parser.add_argument('--OutVideoW', type=int, default=1920, help='out video width size') parser.add_argument('--OutVideoH', type=int, default=1080, help='out video height size') parser.add_argument('--OutVideoFps', type=int, default=30, help='out video fps ') opt = parser.parse_args() print(opt) check_requirements(exclude=('pycocotools', 'thop')) with torch.no_grad(): if opt.update: # update all models (to fix SourceChangeWarning) for opt.weights in ['yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt']: detect() strip_optimizer(opt.weights) else: detect()