# Copyright (c) OpenMMLab. All rights reserved. import warnings from typing import Tuple import numpy as np import torch from torch import Tensor # References: https://github.com/ZFTurbo/Weighted-Boxes-Fusion def weighted_boxes_fusion( bboxes_list: list, scores_list: list, labels_list: list, weights: list = None, iou_thr: float = 0.55, skip_box_thr: float = 0.0, conf_type: str = 'avg', allows_overflow: bool = False) -> Tuple[Tensor, Tensor, Tensor]: """weighted boxes fusion is a method for fusing predictions from different object detection models, which utilizes confidence scores of all proposed bounding boxes to construct averaged boxes. Args: bboxes_list(list): list of boxes predictions from each model, each box is 4 numbers. scores_list(list): list of scores for each model labels_list(list): list of labels for each model weights: list of weights for each model. Default: None, which means weight == 1 for each model iou_thr: IoU value for boxes to be a match skip_box_thr: exclude boxes with score lower than this variable. 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. allows_overflow: false if we want confidence score not exceed 1.0. Returns: bboxes(Tensor): boxes coordinates (Order of boxes: x1, y1, x2, y2). scores(Tensor): confidence scores labels(Tensor): boxes labels """ if weights is None: weights = np.ones(len(bboxes_list)) if len(weights) != len(bboxes_list): print('Warning: incorrect number of weights {}. Must be: ' '{}. Set weights equal to 1.'.format( len(weights), len(bboxes_list))) weights = np.ones(len(bboxes_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 = prefilter_boxes(bboxes_list, scores_list, labels_list, weights, skip_box_thr) if len(filtered_boxes) == 0: return torch.Tensor(), torch.Tensor(), torch.Tensor() overall_boxes = [] for label in filtered_boxes: boxes = filtered_boxes[label] new_boxes = [] weighted_boxes = np.empty((0, 8)) # Clusterize boxes for j in range(0, len(boxes)): index, best_iou = find_matching_box_fast(weighted_boxes, boxes[j], iou_thr) if index != -1: new_boxes[index].append(boxes[j]) weighted_boxes[index] = get_weighted_box( new_boxes[index], conf_type) else: new_boxes.append([boxes[j].copy()]) weighted_boxes = np.vstack((weighted_boxes, boxes[j].copy())) # Rescale confidence based on number of models and boxes for i in range(len(new_boxes)): clustered_boxes = new_boxes[i] if conf_type == 'box_and_model_avg': clustered_boxes = np.array(clustered_boxes) # 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': clustered_boxes = np.array(clustered_boxes) # 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 conf_type == 'max': weighted_boxes[i, 1] = weighted_boxes[i, 1] / weights.max() elif not allows_overflow: weighted_boxes[i, 1] = weighted_boxes[i, 1] * min( len(weights), len(clustered_boxes)) / weights.sum() else: weighted_boxes[i, 1] = weighted_boxes[i, 1] * len( clustered_boxes) / weights.sum() overall_boxes.append(weighted_boxes) overall_boxes = np.concatenate(overall_boxes, axis=0) overall_boxes = overall_boxes[overall_boxes[:, 1].argsort()[::-1]] bboxes = torch.Tensor(overall_boxes[:, 4:]) scores = torch.Tensor(overall_boxes[:, 1]) labels = torch.Tensor(overall_boxes[:, 0]).int() return bboxes, scores, labels def prefilter_boxes(boxes, scores, labels, weights, thr): 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])): score = scores[t][j] if 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 (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 ] if label not in new_boxes: new_boxes[label] = [] new_boxes[label].append(b) # 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]) new_boxes[k] = current_boxes[current_boxes[:, 1].argsort()[::-1]] return new_boxes def get_weighted_box(boxes, conf_type='avg'): 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] conf_list.append(b[1]) w += b[2] box[0] = boxes[0][0] if conf_type in ('avg', 'box_and_model_avg', 'absent_model_aware_avg'): box[1] = conf / len(boxes) elif conf_type == 'max': box[1] = np.array(conf_list).max() box[2] = w box[3] = -1 box[4:] /= conf return box def find_matching_box_fast(boxes_list, new_box, match_iou): def bb_iou_array(boxes, new_box): # bb intersection over union xA = np.maximum(boxes[:, 0], new_box[0]) yA = np.maximum(boxes[:, 1], new_box[1]) xB = np.minimum(boxes[:, 2], new_box[2]) yB = np.minimum(boxes[:, 3], new_box[3]) interArea = np.maximum(xB - xA, 0) * np.maximum(yB - yA, 0) # compute the area of both the prediction and ground-truth rectangles boxAArea = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1]) boxBArea = (new_box[2] - new_box[0]) * (new_box[3] - new_box[1]) iou = interArea / (boxAArea + boxBArea - interArea) return iou if boxes_list.shape[0] == 0: return -1, match_iou boxes = boxes_list ious = bb_iou_array(boxes[:, 4:], new_box[4:]) ious[boxes[:, 0] != new_box[0]] = -1 best_idx = np.argmax(ious) best_iou = ious[best_idx] if best_iou <= match_iou: best_iou = match_iou best_idx = -1 return best_idx, best_iou