访问官网
Github
Huggingface
论文双边参考用于高分辨率二分图像分割
4 阿尔托大学 5 上海人工智能实验室 6 特伦托大学
| DIS-样例_1 | DIS-样例_2 |
|---|---|
本仓库是"双边参考用于高分辨率二分图像分割"（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。您可以通过以下链接访问此服务：
我们的BiRefNet在许多类似的高分辨率任务上达到了最先进的水平：
DIS: 
本工作时DIS论文与代码的对比图：
COD:
本工作时COD论文与代码的对比图：
HRSOD: 
本工作时HRSOD论文与代码的对比图：
尝试我们的在线演示进行推理：
- 给定权重的推理和评估：
- 可调分辨率的带图形界面的在线推理：
- 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上更简单的单图像推理演示。
- 感谢fal.ai/birefnet:这个
-
更多视觉比较
- 感谢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)。
