文章地址https://arxiv.org/abs/1910.13302,
代码https://github.com/ZFTurbo/Weighted-Boxes-Fusion
加权框融合作为目标检测类比赛中常用的tricks,在这里简单做一下代码上的解析,有错误的地方还请各位指正
源码
# coding: utf-8
__author__ = 'ZFTurbo: https://kaggle.com/zfturbo'
import warnings
import numpy as np
from numba import jit
@jit(nopython=True)
def bb_intersection_over_union(A, B) -> float:
xA = max(A[0], B[0])
yA = max(A[1], B[1])
xB = min(A[2], B[2])
yB = min(A[3], B[3])
# compute the area of intersection rectangle
interArea = max(0, xB - xA) * max(0, yB - yA)
if interArea == 0:
return 0.0
# compute the area of both the prediction and ground-truth rectangles
boxAArea = (A[2] - A[0]) * (A[3] - A[1])
boxBArea = (B[2] - B[0]) * (B[3] - B[1])
iou = interArea / float(boxAArea + boxBArea - interArea)
return iou
def prefilter_boxes(boxes, scores, labels, weights, thr):
# Create dict with boxes stored by its label
new_boxes = dict()
for t in range(len(boxes)):
if len(boxes[t]) != len(scores[t]):
print('Error. Length of boxes arrays not equal to length of scores array: {} != {}'.format(len(boxes[t]), len(scores[t])))
exit()
if len(boxes[t]) != len(labels[t]):
print('Error. Length of boxes arrays not equal to length of labels array: {} != {}'.format(len(boxes[t]), len(labels[t])))
exit()
for j in range(len(boxes[t])): # 遍历第一个模型生成的boxes
score = scores[t][j] # 对应的score
if score < thr: # 过滤score小于 thr的框
continue
label = int(labels[t][j]) # 确定种类标签
box_part = boxes[t][j]
x1 = float(box_part[0])
y1 = float(box_part[1])
x2 = float(box_part[2])
y2 = float(box_part[3])
# Box data checks
if x2 < x1:
warnings.warn('X2 < X1 value in box. Swap them.')
x1, x2 = x2, x1
if y2 < y1:
warnings.warn('Y2 < Y1 value in box. Swap them.')
y1, y2 = y2, y1
if x1 < 0:
warnings.warn('X1 < 0 in box. Set it to 0.')
x1 = 0
if x1 > 1:
warnings.warn('X1 > 1 in box. Set it to 1. Check that you normalize boxes in [0, 1] range.')
x1 = 1
if x2 < 0:
warnings.warn('X2 < 0 in box. Set it to 0.')
x2 = 0
if x2 > 1:
warnings.warn('X2 > 1 in box. Set it to 1. Check that you normalize boxes in [0, 1] range.')
x2 = 1
if y1 < 0:
warnings.warn('Y1 < 0 in box. Set it to 0.')
y1 = 0
if y1 > 1:
warnings.warn('Y1 > 1 in box. Set it to 1. Check that you normalize boxes in [0, 1] range.')
y1 = 1
if y2 < 0:
warnings.warn('Y2 < 0 in box. Set it to 0.')
y2 = 0
if y2 > 1:
warnings.warn('Y2 > 1 in box. Set it to 1. Check that you normalize boxes in [0, 1] range.')
y2 = 1
if (x2 - x1) * (y2 - y1) == 0.0:
warnings.warn("Zero area box skipped: {}.".format(box_part))
continue
# [label, score, weight, model index, x1, y1, x2, y2]
b = [int(label), float(score) * weights[t], weights[t], t, x1, y1, x2, y2] # t --> 第t个模型
if label not in new_boxes:
new_boxes[label] = []
new_boxes[label].append(b) # 建立字典new_boxes -->{label: [[label, score, weight, model index, x1, y1, x2, y2], ...]}
# Sort each list in dict by score and transform it to numpy array
for k in new_boxes:
current_boxes = np.array(new_boxes[k]) # current_boxes --> [label, score, weight, model index, x1, y1, x2, y2], ...]
new_boxes[k] = current_boxes[current_boxes[:, 1].argsort()[::-1]] # 根据score排序
return new_boxes
def get_weighted_box(boxes, conf_type='avg'): # boxes --> [label, score, weight, model index, x1, y1, x2, y2]
"""
Create weighted box for set of boxes
:param boxes: set of boxes to fuse
:param conf_type: type of confidence one of 'avg' or 'max'
:return: weighted box (label, score, weight, x1, y1, x2, y2)
"""
box = np.zeros(8, dtype=np.float32)
conf = 0
conf_list = []
w = 0
for b in boxes:
box[4:] += (b[1] * b[4:])
conf += b[1] # score
conf_list.append(b[1])
w += b[2] # weight
box[0] = boxes[0][0] # label
if conf_type == 'avg':
box[1] = conf / len(boxes) # score为平均值
elif conf_type == 'max':
box[1] = np.array(conf_list).max()
elif conf_type in ['box_and_model_avg', 'absent_model_aware_avg']:
box[1] = conf / len(boxes)
box[2] = w
box[3] = -1 # model index field is retained for consistensy but is not used.
box[4:] /= conf
return box
def find_matching_box(boxes_list, new_box, match_iou): # boxes_list第一次时为空,第二次时为[[first]], new_box --> [[label, score, weight, model index, x1, y1, x2, y2], ...]中遍历
best_iou = match_iou
best_index = -1
for i in range(len(boxes_list)):
box = boxes_list[i]
if box[0] != new_box[0]: # 判断label值是否相同
continue
iou = bb_intersection_over_union(box[4:], new_box[4:])
if iou > best_iou:
best_index = i
best_iou = iou
return best_index, best_iou
def weighted_boxes_fusion(boxes_list, scores_list, labels_list, weights=None, iou_thr=0.55, skip_box_thr=0.0, conf_type='avg', allows_overflow=False): # 0.6, 0.5
'''
:param boxes_list: list of boxes predictions from each model, each box is 4 numbers.
It has 3 dimensions (models_number, model_preds, 4)
Order of boxes: x1, y1, x2, y2. We expect float normalized coordinates [0; 1]
:param scores_list: list of scores for each model
:param labels_list: list of labels for each model
:param weights: list of weights for each model. Default: None, which means weight == 1 for each model
:param iou_thr: IoU value for boxes to be a match
:param skip_box_thr: exclude boxes with score lower than this variable
:param conf_type: how to calculate confidence in weighted boxes. 'avg': average value, 'max': maximum value, 'box_and_model_avg': box and model wise hybrid weighted average, 'absent_model_aware_avg': weighted average that takes into account the absent model.
:param allows_overflow: false if we want confidence score not exceed 1.0
:return: boxes: boxes coordinates (Order of boxes: x1, y1, x2, y2).
:return: scores: confidence scores
:return: labels: boxes labels
'''
if weights is None:
weights = np.ones(len(boxes_list))
if len(weights) != len(boxes_list):
print('Warning: incorrect number of weights {}. Must be: {}. Set weights equal to 1.'.format(len(weights), len(boxes_list)))
weights = np.ones(len(boxes_list))
weights = np.array(weights)
if conf_type not in ['avg', 'max', 'box_and_model_avg', 'absent_model_aware_avg']:
print('Unknown conf_type: {}. Must be "avg", "max" or "box_and_model_avg", or "absent_model_aware_avg"'.format(conf_type))
exit()
# filtered_boxes -->{label: array([[label, score, weight, model index, x1, y1, x2, y2], ...])}其中各个key中已根据score大小进行排序
filtered_boxes = prefilter_boxes(boxes_list, scores_list, labels_list, weights, skip_box_thr)
# print("filtered_boxes", filtered_boxes)
if len(filtered_boxes) == 0:
return np.zeros((0, 4)), np.zeros((0,)), np.zeros((0,))
overall_boxes = []
for label in filtered_boxes:
boxes = filtered_boxes[label] # boxes --> array([[label, score, weight, model index, x1, y1, x2, y2], ...])
new_boxes = []
weighted_boxes = []
# Clusterize boxes
for j in range(0, len(boxes)):
index, best_iou = find_matching_box(weighted_boxes, boxes[j], iou_thr) # 第一次循环返回值为-1和iou_thr
if index != -1:
new_boxes[index].append(boxes[j])
# print("new_boxes %d" % j, new_boxes)
weighted_boxes[index] = get_weighted_box(new_boxes[index], conf_type)
else:
new_boxes.append([boxes[j].copy()]) # new_boxes和weight_boxes填充值为boxes[0]
weighted_boxes.append(boxes[j].copy())
# Rescale confidence based on number of models and boxes
for i in range(len(new_boxes)):
clustered_boxes = np.array(new_boxes[i])
if conf_type == 'box_and_model_avg':
# weighted average for boxes
weighted_boxes[i][1] = weighted_boxes[i][1] * len(clustered_boxes) / weighted_boxes[i][2]
# identify unique model index by model index column
_, idx = np.unique(clustered_boxes[:, 3], return_index=True)
# rescale by unique model weights
weighted_boxes[i][1] = weighted_boxes[i][1] * clustered_boxes[idx, 2].sum() / weights.sum()
elif conf_type == 'absent_model_aware_avg':
# get unique model index in the cluster
models = np.unique(clustered_boxes[:, 3]).astype(int)
# create a mask to get unused model weights
mask = np.ones(len(weights), dtype=bool)
mask[models] = False
# absent model aware weighted average
weighted_boxes[i][1] = weighted_boxes[i][1] * len(clustered_boxes) / (weighted_boxes[i][2] + weights[mask].sum())
elif not allows_overflow:
weighted_boxes[i][1] = weighted_boxes[i][1] * min(weights.sum(), len(clustered_boxes)) / weights.sum() # 对score进行操作
else:
weighted_boxes[i][1] = weighted_boxes[i][1] * len(clustered_boxes) / weights.sum()
overall_boxes.append(np.array(weighted_boxes))
overall_boxes = np.concatenate(overall_boxes, axis=0)
overall_boxes = overall_boxes[overall_boxes[:, 1].argsort()[::-1]]
boxes = overall_boxes[:, 4:]
scores = overall_boxes[:, 1]
labels = overall_boxes[:, 0]
return boxes, scores, labels
官方给的例子
boxes_list = [[
[0.00, 0.51, 0.81, 0.91],
[0.10, 0.31, 0.71, 0.61],
[0.01, 0.32, 0.83, 0.93],
[0.02, 0.53, 0.11, 0.94],
[0.03, 0.24, 0.12, 0.35],
],[
[0.04, 0.56, 0.84, 0.92],
[0.12, 0.33, 0.72, 0.64],
[0.38, 0.66, 0.79, 0.95],
[0.08, 0.49, 0.21, 0.89],
]]
scores_list = [[0.9, 0.8, 0.2, 0.4, 0.7], [0.5, 0.8, 0.7, 0.3]]
# labels_list = [[0, 1, 0, 1, 1], [1, 1, 1, 0]]
# labels_list = [[8, 9, 8, 9, 9], [9, 9, 9, 8]]
labels_list = [[0, 0, 0, 0, 0], [0, 0, 0, 0]]
weights = [1, 1]
iou_thr = 0.7
skip_box_thr = 0.0001
sigma = 0.1
# boxes, scores, labels = nms(boxes_list, scores_list, labels_list, weights=weights, iou_thr=iou_thr)
# boxes, scores, labels = soft_nms(boxes_list, scores_list, labels_list, weights=weights, iou_thr=iou_thr, sigma=sigma, thresh=skip_box_thr)
# boxes, scores, labels = non_maximum_weighted(boxes_list, scores_list, labels_list, weights=weights, iou_thr=iou_thr, skip_box_thr=skip_box_thr)
boxes, scores, labels = weighted_boxes_fusion(boxes_list, scores_list, labels_list, weights=weights, iou_thr=iou_thr, skip_box_thr=skip_box_thr)
print(boxes)
print(scores)
print(labels)
源码流程解析
主函数weighted_boxes_fusion(boxes_list, scores_list, labels_list, weights=None, iou_thr=0.55, skip_box_thr=0.0, conf_type=‘avg’, allows_overflow=False),其中iou_thr对应iou阈值;skip_box_thr是对应每个框的得分,如果小于这个值就舍去这个框;weights是对应不同模型取不同的权重,如果是None,则默认每个模型权重一样都是1;
filtered_boxes = prefilter_boxes(boxes_list, scores_list, labels_list, weights, skip_box_thr)
通过该函数进行过滤,如果score<skip_box_thr则过滤掉该box
建立字典,关键字为label,并按score大小对box进行排序
输出为字典filtered_boxes -->{label: array([[label, score, weight, model index, x1, y1, x2, y2], …])}
(PS:为了便于后面理解,将[label, score, weight, model index, x1, y1, x2, y2]这样形式的数据记做element)
其中各个label下的数据已根据score大小进行排序
之后是一个循环,遍历filtered_boxes(即获取每一个label),
新建new_boxes=[ ]和weighted_boxes=[ ]
并在其中嵌套一个循环(即获取每一个label下的所有数据):
new_boxes中会建立多个集合(每个集合中的数据为,通过find_matching_box方法计算该label下的所有element互相计算iou,大于设定值的即放至同一集合中,即认为预测的是同一个物体)
weighted_boxes中的数据为,new_boxes中同一集合下(被认为预测同一个物体的element)做融合(通过get_weighted_box方法)而得到的新的element
接下来的一个循环操作,主要目的是针对element中的score,对score进行融合操作,即更新weighted_boxes中的score值
最后是将weighted_boxes 打包到overall_boxes中,并做排序