diff --git a/concurrency/PullVideoStreamProcess.py b/concurrency/PullVideoStreamProcess.py index 54fe030..4825cd6 100644 --- a/concurrency/PullVideoStreamProcess.py +++ b/concurrency/PullVideoStreamProcess.py @@ -129,7 +129,7 @@ class OnlinePullVideoStreamProcess(PullVideoStreamProcess): frame, pull_p, width, height = pull_read_video_stream(pull_p, pull_url, width, height, width_height_3, w_2, h_2, request_id) if pull_queue.full(): - logger.info("pull拉流队列满了:{}, requestId: {}", os.getppid(), request_id) + #logger.info("pull拉流队列满了:{}, requestId: {}", os.getppid(), request_id) if full_timeout is None: full_timeout = time() if time() - full_timeout > 180: @@ -269,7 +269,7 @@ class OfflinePullVideoStreamProcess(PullVideoStreamProcess): width, height, width_height_3, all_frames, w_2, h_2 = build_video_info(pull_url, request_id) continue if pull_queue.full(): - logger.info("pull拉流队列满了:{}, requestId: {}", os.getppid(), request_id) + #logger.info("pull拉流队列满了:{}, requestId: {}", os.getppid(), request_id) if full_timeout is None: full_timeout = time() if time() - full_timeout > 180: diff --git a/concurrency/PushVideoStreamProcess.py b/concurrency/PushVideoStreamProcess.py index 997d800..56d4852 100644 --- a/concurrency/PushVideoStreamProcess.py +++ b/concurrency/PushVideoStreamProcess.py @@ -36,7 +36,7 @@ class PushStreamProcess(Process): # 传参 self._msg, self._push_queue, self._image_queue, self._push_ex_queue, self._hb_queue, self._context = args self._algStatus = False # 默认关闭 - self._algSwitch = self._context['service']['algSwitch'] + self._algSwitch = self._context['service']['algSwitch'] # 0521: default_enabled = str(self._msg.get("defaultEnabled", "True")).lower() == "true" @@ -139,6 +139,7 @@ class OnPushStreamProcess(PushStreamProcess): # 所有问题的矩阵集合 qs_np = None qs_reurn = [] + bp_np = None for det in push_objs[i]: code, det_result = det # 每个单独模型处理 @@ -167,12 +168,13 @@ class OnPushStreamProcess(PushStreamProcess): else: try: # 应对NaN情况 box, score, cls = xywh2xyxy2(qs) + if cls not in allowedList or score < frame_score: + continue if ModelType.CITY_FIREAREA_MODEL.value[1] == str(code): + # 借score作为points点集 box.append(qs[-1]) except: continue - if cls not in allowedList or score < frame_score: - continue label_array, color = label_arrays[cls], rainbows[cls] if ModelType.CHANNEL2_MODEL.value[1] == str(code) and cls == 2: rr = t.submit(draw_name_joint, box, copy_frame, @@ -190,17 +192,24 @@ class OnPushStreamProcess(PushStreamProcess): if cd is None: det_xywh[code][cls] = [[cls, box, score, label_array, color]] else: - det_xywh[code][cls].append([cls, box, score, label_array, color]) + det_xywh[code][cls].append([cls, box, score, label_array, color]) if qs_np is None: - qs_np = np.array([box[0][0], box[0][1], box[1][0], box[1][1], - box[2][0], box[2][1], box[3][0], box[3][1], + qs_np = np.array([box[0][0], box[0][1], box[1][0], box[1][1], + box[2][0], box[2][1], box[3][0], box[3][1], score, cls, code],dtype=np.float32) else: - result_li = np.array([box[0][0], box[0][1], box[1][0], box[1][1], + result_li = np.array([box[0][0], box[0][1], box[1][0], box[1][1], box[2][0], box[2][1], box[3][0], box[3][1], score, cls, code],dtype=np.float32) qs_np = np.row_stack((qs_np, result_li)) + if ModelType.CITY_FIREAREA_MODEL.value[1] == str(code): + if bp_np is None: + bp_np = np.array([box[0][0], box[0][1], box[-1]], dtype=object) + else: + bp_li = np.array([box[0][0], box[0][1], box[-1]], dtype=object) + bp_np = np.row_stack((bp_np, bp_li)) + if logo: frame = add_water_pic(frame, logo, request_id) copy_frame = add_water_pic(copy_frame, logo, request_id) @@ -231,7 +240,7 @@ class OnPushStreamProcess(PushStreamProcess): if picture_similarity: qs_np_tmp = qs_np_id.copy() b = np.zeros(qs_np.shape[0]) - qs_reurn = np.column_stack((qs_np,b)) + qs_reurn = np.column_stack((qs_np,b)) else: qs_reurn = filterBox(qs_np, qs_np_tmp, pix_dis) if picture_similarity: @@ -259,8 +268,12 @@ class OnPushStreamProcess(PushStreamProcess): label_array, color = label_arrays[cls], rainbows[cls] box = [(int(q[0]), int(q[1])), (int(q[2]), int(q[3])), (int(q[4]), int(q[5])), (int(q[6]), int(q[7]))] - if ModelType.CITY_FIREAREA_MODEL.value[1] == str(code): - box.append(qs[-1]) + if bp_np is not None: + if len(bp_np.shape)==1: + bp_np = bp_np[np.newaxis, ...] + for bp in bp_np: + if np.array_equal(bp[:2], np.array([int(q[0]), int(q[1])])): + box.append(bp[-1]) is_new = False if q[11] == 1: is_new = True @@ -284,7 +297,7 @@ class OnPushStreamProcess(PushStreamProcess): if push_r[0] == 2: logger.info("拉流进程收到控制命令为:{}, requestId: {}",push_r[1] ,request_id) if 'algStart' == push_r[1]: self._algStatus = True;logger.info("算法识别开启, requestId: {}", request_id) - if 'algStop' == push_r[1]: self._algStatus = False;logger.info("算法识别关闭, requestId: {}", request_id) + if 'algStop' == push_r[1]: self._algStatus = False;logger.info("算法识别关闭, requestId: {}", request_id) if 'stop' == push_r[1]: logger.info("停止推流进程, requestId: {}", request_id) break @@ -385,9 +398,9 @@ class OffPushStreamProcess(PushStreamProcess): # 所有问题的矩阵集合 qs_np = None qs_reurn = [] + bp_np = None for det in push_objs[i]: code, det_result = det - # 每个单独模型处理 # 模型编号、100帧的所有问题, 检测目标、颜色、文字图片 if len(det_result) > 0: @@ -415,10 +428,11 @@ class OffPushStreamProcess(PushStreamProcess): else: box, score, cls = xywh2xyxy2(qs) - if ModelType.CITY_FIREAREA_MODEL.value[1] == str(code): - box.append(qs[-1]) if cls not in allowedList or score < frame_score: continue + if ModelType.CITY_FIREAREA_MODEL.value[1] == str(code): + box.append(qs[-1]) + label_array, color = label_arrays[cls], rainbows[cls] if ModelType.CHANNEL2_MODEL.value[1] == str(code) and cls == 2: rr = t.submit(draw_name_joint, box, copy_frame, draw_config[code]["label_dict"], score, color, font_config, qs[6]) @@ -426,7 +440,6 @@ class OffPushStreamProcess(PushStreamProcess): rr = t.submit(draw_painting_joint, box, copy_frame, label_array, score, color, font_config) thread_p.append(rr) - if det_xywh.get(code) is None: det_xywh[code] = {} cd = det_xywh[code].get(cls) @@ -445,6 +458,13 @@ class OffPushStreamProcess(PushStreamProcess): score, cls, code],dtype=np.float32) qs_np = np.row_stack((qs_np, result_li)) + if ModelType.CITY_FIREAREA_MODEL.value[1]== str(code): + if bp_np is None: + bp_np = np.array([box[0][0], box[0][1],box[-1]],dtype=object) + else: + bp_li = np.array([box[0][0], box[0][1],box[-1]],dtype=object) + bp_np = np.row_stack((bp_np, bp_li)) + if logo: frame = add_water_pic(frame, logo, request_id) copy_frame = add_water_pic(copy_frame, logo, request_id) @@ -472,7 +492,7 @@ class OffPushStreamProcess(PushStreamProcess): if picture_similarity: qs_np_tmp = qs_np_id.copy() b = np.zeros(qs_np.shape[0]) - qs_reurn = np.column_stack((qs_np,b)) + qs_reurn = np.column_stack((qs_np,b)) else: qs_reurn = filterBox(qs_np, qs_np_tmp, pix_dis) if picture_similarity: @@ -501,8 +521,12 @@ class OffPushStreamProcess(PushStreamProcess): label_array, color = label_arrays[cls], rainbows[cls] box = [(int(q[0]), int(q[1])), (int(q[2]), int(q[3])), (int(q[4]), int(q[5])), (int(q[6]), int(q[7]))] - if ModelType.CITY_FIREAREA_MODEL.value[1] == str(code): - box.append(qs[-1]) + if bp_np is not None: + if len(bp_np.shape)==1: + bp_np = bp_np[np.newaxis, ...] + for bp in bp_np: + if np.array_equal(bp[:2], np.array([int(q[0]), int(q[1])])): + box.append(bp[-1]) is_new = False if q[11] == 1: is_new = True @@ -524,7 +548,7 @@ class OffPushStreamProcess(PushStreamProcess): if push_r[0] == 2: logger.info("拉流进程收到控制命令为:{}, requestId: {}",push_r[1] ,request_id) if 'algStart' == push_r[1]: self._algStatus = True;logger.info("算法识别开启, requestId: {}", request_id) - if 'algStop' == push_r[1]: self._algStatus = False;logger.info("算法识别关闭, requestId: {}", request_id) + if 'algStop' == push_r[1]: self._algStatus = False;logger.info("算法识别关闭, requestId: {}", request_id) if 'stop' == push_r[1]: logger.info("停止推流进程, requestId: {}", request_id) break diff --git a/enums/ModelTypeEnum.py b/enums/ModelTypeEnum.py index 9f7febb..c506ea8 100644 --- a/enums/ModelTypeEnum.py +++ b/enums/ModelTypeEnum.py @@ -14,6 +14,7 @@ from utilsK.drownUtils import mixDrowing_water_postprocess from utilsK.noParkingUtils import mixNoParking_road_postprocess from utilsK.illParkingUtils import illParking_postprocess from utilsK.pannelpostUtils import pannel_post_process +from utilsK.securitypostUtils import security_post_process from stdc import stdcModel from yolov5 import yolov5Model from p2pNet import p2NnetModel @@ -1044,7 +1045,7 @@ class ModelType(Enum): CITY_FIREAREA_MODEL = ("30", "307", "火焰面积模型", 'FireArea', lambda device, gpuName: { 'device': device, 'gpu_name': gpuName, - 'labelnames': ["火焰面积"], + 'labelnames': ["火焰"], 'seg_nclass': 2, # 分割模型类别数目，默认2类 'segRegionCnt': 0, 'trtFlag_det': True, @@ -1068,6 +1069,27 @@ class ModelType(Enum): }) + CITY_SECURITY_MODEL = ("30", "308", "安防模型", 'SECURITY', lambda device, gpuName: { + 'labelnames': ["带安全帽","安全帽","攀爬","斗殴","未戴安全帽"], + 'postProcess': {'function': security_post_process, 'pars': {'objs': [0,1],'iou':0.25,'unhelmet':4}}, + 'models': + [ + { + 'weight': "../weights/trt/AIlib2/security/yolov5_%s_fp16.engine" % (gpuName), ###检测模型路径 + 'name': 'yolov5', + 'model': yolov5Model, + 'par': {'half': True, 'device': 'cuda:0', 'conf_thres': 0.25, 'iou_thres': 0.45, + 'segRegionCnt': 1, 'trtFlag_det': True, 'trtFlag_seg': False}, + } + + ], + 'postFile': { + "rainbows": COLOR + }, + 'fiterList': [0,1], + "score_byClass": {"0": 0.50} + }) + @staticmethod def checkCode(code): for model in ModelType: diff --git a/util/ModelUtils.py b/util/ModelUtils.py index 7123efb..d3e9f7b 100644 --- a/util/ModelUtils.py +++ b/util/ModelUtils.py @@ -772,4 +772,11 @@ MODEL_CONFIG = { lambda x, y, z: one_label(x, y, z), lambda x: forest_process(x) ), + # 加载安防模型 + ModelType.CITY_SECURITY_MODEL.value[1]: ( + lambda x, y, r, t, z, h: cityManagementModel(x, y, r, ModelType.CITY_SECURITY_MODEL, t, z, h), + ModelType.CITY_SECURITY_MODEL, + lambda x, y, z: one_label(x, y, z), + lambda x: detSeg_demo2(x) + ), } diff --git a/util/PlotsUtils.py b/util/PlotsUtils.py index 57637a4..72b5f76 100644 --- a/util/PlotsUtils.py +++ b/util/PlotsUtils.py @@ -5,7 +5,7 @@ import unicodedata from loguru import logger FONT_PATH = "../AIlib2/conf/platech.ttf" -zhFont = ImageFont.truetype(FONT_PATH, 20, encoding="utf-8") +zhFont = ImageFont.truetype(FONT_PATH, 20, encoding="utf-8") def get_label_array(color=None, label=None, font=None, fontSize=40, unify=False): if unify: @@ -50,6 +50,48 @@ def get_label_array_dict(colors, fontSize=40, fontPath="platech.ttf"): zh_dict[code] = arr return zh_dict +def get_label_left(x0,y1,label_array,img): + imh, imw = img.shape[0:2] + lh, lw = label_array.shape[0:2] + # x1 框框左上x位置 + 描述的宽 + # y0 框框左上y位置 - 描述的高 + x1, y0 = x0 + lw, y1 - lh + # 如果y0小于0, 说明超过上边框 + if y0 < 0: + y0 = 0 + # y1等于文字高度 + y1 = y0 + lh + # 如果y1框框的高大于图片高度 + if y1 > imh: + # y1等于图片高度 + y1 = imh + # y0等于y1减去文字高度 + y0 = y1 - lh + # 如果x0小于0 + if x0 < 0: + x0 = 0 + x1 = x0 + lw + if x1 > imw: + x1 = imw + x0 = x1 - lw + return x0,y0,x1,y1 + +def get_label_right(x1,y0,label_array): + lh, lw = label_array.shape[0:2] + # x1 框框右上x位置 + 描述的宽 + # y0 框框右上y位置 - 描述的高 + x0, y1 = x1 - lw, y0 - lh + # 如果y0小于0, 说明超过上边框 + if y0 < 0 or y1 < 0: + y1 = 0 + # y1等于文字高度 + y0 = y1 + lh + # 如果x0小于0 + if x0 < 0 or x1 < 0: + x0 = 0 + x1 = x0 + lw + + return x0,y1,x1,y0 def xywh2xyxy(box): if not isinstance(box[0], (list, tuple, np.ndarray)): @@ -78,48 +120,41 @@ def xy2xyxy(box): def draw_painting_joint(box, img, label_array, score=0.5, color=None, config=None, isNew=False): # 识别问题描述图片的高、宽 # 图片的长度和宽度 - imh, imw = img.shape[0:2] - points = np.array([]) + lh, lw = label_array.shape[0:2] + tl = config[0] if isinstance(box[-1], np.ndarray): font = ImageFont.truetype(FONT_PATH, 10, encoding='utf-8') points = box[-1] arrea = cv2.contourArea(points) - label = '火焰面积:%s'%arrea + label = '火焰' + #arealabel = '面积:%s'%f"{arrea:.1e}" + arealabel = '面积:%s'%arrea + label_array_area = get_label_array(color, arealabel, font, 10) label_array = get_label_array(color, label, font, 10) - else: - label = ' %.2f' % score - lh, lw = label_array.shape[0:2] - box = xywh2xyxy(box[:4]) + lh_area, lw_area = label_array_area.shape[0:2] + box = box[:4] + # 框框左上的位置 + x0, y1 = box[0][0], max(box[0][1]-lh_area-3,0) + x1, y0 = box[1][0], box[1][1] + x0_label,y0_label,x1_label,y1_label = get_label_left(x0, y1,label_array,img) + x0_area, y0_area, x1_area, y1_area = get_label_right(x1, y0, label_array_area) + img[y0_label:y1_label, x0_label:x1_label, :] = label_array + img[y0_area:y1_area, x0_area:x1_area, :] = label_array_area + #cv2.drawContours(img, points, -1, color, tl) + cv2.polylines(img, [points], False, color, 1) + if lw_area < box[1][0]-box[0][0]: + box = [(x0, y1), (x1, y1),(x1, box[2][1]),(x0, box[2][1])] + else: + box = [(x0_label, y1), (x1, y1), (x1, box[2][1]), (x0_label, box[2][1])] + box = np.asarray(box, np.int32) + cv2.polylines(img, [box], True, color, 1) + return img, box + + label = ' %.2f' % score + box = xywh2xyxy(box) # 框框左上的位置 x0, y1 = box[0][0], box[0][1] - # if score_location == 'leftTop': - # x0, y1 = box[0][0], box[0][1] - # # 框框左下的位置 - # elif score_location == 'leftBottom': - # x0, y1 = box[3][0], box[3][1] - # else: - # x0, y1 = box[0][0], box[0][1] - # x1 框框左上x位置 + 描述的宽 - # y0 框框左上y位置 - 描述的高 - x1, y0 = x0 + lw, y1 - lh - # 如果y0小于0, 说明超过上边框 - if y0 < 0: - y0 = 0 - # y1等于文字高度 - y1 = y0 + lh - # 如果y1框框的高大于图片高度 - if y1 > imh: - # y1等于图片高度 - y1 = imh - # y0等于y1减去文字高度 - y0 = y1 - lh - # 如果x0小于0 - if x0 < 0: - x0 = 0 - x1 = x0 + lw - if x1 > imw: - x1 = imw - x0 = x1 - lw + x0, y0, x1, y1 = get_label_left(x0, y1, label_array, img) # box_tl = max(int(round(imw / 1920 * 3)), 1) or round(0.002 * (imh + imw) / 2) + 1 ''' 1. img（array） 为ndarray类型（可以为cv.imread）直接读取的数据 @@ -129,14 +164,9 @@ def draw_painting_joint(box, img, label_array, score=0.5, color=None, config=Non 5. thickness（int）：画线的粗细 6. shift：顶点坐标中小数的位数 ''' - tl = config[0] img[y0:y1, x0:x1, :] = label_array - if points.size != 0: - cv2.drawContours(img, points, -1, color, 2) - return img, box - else: - box1 = np.asarray(box, np.int32) - cv2.polylines(img, [box1], True, color, tl) + box1 = np.asarray(box, np.int32) + cv2.polylines(img, [box1], True, color, tl) pts_cls = [(x0, y0), (x1, y1)] # 把英文字符score画到类别旁边 # tl = max(int(round(imw / 1920 * 3)), 1) or round(0.002 * (imh + imw) / 2) + 1