You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

modelEval.py 4.8KB

1 year ago
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110
  1. import sys
  2. sys.path.extend(['/home/thsw2/WJ/src/yolov5'])
  3. import utils,json,time,torch
  4. import numpy as np
  5. from segutils.segmodel import SegModel,get_largest_contours
  6. from models.experimental import attempt_load
  7. from utils.torch_utils import select_device, load_classifier, time_synchronized
  8. import subprocess as sp
  9. import cv2
  10. from utils.datasets import LoadStreams, LoadImages
  11. from queRiver import get_labelnames,get_label_arrays,post_process_,save_problem_images,time_str
  12. def get_total_cnt(inSource):
  13. cap=cv2.VideoCapture(inSource)
  14. cnt=cap.get(7)
  15. cap.release()
  16. return cnt
  17. def onlineModelProcess(parIn ):
  18. streamName = parIn['streamName']
  19. childCallback=parIn['callback']
  20. try:
  21. inSource,outSource=parIn['inSource'],parIn['outSource']
  22. weights='../yolov5/weights/1230_last.pt'
  23. device = select_device('0')
  24. half = device.type != 'cpu' # half precision only supported on CUDA
  25. model = attempt_load(weights, map_location=device) # load FP32 model
  26. if half: model.half()
  27. seg_nclass = 2
  28. weights = '../yolov5/weights/segmentation/BiSeNet/checkpoint.pth'
  29. segmodel = SegModel(nclass=seg_nclass,weights=weights,device=device)
  30. jsonfile='../yolov5/config/queRiver.json'
  31. with open(jsonfile,'r') as fp:
  32. parAll = json.load(fp)
  33. resource=parAll['prep_process']['source']
  34. if outSource:
  35. command=['ffmpeg','-y','-f', 'rawvideo','-vcodec','rawvideo','-pix_fmt', 'bgr24',
  36. '-s', "{}x{}".format(parAll["push_process"]['OutVideoW'],parAll["push_process"]['OutVideoH']),# 图片分辨率
  37. '-r', str(30),# 视频帧率
  38. '-i', '-','-c:v', 'libx264','-pix_fmt', 'yuv420p',
  39. '-f', 'flv',outSource
  40. ]
  41. txtname='mintors/%s.txt'%( time.strftime("%Y-%m-%d", time.localtime()) )
  42. fp_out = open( txtname,'a+' )
  43. outstr='%s stream:%s starts \n'%( time_str(),parAll['push_process']['rtmpUrl'])
  44. fp_out.write(outstr);fp_out.flush()
  45. # 管道配置,其中用到管道
  46. if outSource:
  47. ppipe = sp.Popen(command, stdin=sp.PIPE)
  48. ##后处理参数
  49. par=parAll['post_process']
  50. conf_thres,iou_thres,classes=par['conf_thres'],par['iou_thres'],par['classes']
  51. labelnames=par['labelnames']
  52. rainbows=par['rainbows']
  53. fpsample = par['fpsample']
  54. names=get_labelnames(labelnames)
  55. label_arraylist = get_label_arrays(names,rainbows,outfontsize=40)
  56. dataset = LoadStreams(inSource, img_size=640, stride=32)
  57. if (inSource.endswith('.MP4')) or (inSource.endswith('.mp4')):
  58. totalcnt=get_total_cnt(inSource)
  59. childCallback.send('####model load success####')
  60. iframe = 0;post_results=[];time_beg=time.time()
  61. print('###line71 modelEval.py',totalcnt,len(dataset), inSource)
  62. for path, img, im0s, vid_cap in dataset:
  63. print(path)
  64. if not path:childCallback.send('####strem ends####'); break###断流或者到终点
  65. if not outSource:###如果不推流,则显示进度条
  66. view_bar(iframe,totalcnt,time_beg )
  67. time0=time.time()
  68. iframe +=1
  69. time1=time.time()
  70. img = torch.from_numpy(img).to(device)
  71. img = img.half() if half else img.float() # uint8 to fp16/32
  72. img /= 255.0 # 0 - 255 to 0.0 - 1.0
  73. time2 = time.time()
  74. pred = model(img,augment=False)[0]
  75. time3 = time.time()
  76. seg_pred,segstr = segmodel.eval(im0s[0] )
  77. time4 = time.time()
  78. datas = [path, img, im0s, vid_cap,pred,seg_pred,iframe]
  79. p_result,timeOut = post_process_(datas,conf_thres, iou_thres,names,label_arraylist,rainbows,iframe)
  80. ##每隔 fpsample帧处理一次,如果有问题就保存图片
  81. if (iframe % fpsample == 0) and (len(post_results)>0) :
  82. parImage=save_problem_images(post_results,iframe,names,streamName=streamName)
  83. #parOut = {}; parOut['imgOR'] = img_send; parOut['imgAR'] = img_send; parOut['uid']=uid
  84. #parOut['imgORname']=os.path.basename(outnameOR);parOut['imgARname']=os.path.basename(outnameAR);
  85. #parOut['time_str'] = time_str;parOut['type'] = names[cls_max]
  86. post_results=[]
  87. if len(p_result[2] )>0: ##
  88. post_results.append(p_result)
  89. image_array = p_result[1]
  90. if outSource:
  91. ppipe.stdin.write(image_array.tostring())
  92. except Exception as e:
  93. childCallback.send(e) #将异常通过管道送出