Set of methods to ensemble boxes from different object detection models, including implementation of "Weighted boxes fusion (WBF)" method.
MIT License
Repository contains Python implementation of several methods for ensembling boxes from object detection models:
Python 3.*, Numpy, Numba
pip install ensemble-boxes
Coordinates for boxes expected to be normalized e.g in range [0; 1]. Order: x1, y1, x2, y2.
Example of boxes ensembling for 2 models below.
from ensemble_boxes import *
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]]
weights = [2, 1]
iou_thr = 0.5
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)
If you need to apply NMS or any other method to single model predictions you can call function like that:
from ensemble_boxes import *
# Merge boxes for single model predictions
boxes, scores, labels = weighted_boxes_fusion([boxes_list], [scores_list], [labels_list], weights=None, method=method, iou_thr=iou_thr, thresh=thresh)
More examples can be found in example.py
There is support for 3D boxes in WBF method with weighted_boxes_fusion_3d
function. Check example of usage in example_3d.py
There is support for 1D line segments in WBF method with weighted_boxes_fusion_1d
function. Check example of usage in example_1d.py. It was reported that 1D variant can be useful in Named-entity recognition (NER) type of tasks for Natural Language Processing (NLP) problems. Check discussion here.
If you find this code useful please cite:
@article{solovyev2021weighted,
title={Weighted boxes fusion: Ensembling boxes from different object detection models},
author={Solovyev, Roman and Wang, Weimin and Gabruseva, Tatiana},
journal={Image and Vision Computing},
pages={1-6},
year={2021},
publisher={Elsevier}
}