python合并bbox(heatmap定位)

参考:

https://stackoverflow.com/questions/55593506/merge-the-bounding-boxes-near-by-into-one

#Generate two text boxes a larger one that covers them
def merge_boxes(box1, box2):
    return [min(box1[0], box2[0]),
         min(box1[1], box2[1]),
         max(box1[2], box2[2]),
         max(box1[3], box2[3])]

#Computer a Matrix similarity of distances of the text and object
def calc_sim(text, obj):
    # text: ymin, xmin, ymax, xmax
    # obj: ymin, xmin, ymax, xmax
    text_ymin, text_xmin, text_ymax, text_xmax = text
    obj_ymin, obj_xmin, obj_ymax, obj_xmax = obj
    x_dist = min(abs(text_xmin-obj_xmin), abs(text_xmin-obj_xmax), abs(text_xmax-obj_xmin), abs(text_xmax-obj_xmax))
    y_dist = min(abs(text_ymin-obj_ymin), abs(text_ymin-obj_ymax), abs(text_ymax-obj_ymin), abs(text_ymax-obj_ymax))
    dist = x_dist + y_dist
    return dist

#Principal algorithm for merge text
def merge_algo(texts_boxes):
    # Distance definition  between text to be merge
    dist_limit = 40
    for i, text_box_1 in enumerate(texts_boxes):
        for j, text_box_2 in enumerate(texts_boxes):
            if j <= i:
                continue
            # Create a new box if a distances is less than disctance limit defined
            if calc_sim(text_box_1, text_box_2) < dist_limit:
                # Create a new box
                new_box = merge_boxes(text_box_1, text_box_2)
                texts_boxes[i] = new_box
                #delete previous text boxes
                del texts_boxes[j]
                #return a new boxes and new text string that are close
                return True, texts_boxes
    return False, texts_boxes

通过定位heatmap高亮区域,大于某个阈值的点,取60*60的bbox搜索该bbox内满足大于阈值的个数
合并所有bbox,定位高亮区域

            bbox_list = []
            H,W = img_in_arr.shape[:2]
            index = np.argwhere(score_map_arr > 20)
            for m_key, m_val in enumerate(index):
                y = m_val[0]
                x = m_val[1]
                x1 = 0 if x-30 < 0 else x-30
                y1 = 0 if y-30 < 0 else y-30
                x2 = W if x+30 > W else x+30
                y2 = H if y+30 > H else y+30
                # img_crop = img_in_arr[y1:y2, x1:x2, :].copy()
                score_crop = score_map_arr[y1:y2, x1:x2].copy()

                if (np.sum(score_crop > 20) > 30):
                    bbox_list.append([y1,x1,y2,x2])

            need_to_merge = True
            while need_to_merge:
                need_to_merge, texts_boxes_copied = merge_algo(bbox_list)

            for val in texts_boxes_copied:
                x1, y1, x2, y2 = val
                print([x1,y1,x2,y2])
                cv2.rectangle(img_in_arr, (x1, y1), (x2, y2), (0, 255, 0), 4)
                cv2.imwrite("./result/result_{}.png".format(save_flag1), img_in_arr)
                save_flag1 += 1