@@ -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) | |||
# 接收停止指令 |
@@ -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) | |||
# 接收停止指令 |
@@ -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 |