3 changed files with 236 additions and 83 deletions
+ 99
- 43
concurrency/PushVideoStreamProcess.py
View File
| @@ -9,6 +9,7 @@ from time import time, sleep | |||
| from traceback import format_exc | |||
| import cv2 | |||
| import numpy as np | |||
| import psutil | |||
| from loguru import logger | |||
| @@ -20,7 +21,7 @@ from util.Cv2Utils import video_conjuncing, write_or_video, write_ai_video, push | |||
| from util.ImageUtils import url2Array, add_water_pic | |||
| from util.LogUtils import init_log | |||
| from util.PlotsUtils import draw_painting_joint, xywh2xyxy2 | |||
| from util.PlotsUtils import draw_painting_joint, filterBox, xywh2xyxy2 | |||
| from util.QueUtil import get_no_block_queue, put_queue, clear_queue | |||
| @@ -82,12 +83,10 @@ class OnPushStreamProcess(PushStreamProcess): | |||
| self._hb_queue | |||
| or_video_file, ai_video_file, push_p, ex = None, None, None, None | |||
| ex_status = True | |||
| high_score_image = {} | |||
| # 相似度, 默认值0.65 | |||
| similarity = context["service"]["filter"]["similarity"] | |||
| # 图片相似度开关 | |||
| picture_similarity = bool(context["service"]["filter"]["picture_similarity"]) | |||
| frame_step = int(context["service"]["filter"]["frame_step"]) | |||
| qs_np_tmp = None | |||
| pix_dis = 60 | |||
| try: | |||
| init_log(base_dir, env) | |||
| logger.info("开始实时启动推流进程！requestId:{}", request_id) | |||
| @@ -118,6 +117,10 @@ class OnPushStreamProcess(PushStreamProcess): | |||
| # 复制帧用来画图 | |||
| copy_frame = frame.copy() | |||
| det_xywh, thread_p = {}, [] | |||
| det_xywh2 = {} | |||
| # 所有问题的矩阵集合 | |||
| qs_np = None | |||
| qs_reurn = [] | |||
| for det in push_objs[i]: | |||
| code, det_result = det | |||
| # 每个单独模型处理 | |||
| @@ -137,8 +140,15 @@ class OnPushStreamProcess(PushStreamProcess): | |||
| cd = det_xywh[code].get(cls) | |||
| if cd is None: | |||
| det_xywh[code][cls] = [[cls, box, score, label_array, color]] | |||
| 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: | |||
| det_xywh[code][cls].append([cls, box, score, label_array, color]) | |||
| 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 logo: | |||
| frame = add_water_pic(frame, logo, request_id) | |||
| copy_frame = add_water_pic(copy_frame, logo, request_id) | |||
| @@ -155,23 +165,42 @@ class OnPushStreamProcess(PushStreamProcess): | |||
| push_stream_result = t.submit(push_video_stream, frame_merge, push_p, push_url, | |||
| p_push_status, request_id) | |||
| # 如果有问题, 走下面的逻辑 | |||
| if len(det_xywh) > 0: | |||
| flag = True | |||
| if picture_similarity and len(high_score_image) > 0: | |||
| hash1 = ImageUtils.dHash(high_score_image.get("or_frame")) | |||
| hash2 = ImageUtils.dHash(frame) | |||
| dist = ImageUtils.Hamming_distance(hash1, hash2) | |||
| similarity_1 = 1 - dist * 1.0 / 64 | |||
| if similarity_1 >= similarity: | |||
| flag = False | |||
| if len(high_score_image) > 0: | |||
| diff_frame_num = frame_index_list[i] - high_score_image.get("current_frame") | |||
| if diff_frame_num < frame_step: | |||
| flag = False | |||
| if flag: | |||
| high_score_image["or_frame"] = frame | |||
| high_score_image["current_frame"] = frame_index_list[i] | |||
| put_queue(image_queue, (1, [det_xywh, frame, frame_index_list[i], all_frames, draw_config["font_config"]])) | |||
| if qs_np is not None: | |||
| if len(qs_np.shape) == 1: | |||
| qs_np = qs_np[np.newaxis,...] | |||
| qs_np_id = qs_np.copy() | |||
| b = np.ones(qs_np_id.shape[0]) | |||
| qs_np_id = np.column_stack((qs_np_id,b)) | |||
| if qs_np_tmp is None: | |||
| 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)) | |||
| else: | |||
| qs_reurn = filterBox(qs_np, qs_np_tmp, pix_dis) | |||
| if picture_similarity: | |||
| qs_np_tmp = np.append(qs_np_tmp,qs_np_id,axis=0) | |||
| qs_np_tmp[:, 11] += 1 | |||
| qs_np_tmp = np.delete(qs_np_tmp, np.where((qs_np_tmp[:, 11] >= 75))[0], axis=0) | |||
| for q in qs_reurn: | |||
| if q[11] == 0: | |||
| cls = int(q[9]) | |||
| code = str(int(q[10])).zfill(3) | |||
| if det_xywh2.get(code) is None: | |||
| det_xywh2[code] = {} | |||
| cd = det_xywh2[code].get(cls) | |||
| score = q[8] | |||
| rainbows, label_arrays = draw_config[code]["rainbows"], draw_config[code]["label_arrays"] | |||
| 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 cd is None: | |||
| det_xywh2[code][cls] = [[cls, box, score, label_array, color]] | |||
| else: | |||
| det_xywh2[code][cls].append([cls, box, score, label_array, color]) | |||
| if len(det_xywh2) > 0: | |||
| put_queue(image_queue, (1, [det_xywh2, frame, frame_index_list[i], all_frames, draw_config["font_config"]])) | |||
| push_p = push_stream_result.result(timeout=60) | |||
| ai_video_file = write_ai_video_result.result(timeout=60) | |||
| or_video_file = write_or_video_result.result(timeout=60) | |||
| @@ -228,12 +257,10 @@ class OffPushStreamProcess(PushStreamProcess): | |||
| frame_score = context["service"]["filter"]["frame_score"] | |||
| ex = None | |||
| ex_status = True | |||
| high_score_image = {} | |||
| # 相似度, 默认值0.65 | |||
| similarity = context["service"]["filter"]["similarity"] | |||
| # 图片相似度开关 | |||
| picture_similarity = bool(context["service"]["filter"]["picture_similarity"]) | |||
| frame_step = int(context["service"]["filter"]["frame_step"]) | |||
| qs_np_tmp = None | |||
| pix_dis = 60 | |||
| try: | |||
| init_log(base_dir, env) | |||
| logger.info("开始启动离线推流进程！requestId:{}", request_id) | |||
| @@ -271,6 +298,10 @@ class OffPushStreamProcess(PushStreamProcess): | |||
| copy_frame = frame.copy() | |||
| # 所有问题记录字典 | |||
| det_xywh, thread_p = {}, [] | |||
| det_xywh2 = {} | |||
| # 所有问题的矩阵集合 | |||
| qs_np = None | |||
| qs_reurn = [] | |||
| for det in push_objs[i]: | |||
| code, det_result = det | |||
| # 每个单独模型处理 | |||
| @@ -290,8 +321,15 @@ class OffPushStreamProcess(PushStreamProcess): | |||
| cd = det_xywh[code].get(cls) | |||
| if cd is None: | |||
| det_xywh[code][cls] = [[cls, box, score, label_array, color]] | |||
| 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: | |||
| det_xywh[code][cls].append([cls, box, score, label_array, color]) | |||
| 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 logo: | |||
| frame = add_water_pic(frame, logo, request_id) | |||
| copy_frame = add_water_pic(copy_frame, logo, request_id) | |||
| @@ -306,23 +344,41 @@ class OffPushStreamProcess(PushStreamProcess): | |||
| push_stream_result = t.submit(push_video_stream, frame_merge, push_p, push_url, | |||
| p_push_status, request_id) | |||
| if len(det_xywh) > 0: | |||
| flag = True | |||
| if picture_similarity and len(high_score_image) > 0: | |||
| hash1 = ImageUtils.dHash(high_score_image.get("or_frame")) | |||
| hash2 = ImageUtils.dHash(frame) | |||
| dist = ImageUtils.Hamming_distance(hash1, hash2) | |||
| similarity_1 = 1 - dist * 1.0 / 64 | |||
| if similarity_1 >= similarity: | |||
| flag = False | |||
| if len(high_score_image) > 0: | |||
| diff_frame_num = frame_index_list[i] - high_score_image.get("current_frame") | |||
| if diff_frame_num < frame_step: | |||
| flag = False | |||
| if flag: | |||
| high_score_image["or_frame"] = frame | |||
| high_score_image["current_frame"] = frame_index_list[i] | |||
| put_queue(image_queue, (1, [det_xywh, frame, frame_index_list[i], all_frames, draw_config["font_config"]])) | |||
| if qs_np is not None: | |||
| if len(qs_np.shape) == 1: | |||
| qs_np = qs_np[np.newaxis,...] | |||
| qs_np_id = qs_np.copy() | |||
| b = np.ones(qs_np_id.shape[0]) | |||
| qs_np_id = np.column_stack((qs_np_id,b)) | |||
| if qs_np_tmp is None: | |||
| 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)) | |||
| else: | |||
| qs_reurn = filterBox(qs_np, qs_np_tmp, pix_dis) | |||
| if picture_similarity: | |||
| qs_np_tmp = np.append(qs_np_tmp,qs_np_id,axis=0) | |||
| qs_np_tmp[:, 11] += 1 | |||
| qs_np_tmp = np.delete(qs_np_tmp, np.where((qs_np_tmp[:, 11] >= 75))[0], axis=0) | |||
| for q in qs_reurn: | |||
| if q[11] == 0: | |||
| cls = int(q[9]) | |||
| code = str(int(q[10])).zfill(3) | |||
| if det_xywh2.get(code) is None: | |||
| det_xywh2[code] = {} | |||
| cd = det_xywh2[code].get(cls) | |||
| score = q[8] | |||
| rainbows, label_arrays = draw_config[code]["rainbows"], draw_config[code]["label_arrays"] | |||
| 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 cd is None: | |||
| det_xywh2[code][cls] = [[cls, box, score, label_array, color]] | |||
| else: | |||
| det_xywh2[code][cls].append([cls, box, score, label_array, color]) | |||
| if len(det_xywh2) > 0: | |||
| put_queue(image_queue, (1, [det_xywh2, frame, frame_index_list[i], all_frames, draw_config["font_config"]])) | |||
| push_p = push_stream_result.result(timeout=60) | |||
| ai_video_file = write_ai_video_result.result(timeout=60) | |||
| # 接收停止指令 | |||
+ 99
- 40
concurrency/PushVideoStreamProcess2.py
View File
| @@ -9,6 +9,7 @@ from time import time, sleep | |||
| from traceback import format_exc | |||
| import cv2 | |||
| import numpy as np | |||
| import psutil | |||
| from loguru import logger | |||
| @@ -20,7 +21,7 @@ from util.Cv2Utils import video_conjuncing, write_or_video, write_ai_video, push | |||
| from util.ImageUtils import url2Array, add_water_pic | |||
| from util.LogUtils import init_log | |||
| from util.PlotsUtils import draw_painting_joint, xywh2xyxy2 | |||
| from util.PlotsUtils import draw_painting_joint, filterBox, xywh2xyxy2 | |||
| from util.QueUtil import get_no_block_queue, put_queue, clear_queue | |||
| @@ -60,12 +61,10 @@ class OnPushStreamProcess2(PushStreamProcess2): | |||
| frame_score = context["service"]["filter"]["frame_score"] | |||
| ex = None | |||
| ex_status = True | |||
| high_score_image = {} | |||
| # 相似度, 默认值0.65 | |||
| similarity = context["service"]["filter"]["similarity"] | |||
| # 图片相似度开关 | |||
| picture_similarity = bool(context["service"]["filter"]["picture_similarity"]) | |||
| frame_step = int(context["service"]["filter"]["frame_step"]) | |||
| qs_np_tmp = None | |||
| pix_dis = 60 | |||
| try: | |||
| init_log(base_dir, env) | |||
| logger.info("开始启动推流进程！requestId:{}", request_id) | |||
| @@ -101,6 +100,10 @@ class OnPushStreamProcess2(PushStreamProcess2): | |||
| copy_frame = frame.copy() | |||
| # 所有问题记录字典 | |||
| det_xywh, thread_p = {}, [] | |||
| det_xywh2 = {} | |||
| # 所有问题的矩阵集合 | |||
| qs_np = None | |||
| qs_reurn = [] | |||
| # [模型1识别数组, 模型2识别数组, 模型3识别数组] | |||
| for s_det_list in push_objs: | |||
| code, det_result = s_det_list[0], s_det_list[1][i] | |||
| @@ -115,12 +118,19 @@ class OnPushStreamProcess2(PushStreamProcess2): | |||
| 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] = {} | |||
| det_xywh[code] = {} | |||
| cd = det_xywh[code].get(cls) | |||
| if cd is None: | |||
| det_xywh[code][cls] = [[cls, box, score, label_array, color]] | |||
| 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: | |||
| det_xywh[code][cls].append([cls, box, score, label_array, color]) | |||
| 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 logo: | |||
| frame = add_water_pic(frame, logo, request_id) | |||
| copy_frame = add_water_pic(copy_frame, logo, request_id) | |||
| @@ -137,23 +147,42 @@ class OnPushStreamProcess2(PushStreamProcess2): | |||
| push_p_result = t.submit(push_video_stream, frame_merge, push_p, push_url, | |||
| p_push_status, | |||
| request_id) | |||
| if len(det_xywh) > 0: | |||
| flag = True | |||
| if picture_similarity and len(high_score_image) > 0: | |||
| hash1 = ImageUtils.dHash(high_score_image.get("or_frame")) | |||
| hash2 = ImageUtils.dHash(frame) | |||
| dist = ImageUtils.Hamming_distance(hash1, hash2) | |||
| similarity_1 = 1 - dist * 1.0 / 64 | |||
| if similarity_1 >= similarity: | |||
| flag = False | |||
| if len(high_score_image) > 0: | |||
| diff_frame_num = frame_index_list[i] - high_score_image.get("current_frame") | |||
| if diff_frame_num < frame_step: | |||
| flag = False | |||
| if flag: | |||
| high_score_image["or_frame"] = frame | |||
| high_score_image["current_frame"] = frame_index_list[i] | |||
| put_queue(image_queue, (1, [det_xywh, frame, frame_index_list[i], all_frames, draw_config["font_config"]])) | |||
| if qs_np is not None: | |||
| if len(qs_np.shape) == 1: | |||
| qs_np = qs_np[np.newaxis,...] | |||
| qs_np_id = qs_np.copy() | |||
| b = np.ones(qs_np_id.shape[0]) | |||
| qs_np_id = np.column_stack((qs_np_id,b)) | |||
| if qs_np_tmp is None: | |||
| 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)) | |||
| else: | |||
| qs_reurn = filterBox(qs_np, qs_np_tmp, pix_dis) | |||
| if picture_similarity: | |||
| qs_np_tmp = np.append(qs_np_tmp,qs_np_id,axis=0) | |||
| qs_np_tmp[:, 11] += 1 | |||
| qs_np_tmp = np.delete(qs_np_tmp, np.where((qs_np_tmp[:, 11] >= 75))[0], axis=0) | |||
| for q in qs_reurn: | |||
| if q[11] == 0: | |||
| cls = int(q[9]) | |||
| code = str(int(q[10])).zfill(3) | |||
| if det_xywh2.get(code) is None: | |||
| det_xywh2[code] = {} | |||
| cd = det_xywh2[code].get(cls) | |||
| score = q[8] | |||
| rainbows, label_arrays = draw_config[code]["rainbows"], draw_config[code]["label_arrays"] | |||
| 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 cd is None: | |||
| det_xywh2[code][cls] = [[cls, box, score, label_array, color]] | |||
| else: | |||
| det_xywh2[code][cls].append([cls, box, score, label_array, color]) | |||
| if len(det_xywh2) > 0: | |||
| put_queue(image_queue, (1, [det_xywh2, frame, frame_index_list[i], all_frames, draw_config["font_config"]])) | |||
| push_p = push_p_result.result(timeout=60) | |||
| ai_video_file = write_ai_video_result.result(timeout=60) | |||
| or_video_file = write_or_video_result.result(timeout=60) | |||
| @@ -216,6 +245,8 @@ class OffPushStreamProcess2(PushStreamProcess2): | |||
| # 图片相似度开关 | |||
| picture_similarity = bool(context["service"]["filter"]["picture_similarity"]) | |||
| frame_step = int(context["service"]["filter"]["frame_step"]) | |||
| qs_np_tmp = None | |||
| pix_dis = 60 | |||
| try: | |||
| init_log(base_dir, env) | |||
| logger.info("开始启动离线推流进程！requestId:{}", request_id) | |||
| @@ -253,6 +284,9 @@ class OffPushStreamProcess2(PushStreamProcess2): | |||
| copy_frame = frame.copy() | |||
| # 所有问题记录字典 | |||
| det_xywh, thread_p = {}, [] | |||
| det_xywh2 = {} | |||
| qs_np = None | |||
| qs_reurn = [] | |||
| for s_det_list in push_objs: | |||
| code, det_result = s_det_list[0], s_det_list[1][i] | |||
| if len(det_result) > 0: | |||
| @@ -270,8 +304,15 @@ class OffPushStreamProcess2(PushStreamProcess2): | |||
| cd = det_xywh[code].get(cls) | |||
| if cd is None: | |||
| det_xywh[code][cls] = [[cls, box, score, label_array, color]] | |||
| 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: | |||
| det_xywh[code][cls].append([cls, box, score, label_array, color]) | |||
| 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 logo: | |||
| frame = add_water_pic(frame, logo, request_id) | |||
| copy_frame = add_water_pic(copy_frame, logo, request_id) | |||
| @@ -286,23 +327,41 @@ class OffPushStreamProcess2(PushStreamProcess2): | |||
| push_p_result = t.submit(push_video_stream, frame_merge, push_p, push_url, | |||
| p_push_status, | |||
| request_id) | |||
| if len(det_xywh) > 0: | |||
| flag = True | |||
| if picture_similarity and len(high_score_image) > 0: | |||
| hash1 = ImageUtils.dHash(high_score_image.get("or_frame")) | |||
| hash2 = ImageUtils.dHash(frame) | |||
| dist = ImageUtils.Hamming_distance(hash1, hash2) | |||
| similarity_1 = 1 - dist * 1.0 / 64 | |||
| if similarity_1 >= similarity: | |||
| flag = False | |||
| if len(high_score_image) > 0: | |||
| diff_frame_num = frame_index_list[i] - high_score_image.get("current_frame") | |||
| if diff_frame_num < frame_step: | |||
| flag = False | |||
| if flag: | |||
| high_score_image["or_frame"] = frame | |||
| high_score_image["current_frame"] = frame_index_list[i] | |||
| put_queue(image_queue, (1, [det_xywh, frame, frame_index_list[i], all_frames, draw_config["font_config"]])) | |||
| if qs_np is not None: | |||
| if len(qs_np.shape) == 1: | |||
| qs_np = qs_np[np.newaxis,...] | |||
| qs_np_id = qs_np.copy() | |||
| b = np.ones(qs_np_id.shape[0]) | |||
| qs_np_id = np.column_stack((qs_np_id,b)) | |||
| if qs_np_tmp is None: | |||
| 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)) | |||
| else: | |||
| qs_reurn = filterBox(qs_np, qs_np_tmp, pix_dis) | |||
| if picture_similarity: | |||
| qs_np_tmp = np.append(qs_np_tmp,qs_np_id,axis=0) | |||
| qs_np_tmp[:, 11] += 1 | |||
| qs_np_tmp = np.delete(qs_np_tmp, np.where((qs_np_tmp[:, 11] >= 75))[0], axis=0) | |||
| for q in qs_reurn: | |||
| if q[11] == 0: | |||
| cls = int(q[9]) | |||
| code = str(int(q[10])).zfill(3) | |||
| if det_xywh2.get(code) is None: | |||
| det_xywh2[code] = {} | |||
| cd = det_xywh2[code].get(cls) | |||
| score = q[8] | |||
| rainbows, label_arrays = draw_config[code]["rainbows"], draw_config[code]["label_arrays"] | |||
| 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 cd is None: | |||
| det_xywh2[code][cls] = [[cls, box, score, label_array, color]] | |||
| else: | |||
| det_xywh2[code][cls].append([cls, box, score, label_array, color]) | |||
| if len(det_xywh2) > 0: | |||
| put_queue(image_queue, (1, [det_xywh2, frame, frame_index_list[i], all_frames, draw_config["font_config"]])) | |||
| push_p = push_p_result.result(timeout=60) | |||
| ai_video_file = write_ai_video_result.result(timeout=60) | |||
| # 接收停止指令 | |||
+ 38
- 0
util/PlotsUtils.py
View File
| @@ -134,3 +134,41 @@ def draw_painting_joint(box, img, label_array, score=0.5, color=None, config=Non | |||
| ''' | |||
| cv2.putText(img, label, p3, 0, config[3], [225, 255, 255], thickness=config[4], lineType=cv2.LINE_AA) | |||
| return img, box | |||
| def filterBox(det0, det1, pix_dis): | |||
| # det0为 (m1, 11) 矩阵 | |||
| # det1为 (m2, 12) 矩阵 | |||
| if len(det0.shape) == 1: | |||
| det0 = det0[np.newaxis,...] | |||
| if len(det1.shape) == 1: | |||
| det1 = det1[np.newaxis,...] | |||
| det1 = det1[...,0:11].copy() | |||
| m, n = det0.size, det1.size | |||
| if not m: | |||
| return det0 | |||
| # 在det0的列方向加一个元素flag代表该目标框中心点是否在之前目标框内(0代表不在，其他代表在) | |||
| flag = np.zeros([len(det0), 1]) | |||
| det0 = np.concatenate([det0, flag], axis=1) | |||
| det0_copy = det0.copy() | |||
| # det1_copy = det1.copy() | |||
| if not n: | |||
| return det0 | |||
| # det0转成 (m1, m2, 12) 的矩阵 | |||
| # det1转成 (m1, m2, 12) 的矩阵 | |||
| # det0与det1在第3维方向上拼接(6 + 7 = 13) | |||
| det0 = det0[:, np.newaxis, :].repeat(det1.shape[0], 1) | |||
| det1 = det1[np.newaxis, ...].repeat(det0.shape[0], 0) | |||
| joint_det = np.concatenate((det1, det0), axis=2) | |||
| # 分别求det0和det1的x1, y1, x2, y2(水平框的左上右下角点) | |||
| x1, y1, x2, y2 = joint_det[..., 0], joint_det[..., 1], joint_det[..., 4], joint_det[..., 5] | |||
| x3, y3, x4, y4 = joint_det[..., 11], joint_det[..., 12], joint_det[..., 15], joint_det[..., 16] | |||
| x2_c, y2_c = (x1+x2)//2, (y1+y2)//2 | |||
| x_c, y_c = (x3+x4)//2, (y3+y4)//2 | |||
| dis = (x2_c - x_c)**2 + (y2_c - y_c)**2 | |||
| mask = (joint_det[..., 9] == joint_det[..., 20]) & (dis <= pix_dis**2) | |||
| # 类别相同 & 中心点在上一帧的框内 判断为True | |||
| res = np.sum(mask, axis=1) | |||
| det0_copy[..., -1] = res | |||
| return det0_copy | |||
Loading…