BiRefNet

<h1 align="center">双边参考用于高分辨率二分图像分割</h1> <div align='center'> <a href='https://scholar.google.com/citations?user=TZRzWOsAAAAJ' target='_blank'>郑鹏</a> 1,4,5,6，  <a href='https://scholar.google.com/citations?user=0uPb8MMAAAAJ' target='_blank'>高德宏</a> 2，  <a href='https://scholar.google.com/citations?user=kakwJ5QAAAAJ' target='_blank'>范登平</a> 1*，  <a href='https://scholar.google.com/citations?user=9cMQrVsAAAAJ' target='_blank'>刘丽</a> 3，  <a href='https://scholar.google.com/citations?user=qQP6WXIAAAAJ' target='_blank'>Jorma Laaksonen</a> 4，  <a href='https://scholar.google.com/citations?user=pw_0Z_UAAAAJ' target='_blank'>欧阳万里</a> 5，  <a href='https://scholar.google.com/citations?user=stFCYOAAAAAJ' target='_blank'>Nicu Sebe</a> 6 </div> <div align='center'> 1 南开大学&ensp; 2 西北工业大学&ensp; 3 国防科技大学&ensp; 4 阿尔托大学&ensp; 5 上海人工智能实验室&ensp; 6 特伦托大学&ensp; </div> <div align="center" style="display: flex; justify-content: center; flex-wrap: wrap;"> <a href='https://arxiv.org/pdf/2401.03407'><img src='https://img.shields.io/badge/arXiv-论文-red'></a>&ensp; <a href='https://drive.google.com/file/d/1aBnJ_R9lbnC2dm8dqD0-pzP2Cu-U1Xpt/view'><img src='https://img.shields.io/badge/中文版-论文-red'></a>&ensp; <a href='https://www.birefnet.top'><img src='https://img.shields.io/badge/页面-项目-red'></a>&ensp; <a href='https://drive.google.com/drive/folders/1s2Xe0cjq-2ctnJBR24563yMSCOu4CcxM'><img src='https://img.shields.io/badge/谷歌云盘-相关资料-green'></a>&ensp; <a href='LICENSE'><img src='https://img.shields.io/badge/许可证-MIT-yellow'></a>&ensp; <a href='https://huggingface.co/spaces/ZhengPeng7/BiRefNet_demo'><img src='https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-空间-blue'></a>&ensp; <a href='https://huggingface.co/ZhengPeng7/BiRefNet'><img src='https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-模型-blue'></a>&ensp; </div> <div align="center" style="display: flex; justify-content: center; flex-wrap: wrap;"> <a href='https://colab.research.google.com/drive/14Dqg7oeBkFEtchaHLNpig2BcdkZEogba'><img src='https://img.shields.io/badge/单图推理-F9AB00?style=for-the-badge&logo=googlecolab&color=525252'></a>&ensp; <a href='https://colab.research.google.com/drive/1MaEiBfJ4xIaZZn0DqKrhydHB8X97hNXl'><img src='https://img.shields.io/badge/推理与评估-F9AB00?style=for-the-badge&logo=googlecolab&color=525252'></a>&ensp; <a href='https://colab.research.google.com/drive/1B6aKZ3ekcvKMkSBn0N5mCASLUYMp0whK'><img src='https://img.shields.io/badge/框引导分割-F9AB00?style=for-the-badge&logo=googlecolab&color=525252'></a>&ensp; </div>

DIS-样例_1	DIS-样例_2
<img src="https://drive.google.com/thumbnail?id=1ItXaA26iYnE8XQ_GgNLy71MOWePoS2-g&sz=w400" />	<img src="https://drive.google.com/thumbnail?id=1Z-esCujQF_uEa_YJjkibc3NUrW4aR_d4&sz=w400" />

本仓库是"双边参考用于高分辨率二分图像分割"（CAAI AIR 2024）的官方实现。

新闻 :newspaper:

2024年7月30日: 感谢@not-lain在官方huggingface.js 仓库中添加BiRefNet的友善努力。
2024年7月28日: 我们发布了框引导分割的Colab演示。
2024年7月15日: 我们在Hugging Face模型上部署了BiRefNet，用户可以轻松地用一行代码加载它。
2024年6月21日: 我们发布并上传了原始论文的中文版到我的谷歌云盘。
2024年5月28日: 我们建立了一个模型库，其中包含不同大小和不同任务的BiRefNet经过良好训练的权重，包括通用用途、人像分割、DIS、HRSOD、COD等。
2024年5月7日: 我们还发布了单图推理的Colab演示。非常感谢@rishabh063对此的支持。
2024年4月9日: 感谢Features and Labels Inc.部署了一个很酷的在线BiRefNet推理API，并为我提供了强大的GPU资源用于进一步实验！
2024年3月7日: 我们发布了BiRefNet代码、原始论文中所有任务的经过良好训练的权重，以及所有相关资料在我的谷歌云盘文件夹中。同时，我们还在Hugging Face Spaces上部署了我们的BiRefNet，以便更容易在线使用，并发布了推理和评估的Colab演示。
2024年1月7日: 我们在arXiv上发布了我们的论文。

:rocket: 通过HuggingFace用一行代码加载BiRefNet，查看更多：

from transformers import AutoModelForImageSegmentation
birefnet = AutoModelForImageSegmentation.from_pretrained('zhengpeng7/BiRefNet', trust_remote_code=True)

:flight_arrival: 推理伙伴：

我们很高兴能与FAL合作部署BiRefNet的推理API。您可以通过以下链接访问此服务：

https://fal.ai/models/fal-ai/birefnet

我们的BiRefNet在许多类似的高分辨率任务上达到了最先进的水平： DIS:

<details><summary>本工作时DIS论文与代码的对比图：</summary> <img src="https://drive.google.com/thumbnail?id=1DLt6CFXdT1QSWDj_6jRkyZINXZ4vmyRp&sz=w1620" /> <img src="https://drive.google.com/thumbnail?id=1gn5GyKFlJbMIkre1JyEdHDSYcrFmcLD0&sz=w1620" /> <img src="https://drive.google.com/thumbnail?id=16CVYYOtafEeZhHqv0am2Daku1n_exMP6&sz=w1620" /> <img src="https://drive.google.com/thumbnail?id=10K45xwPXmaTG4Ex-29ss9payA9yBnyLn&sz=w1620" /> <img src="https://drive.google.com/thumbnail?id=16EuyqKFJOqwMmagvfnbC9hUurL9pYLLB&sz=w1620" /> </details>

COD:

<details><summary>本工作时COD论文与代码的对比图：</summary> <img src="https://drive.google.com/thumbnail?id=1DLt6CFXdT1QSWDj_6jRkyZINXZ4vmyRp&sz=w1620" /> <img src="https://drive.google.com/thumbnail?id=1gn5GyKFlJbMIkre1JyEdHDSYcrFmcLD0&sz=w1620" /> <img src="https://drive.google.com/thumbnail?id=16CVYYOtafEeZhHqv0am2Daku1n_exMP6&sz=w1620" /> </details>

HRSOD:

<details><summary>本工作时HRSOD论文与代码的对比图：</summary> <img src="https://drive.google.com/thumbnail?id=1hNfQtlTAHT4-AVbk_47852zyRp1NOFLs&sz=w1620" /> <img src="https://drive.google.com/thumbnail?id=1bcVldUAxYkMI3OMTyaP_jNuOugDfYj-d&sz=w1620" /> <img src="https://drive.google.com/thumbnail?id=1p1zgyVz27cGEqQMtOKzm_6zoYK3Sw_Zk&sz=w1620" /> <img src="https://drive.google.com/thumbnail?id=1TubAvcoEbH_mHu3I-AxflnB71nkf35jJ&sz=w1620" /> <img src="https://drive.google.com/thumbnail?id=1A3V9HjVtcMQdnGPwuy-DBVhwKuo0q2lT&sz=w1620" /> </details>

尝试我们的在线演示进行推理：

给定权重的推理和评估：
可调分辨率的带图形界面的在线推理：
Colab上的单图像在线推理：

模型库

为了更广泛地使用我们的BiRefNet,我成功地将原始的学术版本扩展为更通用的版本,以便在实际生活中更好地应用。

建议从官方页面下载数据集。但你也可以下载打包好的数据集:DIS,HRSOD,COD,Backbones。

在[stuff]的exp-TASK_SETTINGS文件夹中可以找到所有模型的性能表现(几乎所有指标)。

原论文中的模型,用于在基准测试中进行比较:

[表格内容略]

使用自定义数据训练的模型(通用、人像),用于实际应用中的一般用途:

[表格内容略]

带框引导的分割:

给定框引导:在Colab中打开

模型效率:

截图来自原论文。所有测试均在单个A100 GPU上进行。

[图片]

第三方创作

考虑到计算能力较弱的边缘设备,我们提供了一个以swin_v1_tiny为骨干网络的轻量级版本,速度提高4倍以上,大小减小5倍以上。详细信息可以在这个issue及其链接中找到。

我们发现已经有一些基于BiRefNet的第三方应用。非常感谢他们对社区的贡献! 选择你喜欢的,点击即可尝试,无需编码:

应用:
- 感谢fal.ai/birefnet:这个fal.ai上的项目在线封装了BiRefNet,提供了更多有用的UI和API选项来调用模型。
- 感谢ZHO-ZHO-ZHO/ComfyUI-BiRefNet-ZHO:这个项目进一步改进了ComfyUI中BiRefNet的UI,特别是对视频数据。
- 感谢viperyl/ComfyUI-BiRefNet:这个项目将BiRefNet打包成ComfyUI节点,使这个最先进的模型更容易被每个人使用。
- 感谢Rishabh提供了在colab上更简单的单图像推理演示。
更多视觉比较
- 感谢twitter.com/ZHOZHO672070对图像中更多背景去除方法的比较。
- 感谢twitter.com/toyxyz3对视频中更多背景去除方法的比较。

使用方法

环境设置

# 使用PyTorch==2.0.1进行更快的训练编译。
conda create -n birefnet python=3.9 -y && conda activate birefnet
pip install -r requirements.txt

数据集准备

从以下链接下载我整理好的组合训练/测试集:DIS--COD--HRSOD，或者从single_ones文件夹或其官方页面下载单个官方数据集。您也可以在我的百度网盘上找到相同的数据集:DIS--COD--HRSOD。

权重准备

从我的谷歌云盘文件夹或其官方页面下载骨干网络权重。

运行

# 训练、测试和评估
./train_test.sh 运行名称 训练用GPU编号 测试用GPU编号
# 示例: ./train_test.sh tmp-proj 0,1,2,3,4,5,6,7 0

# 查看train.sh / test.sh以了解仅训练/测试评估的情况。
# 评估后，运行`gen_best_ep.py`从特定指标(您可以从Sm、wFm、HCE(仅DIS)中选择)中选择最佳检查点。

训练好的权重:

从[stuff]下载BiRefNet-{TASK}-{EPOCH}.pth。相应权重的(预测图/性能/训练日志)信息也可以在同一目录下的exp-BiRefNet-{TASK_SETTINGS}等文件夹中找到。

您也可以从本仓库的发布版本中下载权重。

结果可能与原始论文中的略有不同，您可以在每个exp-xx中的eval_results-BiRefNet-{TASK_SETTINGS}文件夹中查看它们，我们将在接下来的几天内更新它们。由于我使用的成本非常高(A100-80G x 8)，许多人(包括我自己....)无法负担，我只在单个A100-40G上重新训练了BiRefNet，并达到了同等(甚至更好)的性能水平。这意味着您可以直接在单个具有36.5G+内存的GPU上训练模型。顺便说一下，1024x1024的推理需要5.5G的GPU内存。(我个人花了很多钱租用A100-40G来重新训练三个任务的BiRefNet...T_T。希望这能帮到您。)

但如果您有更多更强大的GPU，可以在config.py中设置GPU ID并增加批量大小以加速训练。我们已经在脚本中使所有这类事情都具有适应性，以便在单卡训练和多卡训练之间无缝切换。尽情享受吧:)

我的一些留言:

这个项目最初只为DIS构建。但经过一次又一次的更新，我使它变得越来越大，嵌入了许多功能。最终，您可以将其用于任何二值图像分割任务，如DIS/COD/SOD、医学图像分割、异常分割等。您可以轻松开启/关闭以下功能(通常在config.py中):

多GPU训练:通过一个变量开启/关闭。
骨干网络选择:Swin_v1、PVT_v2、ConvNets、...
加权损失:BCE、IoU、SSIM、MAE、Reg、...
二值分割的对抗损失(在我之前的工作MCCL中提出)。
训练技巧:多尺度监督、冻结骨干网络、多尺度输入...
数据整理器:将所有数据加载到内存中，平滑组合不同数据集以进行组合训练和测试。
... 我真心希望您喜欢这个项目，并在更多工作中使用它来实现新的最先进水平。

定量结果

定性结果

引用

@article{zheng2024birefnet,
  title={Bilateral Reference for High-Resolution Dichotomous Image Segmentation},
  author={Zheng, Peng and Gao, Dehong and Fan, Deng-Ping and Liu, Li and Laaksonen, Jorma and Ouyang, Wanli and Sebe, Nicu},
  journal={CAAI Artificial Intelligence Research},
  year={2024}
}

联系方式

如有任何问题、讨论，甚至投诉，请随时在此处留下问题或发送电子邮件给我(zhengpeng0108@gmail.com)。如果您想进行更多公开交流，也可以加入Discord群组(https://discord.gg/d9NN5sgFrq)或QQ群(https://qm.qq.com/q/y6WPy7WOIK)。