ZFTurbo / Weighted Boxes Fusion
Programming Languages
Labels
Projects that are alternatives of or similar to Weighted Boxes Fusion
Weighted boxes fusion
Repository contains Python implementation of several methods for ensembling boxes from object detection models:
- Non-maximum Suppression (NMS)
- Soft-NMS [1]
- Non-maximum weighted (NMW) [2]
- Weighted boxes fusion (WBF) [3] - new method which gives better results comparing to others
Requirements
Python 3.*, Numpy, Numba
Installation
pip install ensemble-boxes
Usage examples
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.
- First model predicts 5 boxes, second model predicts 4 boxes.
- Confidence scores for each box model 1: [0.9, 0.8, 0.2, 0.4, 0.7]
- Confidence scores for each box model 2: [0.5, 0.8, 0.7, 0.3]
- Labels (classes) for each box model 1: [0, 1, 0, 1, 1]
- Labels (classes) for each box model 2: [1, 1, 1, 0]
- We set weight for 1st model to be 2, and weight for second model to be 1.
- We set intersection over union for boxes to be match: iou_thr = 0.5
- We skip boxes with confidence lower than skip_box_thr = 0.0001
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)
Single model
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
3D version
There is support for 3D boxes in WBF method with weighted_boxes_fusion_3d
function. Check example of usage in example_3d.py
Accuracy and speed comparison
Comparison was made for ensemble of 5 different object detection models predictions trained on Open Images Dataset (500 classes).
Model scores at local validation:
- Model 1: mAP(0.5) 0.5164
- Model 2: mAP(0.5) 0.5019
- Model 3: mAP(0.5) 0.5144
- Model 4: mAP(0.5) 0.5152
- Model 5: mAP(0.5) 0.4910
Method | mAP(0.5) Result | Best params | Elapsed time (sec) |
---|---|---|---|
NMS | 0.5642 | IOU Thr: 0.5 | 47 |
Soft-NMS | 0.5616 | Sigma: 0.1, Confidence Thr: 0.001 | 88 |
NMW | 0.5667 | IOU Thr: 0.5 | 171 |
WBF | 0.5982 | IOU Thr: 0.6 | 249 |
You can download model predictions as well as ground truth labels from here: test_data.zip
Ensemble script for them is available here: example_oid.py
We also published large benchmark based on COCO dataset here.
Description of WBF method and citation
- https://arxiv.org/abs/1910.13302 (updated: 2020.08)
- https://authors.elsevier.com/c/1ca0dxnVK3cWY
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}
}