加权框融合

该仓库包含几种用于整合目标检测模型预测框的 Python 实现方法：

• 非极大值抑制 (NMS)
• 软性非极大值抑制 (Soft-NMS) [1]
• 非极大值加权法 (NMW) [2]
• 加权框融合 (WBF) [3] - 一种较其他方法表现更优的新方法

环境要求

Python 3.*、Numpy、Numba

安装

pip install ensemble-boxes

使用示例

框的坐标应为归一化形式，例如在 [0, 1] 范围内。顺序为：x1, y1, x2, y2。

以下是针对两个模型进行框融合的示例：

• 第一个模型预测 5 个框，第二个模型预测 4 个框。
• 模型 1 各框的置信度分数为：[0.9, 0.8, 0.2, 0.4, 0.7]
• 模型 2 各框的置信度分数为：[0.5, 0.8, 0.7, 0.3]
• 模型 1 各框的标签（类别）为：[0, 1, 0, 1, 1]
• 模型 2 各框的标签（类别）为：[1, 1, 1, 0]
• 我们设置第一个模型的权重为 2，第二个模型的权重为 1。
• 设置框之间的交并比阈值为 iou_thr = 0.5。
• 跳过置信度低于 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)

单模型

如果您需要对单个模型的预测结果应用 NMS 或其他方法，可以这样调用函数：

from ensemble_boxes import *
# 对单模型预测结果进行框合并
boxes, scores, labels = weighted_boxes_fusion([boxes_list], [scores_list], [labels_list], weights=None, method=method, iou_thr=iou_thr, thresh=thresh)

更多示例请参阅 example.py。

三维版本

WBF 方法支持三维框，可通过 weighted_boxes_fusion_3d 函数实现。使用示例请参见 example_3d.py。

一维版本

WBF 方法也支持一维线段，可通过 weighted_boxes_fusion_1d 函数实现。使用示例请参见 example_1d.py。据报道，一维变体在自然语言处理 (NLP) 中的命名实体识别 (NER) 类型任务中非常有用。相关讨论请参见 这里。

基准测试

• Open Images 数据集基准测试（5 个模型）benchmark_oid/README.md
• COCO 数据集基准测试（10 个模型）benchmark_coco/README.md
• NLP 数据集基准测试（10 个模型）benchmark_nlp/README.md - 一维 WBF 变体的示例

WBF 方法说明及引用

• https://arxiv.org/abs/1910.13302（更新于 2020 年 8 月）
• https://authors.elsevier.com/c/1ca0dxnVK3cWY

如果您觉得此代码有用，请引用以下文献：

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

Weighted-Boxes-Fusion 快速上手指南

Weighted-Boxes-Fusion (WBF) 是一个用于目标检测模型结果融合的 Python 库。它提供了多种融合算法（如 NMS、Soft-NMS、NMW），其中核心的 WBF 算法能通过加权平均显著提升多模型或单模型的检测精度。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统：Linux, macOS 或 Windows
• Python 版本：Python 3.*
• 核心依赖：
  • numpy
  • numba (用于加速计算)

安装步骤

推荐使用 pip 直接安装官方发布包。国内用户可指定清华源或阿里源以加速下载。

# 使用默认源安装
pip install ensemble-boxes

# 或使用国内镜像源加速安装（推荐）
pip install ensemble-boxes -i https://pypi.tuna.tsinghua.edu.cn/simple

基本使用

1. 数据格式说明

输入数据的坐标需归一化到 [0, 1] 区间，顺序为 [x1, y1, x2, y2]。 你需要准备以下列表（支持多模型融合）：

• boxes_list: 每个模型的预测框列表。
• scores_list: 每个模型对应框的置信度分数。
• labels_list: 每个模型对应框的类别标签。
• weights: 每个模型的权重（列表），权重越高，该模型的预测结果对最终结果影响越大。

2. 代码示例

以下示例演示如何融合两个模型的预测结果，并应用 WBF 算法：

from ensemble_boxes import *

# 定义两个模型的预测框 (归一化坐标: x1, y1, x2, y2)
boxes_list = [
    [ # 模型 1 的 5 个框
        [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],
    ],
    [ # 模型 2 的 4 个框
        [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]]

# 设置模型权重 (模型 1 权重为 2, 模型 2 权重为 1)
weights = [2, 1]

# 参数设置
iou_thr = 0.5          # IoU 阈值，超过此值视为同一目标
skip_box_thr = 0.0001  # 忽略低于此置信度的框

# 执行 Weighted Boxes Fusion (WBF)
boxes, scores, labels = weighted_boxes_fusion(
    boxes_list, 
    scores_list, 
    labels_list, 
    weights=weights, 
    iou_thr=iou_thr, 
    skip_box_thr=skip_box_thr
)

# 输出结果为融合后的 boxes, scores, labels
print(f"融合后框数量：{len(boxes)}")

3. 单模型使用

如果你只有一个模型的预测结果，想利用 WBF 进行去重和优化（替代传统 NMS），只需将输入包裹一层列表，并将 weights 设为 None：

from ensemble_boxes import *

# 假设 boxes_list, scores_list, labels_list 是单个模型的数据
# 注意：外层需要再加一层列表，例如 [boxes_list]
boxes, scores, labels = weighted_boxes_fusion(
    [boxes_list], 
    [scores_list], 
    [labels_list], 
    weights=None, 
    iou_thr=0.5, 
    skip_box_thr=0.0001
)

提示：该库还支持 3D 框 (weighted_boxes_fusion_3d) 和 1D 线段 (weighted_boxes_fusion_1d, 适用于 NLP 任务) 的融合，具体用法可参考库内的 example 文件。