Spaces:
Runtime error
Runtime error
| # 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 <https://arxiv.org/abs/1910.13302> 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 | |