/Weighted-Boxes-Fusion

Set of methods to ensemble boxes from different object detection models, including implementation of "Weighted boxes fusion (WBF)" method.

Primary LanguagePythonMIT LicenseMIT

DOI

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

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}
}