111 lines
4.8 KiB
Python
111 lines
4.8 KiB
Python
|
import sys
|
|||
|
|
sys.path.extend(['/home/thsw2/WJ/src/yolov5'])
|
|||
|
|
import utils,json,time,torch
|
|||
|
|
import numpy as np
|
|||
|
|
from segutils.segmodel import SegModel,get_largest_contours
|
|||
|
|
from models.experimental import attempt_load
|
|||
|
|
from utils.torch_utils import select_device, load_classifier, time_synchronized
|
|||
|
|
import subprocess as sp
|
|||
|
|
import cv2
|
|||
|
|
from utils.datasets import LoadStreams, LoadImages
|
|||
|
|
from queRiver import get_labelnames,get_label_arrays,post_process_,save_problem_images,time_str
|
|||
|
|
def get_total_cnt(inSource):
|
|||
|
|
cap=cv2.VideoCapture(inSource)
|
|||
|
|
cnt=cap.get(7)
|
|||
|
|
cap.release()
|
|||
|
|
return cnt
|
|||
|
|
def onlineModelProcess(parIn ):
|
|||
|
|
streamName = parIn['streamName']
|
|||
|
|
childCallback=parIn['callback']
|
|||
|
|
try:
|
|||
|
|
|
|||
|
|
inSource,outSource=parIn['inSource'],parIn['outSource']
|
|||
|
|
weights='../yolov5/weights/1230_last.pt'
|
|||
|
|
device = select_device('0')
|
|||
|
|
half = device.type != 'cpu' # half precision only supported on CUDA
|
|||
|
|
model = attempt_load(weights, map_location=device) # load FP32 model
|
|||
|
|
if half: model.half()
|
|||
|
|
seg_nclass = 2
|
|||
|
|
weights = '../yolov5/weights/segmentation/BiSeNet/checkpoint.pth'
|
|||
|
|
segmodel = SegModel(nclass=seg_nclass,weights=weights,device=device)
|
|||
|
|
jsonfile='../yolov5/config/queRiver.json'
|
|||
|
|
with open(jsonfile,'r') as fp:
|
|||
|
|
parAll = json.load(fp)
|
|||
|
|
|
|||
|
|
|
|||
|
|
resource=parAll['prep_process']['source']
|
|||
|
|
if outSource:
|
|||
|
|
command=['ffmpeg','-y','-f', 'rawvideo','-vcodec','rawvideo','-pix_fmt', 'bgr24',
|
|||
|
|
'-s', "{}x{}".format(parAll["push_process"]['OutVideoW'],parAll["push_process"]['OutVideoH']),# 图片分辨率
|
|||
|
|
'-r', str(30),# 视频帧率
|
|||
|
|
'-i', '-','-c:v', 'libx264','-pix_fmt', 'yuv420p',
|
|||
|
|
'-f', 'flv',outSource
|
|||
|
|
]
|
|||
|
|
txtname='mintors/%s.txt'%( time.strftime("%Y-%m-%d", time.localtime()) )
|
|||
|
|
fp_out = open( txtname,'a+' )
|
|||
|
|
outstr='%s stream:%s starts \n'%( time_str(),parAll['push_process']['rtmpUrl'])
|
|||
|
|
fp_out.write(outstr);fp_out.flush()
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
# 管道配置，其中用到管道
|
|||
|
|
if outSource:
|
|||
|
|
ppipe = sp.Popen(command, stdin=sp.PIPE)
|
|||
|
|
|
|||
|
|
##后处理参数
|
|||
|
|
par=parAll['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)
|
|||
|
|
|
|||
|
|
dataset = LoadStreams(inSource, img_size=640, stride=32)
|
|||
|
|
if (inSource.endswith('.MP4')) or (inSource.endswith('.mp4')):
|
|||
|
|
totalcnt=get_total_cnt(inSource)
|
|||
|
|
childCallback.send('####model load success####')
|
|||
|
|
iframe = 0;post_results=[];time_beg=time.time()
|
|||
|
|
print('###line71 modelEval.py',totalcnt,len(dataset), inSource)
|
|||
|
|
for path, img, im0s, vid_cap in dataset:
|
|||
|
|
print(path)
|
|||
|
|
if not path:childCallback.send('####strem ends####'); break###断流或者到终点
|
|||
|
|
if not outSource:###如果不推流，则显示进度条
|
|||
|
|
view_bar(iframe,totalcnt,time_beg )
|
|||
|
|
time0=time.time()
|
|||
|
|
iframe +=1
|
|||
|
|
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,segstr = segmodel.eval(im0s[0] )
|
|||
|
|
|
|||
|
|
time4 = time.time()
|
|||
|
|
datas = [path, img, im0s, vid_cap,pred,seg_pred,iframe]
|
|||
|
|
|
|||
|
|
p_result,timeOut = post_process_(datas,conf_thres, iou_thres,names,label_arraylist,rainbows,iframe)
|
|||
|
|
##每隔 fpsample帧处理一次，如果有问题就保存图片
|
|||
|
|
if (iframe % fpsample == 0) and (len(post_results)>0) :
|
|||
|
|
parImage=save_problem_images(post_results,iframe,names,streamName=streamName)
|
|||
|
|
|
|||
|
|
#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]
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
post_results=[]
|
|||
|
|
|
|||
|
|
if len(p_result[2] )>0: ##
|
|||
|
|
post_results.append(p_result)
|
|||
|
|
|
|||
|
|
image_array = p_result[1]
|
|||
|
|
if outSource:
|
|||
|
|
ppipe.stdin.write(image_array.tostring())
|
|||
|
|
except Exception as e:
|
|||
|
|
childCallback.send(e) #将异常通过管道送出
|