nodes/chflame163-ComfyUI_LayerStyle
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

commit BiRefNetUltraV2 and LoadBiRefNetModel nodes, fix repair_dependency scropts

Browse Source
This commit is contained in:
chflame163
2024-09-05 14:28:05 +08:00
parent 0c2312209d
commit f38f395add
66 changed files with 7604 additions and 229 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.MD
+43
View File
@@ -91,6 +91,17 @@ This error is caused by the low version of ```protobuf``` package.
#### For the issues with the above three dependency packages, please double click ```repair_dependency.bat``` (for Official ComfyUI Protable) or ```repair_dependency_aki.bat``` (for ComfyUI-aki-v1.x) in the plugin folder to automatically fix them.
### ModuleNotFoundError: No module named 'supervision'
### ModuleNotFoundError: No module named 'inference.models'
### ImportError: cannot import name 'YOLOWorld' from 'inference.models'
One of the 3 errors encountered when running the Yolo-World node is due to the lack of dependency packages ```inference-cli[yolo-world]``` and ```inference-gpu[yolo-world]``` .
Solution:
* Please double click ```repair_dependency.bat``` (for Official ComfyUI Protable) or ```repair_dependency_aki.bat``` (for ComfyUI-aki-v1.x) in the plugin folder to automatically fix them.
* Or manual install the dependency package ```inference-cli[yolo-world]``` and ```inference-gpu[yolo-world]```.
### ValueError: Trimap did not contain foreground values (xxxx...)
This error is caused by the mask area being too large or too small when using the ```PyMatting``` method to handle the mask edges.
@@ -103,6 +114,9 @@ When this error has occurred, please check the network environment.
## Update
<font size="4">**If the dependency package error after updating, please double clicking ```repair_dependency.bat``` (for Official ComfyUI Protable) or ```repair_dependency_aki.bat``` (for ComfyUI-aki-v1.x) in the plugin folder to reinstall the dependency packages. </font><br />
* Commit [BiRefNetUltraV2](#BiRefNetUltraV2) and [LoadBiRefNetModel](#LoadBiRefNetModel) nodes, that support the use of the latest BiRefNet model.
Download model file from [BaiduNetdisk](https://pan.baidu.com/s/12z3qUuqag3nqpN2NJ5pSzg?pwd=ek65) or [GoogleDrive](https://drive.google.com/drive/folders/1s2Xe0cjq-2ctnJBR24563yMSCOu4CcxM) named ```BiRefNet-general-epoch_244.pth``` to ```ComfyUI/Models/BiRefNet/pth``` folder. You can also download more BiRefNet models and put them here.
* [ExtendCanvasV2](#ExtendCanvasV2) node support negative value input, it means image will be cropped.
* The default title color of nodes is changed to blue-green, and nodes in LayerStyle, LayerColor, LayerMask, LayerUtility, and LayerFilter are distinguished by different colors.
* The Object Detector nodes added sort bbox option, which allows sorting from left to right, top to bottom, and large to small, making object selection more intuitive and convenient. The nodes released yesterday has been abandoned, please manually replace it with the new version node (sorry).
@@ -1811,6 +1825,35 @@ Node options:
* device: Set whether the VitMatte to use cuda.
* max_megapixels: Set the maximum size for VitMate operations.
### <a id="table1">BiRefNetUltra</a>
This node supports the use of the latest BiRefNet model.
*Download model file from [BaiduNetdisk](https://pan.baidu.com/s/12z3qUuqag3nqpN2NJ5pSzg?pwd=ek65) or [GoogleDrive](https://drive.google.com/drive/folders/1s2Xe0cjq-2ctnJBR24563yMSCOu4CcxM) named ```BiRefNet-general-epoch_244.pth``` to ```ComfyUI/Models/BiRefNet/pth``` folder. You can also download more BiRefNet models and put them here.
![image](image/birefnet_ultra_v2_example.jpg)
Node Options:
![image](image/birefnet_ultra_v2_node.jpg)
* image: The input image.
* birefnet_model: The BiRefNet model is input and it is output from the LoadBiRefNetModel node.
* detail_method: Edge processing methods. provides VITMatte, VITMatte(local), PyMatting, GuidedFilter. If the model has been downloaded after the first use of VITMatte, you can use VITMatte (local) afterwards.
* detail_erode: Mask the erosion range inward from the edge. the larger the value, the larger the range of inward repair.
* detail_dilate: The edge of the mask expands outward. the larger the value, the wider the range of outward repair.
* black_point: Edge black sampling threshold.
* white_point: Edge white sampling threshold.
* process_detail: Due to the excellent edge processing of BiRefNet, it is set to False by default here.
* device: Set whether the VitMatte to use cuda.
* max_megapixels: Set the maximum size for VitMate operations.
### <a id="table1">LoadBiRefNetModel</a>
Load the BiRefNet model.
Node Options:
![image](image/load_birefnet_model_node.jpg)
* model: Select the model. List the files in the ```CoomfyUI/models/BiRefNet/pth``` folder for selection.
### <a id="table1">TransparentBackgroundUltra</a>
Using the transparent-background model to remove background has better recognition ability and speed, while also having ultra-high edge details.
README_CN.MD
+46 -2
View File
@@ -76,13 +76,24 @@ git clone https://github.com/chflame163/ComfyUI_LayerStyle.git
### NameError: name 'guidedFilter' is not defined
问题原因同上。
### Cannot import name 'VitMatteImageProcessor' from 'transformers'
### Cannot import name 'VitMatteImageProcessor' from 'transformers'
这个错误是由于```transformers``` 版本过低造成的
### insightface 加载缓慢
这是由于```protobuf``` 版本过低造成的。
#### 以上三个依赖包的问题，请双击运行插件目录下的```repair_dependency.bat```(官方便携包)或者```repair_dependency_aki.bat```(秋叶整合包)自动修复。
#### 以上3个依赖包的问题，请双击运行插件目录下的```repair_dependency.bat```(官方便携包)或者```repair_dependency_aki.bat```(秋叶整合包)自动修复。
### ModuleNotFoundError: No module named 'supervision'
### ModuleNotFoundError: No module named 'inference.models'
### ImportError: cannot import name 'YOLOWorld' from 'inference.models'
运行Yolo-World 节点时出现以上3个错误之一，是由于缺少 ```inference-cli[yolo-world]``` 和 ```inference-gpu[yolo-world]``` 依赖包造成的。
解决方法：
* 请双击运行插件目录下的```repair_dependency.bat```(官方便携包)或者```repair_dependency_aki.bat```(秋叶整合包)自动修复。
* 或者手动安装 ```inference-cli[yolo-world]``` 和 ```inference-gpu[yolo-world]``` 依赖包。
### ValueError: Trimap did not contain foreground values (xxxx...)
这个错误是由于使用PyMatting方法处理遮罩边缘时，遮罩面积过大或者过小引起的。
@@ -105,6 +116,8 @@ git clone https://github.com/chflame163/ComfyUI_LayerStyle.git
## 更新说明
<font size="4">**如果本插件更新后出现依赖包错误，请双击运行插件目录下的```install_requirements.bat```(官方便携包)，或 ```install_requirements_aki.bat```(秋叶整合包) 重新安装依赖包。
* 添加 [BiRefNetUltraV2](#BiRefNetUltraV2) 和 [LoadBiRefNetModel](#LoadBiRefNetModel) 节点，支持使用最新的BiRefNet模型。
从[百度网盘](https://pan.baidu.com/s/12z3qUuqag3nqpN2NJ5pSzg?pwd=ek65) 或 [GoogleDrive](https://drive.google.com/drive/folders/1s2Xe0cjq-2ctnJBR24563yMSCOu4CcxM) 下载 ```BiRefNet-general-epoch_244.pth``` 到 ComfyUI/Models/BiRefNet/pth 文件夹。也可以下载更多的BiRefNet模型放到这里。
* [ExtendCanvasV2](#ExtendCanvasV2) 节点支持负值输入，负值将被裁剪。
* 节点默认标题颜色改为蓝绿色，LayerStyle, LayerColor, LayerMask, LayerUtility 和 LayerFilter 的节点分别用不同的颜色区分。
* 合并 [heshengtao](https://github.com/heshengtao) 提交的PR, 添加[TextImageV2](#TextImageV2)节点，修改图片文字节点的缩放，字体间隔跟随缩放，坐标不再以文字左上角，而是整行文字的中心点。感谢作者的贡献。
@@ -1787,6 +1800,37 @@ RemBgUltra的V2升级版，增加了VITMatte边缘处理方法。
* device: 设置是否使用cuda。
* max_megapixels: 设置vitmatte运算的最大尺寸。
### <a id="table1">BiRefNetUltra</a>
本节点支持使用最新的BiRefNet模型。
*从[百度网盘](https://pan.baidu.com/s/12z3qUuqag3nqpN2NJ5pSzg?pwd=ek65) 或 [GoogleDrive](https://drive.google.com/drive/folders/1s2Xe0cjq-2ctnJBR24563yMSCOu4CcxM) 下载 ```BiRefNet-general-epoch_244.pth``` 到 ```ComfyUI/Models/BiRefNet/pth``` 文件夹。也可以下载更多的BiRefNet模型放到这里。
![image](image/birefnet_ultra_v2_example.jpg)
节点选项说明:
![image](image/birefnet_ultra_v2_node.jpg)
* image: 图片输入。
* birefnet_model: BiRefNet模型输入，模型从LoadBiRefNetModel节点输出。
* detail_method: 边缘处理方法。提供了VITMatte, VITMatte(local), PyMatting, GuidedFilter。如果首次使用VITMatte后模型已经下载，之后可以使用VITMatte(local)。
* detail_erode: 遮罩边缘向内侵蚀范围。数值越大，向内修复的范围越大。
* detail_dilate: 遮罩边缘向外扩张范围。数值越大，向外修复的范围越大。
* black_point: 边缘黑色采样阈值。
* white_point: 边缘白色采样阈值。
* process_detail: 由于BiRefNet的边缘处理已经非常不错，此处默认设为False。
* device: 设置是否使用cuda。
* max_megapixels: 设置vitmatte运算的最大尺寸。
### <a id="table1">LoadBiRefNetModel</a>
加载BiRefNet模型。
节点选项说明:
![image](image/load_birefnet_model_node.jpg)
* model: 选择模型。列出 ```CoomfyUI/models/BiRefNet/pth``` 文件夹下的文件供选择。
### <a id="table1">TransparentBackgroundUltra</a>
使用transparent-background模型去除背景，有更好的识别能力和识别速度，同时具有超高的边缘细节。
py/BiRefNet/.gitignore
+142
View File
@@ -0,0 +1,142 @@
# Custom
e_*
.vscode
ckpt
preds
evaluation/eval-*
nohup.out*
tmp*
*.pth
core-*-python-*
.DS_Store
__MACOSX/
# Byte-compiled / optimized / DLL files
__pycache__/
*.py[cod]
*$py.class
# C extensions
*.so
# Distribution / packaging
.Python
build/
develop-eggs/
dist/
downloads/
eggs/
.eggs/
lib/
lib64/
parts/
sdist/
var/
wheels/
pip-wheel-metadata/
share/python-wheels/
*.egg-info/
.installed.cfg
*.egg
MANIFEST
# PyInstaller
# Usually these files are written by a python script from a template
# before PyInstaller builds the exe, so as to inject date/other infos into it.
*.manifest
*.spec
# Installer logs
pip-log.txt
pip-delete-this-directory.txt
# Unit test / coverage reports
htmlcov/
.tox/
.nox/
.coverage
.coverage.*
.cache
nosetests.xml
coverage.xml
*.cover
*.py,cover
.hypothesis/
.pytest_cache/
# Translations
*.mo
*.pot
# Django stuff:
*.log
local_settings.py
db.sqlite3
db.sqlite3-journal
# Flask stuff:
instance/
.webassets-cache
# Scrapy stuff:
.scrapy
# Sphinx documentation
docs/_build/
# PyBuilder
target/
# Jupyter Notebook
.ipynb_checkpoints
# IPython
profile_default/
ipython_config.py
# pyenv
.python-version
# pipenv
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
# However, in case of collaboration, if having platform-specific dependencies or dependencies
# having no cross-platform support, pipenv may install dependencies that don't work, or not
# install all needed dependencies.
#Pipfile.lock
# PEP 582; used by e.g. github.com/David-OConnor/pyflow
__pypackages__/
# Celery stuff
celerybeat-schedule
celerybeat.pid
# SageMath parsed files
*.sage.py
# Environments
.env
.venv
env/
venv/
ENV/
env.bak/
venv.bak/
# Spyder project settings
.spyderproject
.spyproject
# Rope project settings
.ropeproject
# mkdocs documentation
/site
# mypy
.mypy_cache/
.dmypy.json
dmypy.json
# Pyre type checker
.pyre/
py/BiRefNet/LICENSE
+21
View File
@@ -0,0 +1,21 @@
MIT License
Copyright (c) 2024 ZhengPeng
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
py/BiRefNet/README.md
+316
View File
@@ -0,0 +1,316 @@
<h1 align="center">Bilateral Reference for High-Resolution Dichotomous Image Segmentation</h1>
<div align='center'>
<a href='https://scholar.google.com/citations?user=TZRzWOsAAAAJ' target='_blank'><strong>Peng Zheng</strong></a><sup> 1,4,5,6</sup>,&thinsp;
<a href='https://scholar.google.com/citations?user=0uPb8MMAAAAJ' target='_blank'><strong>Dehong Gao</strong></a><sup> 2</sup>,&thinsp;
<a href='https://scholar.google.com/citations?user=kakwJ5QAAAAJ' target='_blank'><strong>Deng-Ping Fan</strong></a><sup> 1*</sup>,&thinsp;
<a href='https://scholar.google.com/citations?user=9cMQrVsAAAAJ' target='_blank'><strong>Li Liu</strong></a><sup> 3</sup>,&thinsp;
<a href='https://scholar.google.com/citations?user=qQP6WXIAAAAJ' target='_blank'><strong>Jorma Laaksonen</strong></a><sup> 4</sup>,&thinsp;
<a href='https://scholar.google.com/citations?user=pw_0Z_UAAAAJ' target='_blank'><strong>Wanli Ouyang</strong></a><sup> 5</sup>,&thinsp;
<a href='https://scholar.google.com/citations?user=stFCYOAAAAAJ' target='_blank'><strong>Nicu Sebe</strong></a><sup> 6</sup>
</div>
<div align='center'>
<sup>1 </sup>Nankai University&ensp; <sup>2 </sup>Northwestern Polytechnical University&ensp; <sup>3 </sup>National University of Defense Technology&ensp;
<br />
<sup>4 </sup>Aalto University&ensp; <sup>5 </sup>Shanghai AI Laboratory&ensp; <sup>6 </sup>University of Trento&ensp;
</div>
<div align="center" style="display: flex; justify-content: center; flex-wrap: wrap;">
<a href='https://www.sciopen.com/article/pdf/10.26599/AIR.2024.9150038.pdf'><img src='https://img.shields.io/badge/Journal-Paper-red'></a>&ensp;
<a href='https://arxiv.org/pdf/2401.03407'><img src='https://img.shields.io/badge/arXiv-Paper-red'></a>&ensp;
<a href='https://drive.google.com/file/d/1FWvKDWTnK9RsiywfCsIxsnQzqv-dlO5u/view'><img src='https://img.shields.io/badge/中文版-Paper-red'></a>&ensp;
<a href='https://www.birefnet.top'><img src='https://img.shields.io/badge/Page-Project-red'></a>&ensp;
<a href='https://drive.google.com/drive/folders/1s2Xe0cjq-2ctnJBR24563yMSCOu4CcxM'><img src='https://img.shields.io/badge/GDrive-Stuff-green'></a>&ensp;
<a href='LICENSE'><img src='https://img.shields.io/badge/License-MIT-yellow'></a>&ensp;
<a href='https://huggingface.co/spaces/ZhengPeng7/BiRefNet_demo'><img src='https://img.shields.io/badge/%F0%9F%A4%97%20HF-Space-blue'></a>&ensp;
<a href='https://huggingface.co/ZhengPeng7/BiRefNet'><img src='https://img.shields.io/badge/%F0%9F%A4%97%20HF-Model-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/Multiple_Images_Inference-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/Inference_&_Evaluation-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/Box_Guided_Segmentation-F9AB00?style=for-the-badge&logo=googlecolab&color=525252'></a>&ensp;
</div>
| *DIS-Sample_1* | *DIS-Sample_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" /> |
This repo is the official implementation of "[**Bilateral Reference for High-Resolution Dichotomous Image Segmentation**](https://arxiv.org/pdf/2401.03407)" (___CAAI AIR 2024___).
> [!note]
> **We need more GPU resources** to push forward the performance of BiRefNet, especially on *matting* tasks, higher-resolution inference (*2K*), and more *efficient* model design. If you are happy to cooperate, please contact me at zhengpeng0108@gmail.com.
## News :newspaper:
* **`Aug 30, 2024`:** We uploaded notebooks in `tutorials` to run the inference and ONNX conversion locally.
* **`Aug 23, 2024`:** Our BiRefNet is now officially released [online](https://www.sciopen.com/article/10.26599/AIR.2024.9150038) on CAAI AIR journal. And thanks to the [press release](https://www.eurekalert.org/news-releases/1055380).
* **`Aug 19, 2024`:** We uploaded the ONNX model files of all weights in the [GitHub release](https://github.com/ZhengPeng7/BiRefNet/releases/tag/v1) and [GDrive folder](https://drive.google.com/drive/u/0/folders/1kZM55bwsRdS__bdnsXpkmH6QPyza-9-N). Check out the **ONNX conversion** part in [model zoo](https://github.com/ZhengPeng7/BiRefNet?tab=readme-ov-file#model-zoo) for more details.
* **`Jul 30, 2024`:** Thanks to @not-lain for his kind efforts in adding BiRefNet to the official huggingface.js [repo](https://github.com/huggingface/huggingface.js/blob/3a8651fbc6508920475564a692bf0e5b601d9343/packages/tasks/src/model-libraries-snippets.ts#L763).
* **`Jul 28, 2024`:** We released the [Colab demo for box-guided segmentation](https://colab.research.google.com/drive/1B6aKZ3ekcvKMkSBn0N5mCASLUYMp0whK).
* **`Jul 15, 2024`:** We deployed our BiRefNet on [Hugging Face Models](https://huggingface.co/ZhengPeng7/BiRefNet) for users to easily load it in one line code.
* **`Jun 21, 2024`:** We released and uploaded the Chinese version of our original paper to my [GDrive](https://drive.google.com/file/d/1aBnJ_R9lbnC2dm8dqD0-pzP2Cu-U1Xpt/view).
* **`May 28, 2024`:** We hold a [model zoo](https://github.com/ZhengPeng7/BiRefNet?tab=readme-ov-file#model-zoo) with well-trained weights of our BiRefNet in different sizes and for different tasks, including general use, matting segmentation, DIS, HRSOD, COD, etc.
* **`May 7, 2024`:** We also released the [Colab demo for multiple images inference](https://colab.research.google.com/drive/14Dqg7oeBkFEtchaHLNpig2BcdkZEogba). Many thanks to @rishabh063 for his support on it.
* **`Apr 9, 2024`:** Thanks to [Features and Labels Inc.](https://fal.ai/) for deploying a cool online BiRefNet [inference API](https://fal.ai/models/fal-ai/birefnet/playground) and providing me with strong GPU resources for 4 months on more extensive experiments!
* **`Mar 7, 2024`:** We released BiRefNet codes, the well-trained weights for all tasks in the original papers, and all related stuff in my [GDrive folder](https://drive.google.com/drive/folders/1s2Xe0cjq-2ctnJBR24563yMSCOu4CcxM). Meanwhile, we also deployed our BiRefNet on [Hugging Face Spaces](https://huggingface.co/spaces/ZhengPeng7/BiRefNet_demo) for easier online use and released the [Colab demo for inference and evaluation](https://colab.research.google.com/drive/1MaEiBfJ4xIaZZn0DqKrhydHB8X97hNXl).
* **`Jan 7, 2024`:** We released our paper on [arXiv](https://arxiv.org/pdf/2401.03407).
## :rocket: Load BiRefNet in _ONE LINE_ by HuggingFace, check more: [![BiRefNet](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Models-blue)](https://huggingface.co/ZhengPeng7/birefnet)
```python
from transformers import AutoModelForImageSegmentation
birefnet = AutoModelForImageSegmentation.from_pretrained('zhengpeng7/BiRefNet', trust_remote_code=True)
```
## :flight_arrival: Inference Partner:
We are really happy to collaborate with [FAL](https://fal.ai) to deploy the **inference API** of BiRefNet. You can access this service via the link below:
+ https://fal.ai/models/fal-ai/birefnet
Our BiRefNet has achieved SOTA on many similar HR tasks:
**DIS**: [![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/bilateral-reference-for-high-resolution/dichotomous-image-segmentation-on-dis-te1)](https://paperswithcode.com/sota/dichotomous-image-segmentation-on-dis-te1?p=bilateral-reference-for-high-resolution) [![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/bilateral-reference-for-high-resolution/dichotomous-image-segmentation-on-dis-te2)](https://paperswithcode.com/sota/dichotomous-image-segmentation-on-dis-te2?p=bilateral-reference-for-high-resolution) [![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/bilateral-reference-for-high-resolution/dichotomous-image-segmentation-on-dis-te3)](https://paperswithcode.com/sota/dichotomous-image-segmentation-on-dis-te3?p=bilateral-reference-for-high-resolution) [![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/bilateral-reference-for-high-resolution/dichotomous-image-segmentation-on-dis-te4)](https://paperswithcode.com/sota/dichotomous-image-segmentation-on-dis-te4?p=bilateral-reference-for-high-resolution) [![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/bilateral-reference-for-high-resolution/dichotomous-image-segmentation-on-dis-vd)](https://paperswithcode.com/sota/dichotomous-image-segmentation-on-dis-vd?p=bilateral-reference-for-high-resolution)
<details><summary>Figure of Comparison on DIS Papers with Codes (by the time of this work):</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>
<br />
**COD**:[![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/bilateral-reference-for-high-resolution/camouflaged-object-segmentation-on-cod)](https://paperswithcode.com/sota/camouflaged-object-segmentation-on-cod?p=bilateral-reference-for-high-resolution) [![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/bilateral-reference-for-high-resolution/camouflaged-object-segmentation-on-nc4k)](https://paperswithcode.com/sota/camouflaged-object-segmentation-on-nc4k?p=bilateral-reference-for-high-resolution) [![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/bilateral-reference-for-high-resolution/camouflaged-object-segmentation-on-camo)](https://paperswithcode.com/sota/camouflaged-object-segmentation-on-camo?p=bilateral-reference-for-high-resolution) [![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/bilateral-reference-for-high-resolution/camouflaged-object-segmentation-on-chameleon)](https://paperswithcode.com/sota/camouflaged-object-segmentation-on-chameleon?p=bilateral-reference-for-high-resolution)
<details><summary>Figure of Comparison on COD Papers with Codes (by the time of this work):</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>
<br />
**HRSOD**: [![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/bilateral-reference-for-high-resolution/rgb-salient-object-detection-on-davis-s)](https://paperswithcode.com/sota/rgb-salient-object-detection-on-davis-s?p=bilateral-reference-for-high-resolution) [![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/bilateral-reference-for-high-resolution/rgb-salient-object-detection-on-hrsod)](https://paperswithcode.com/sota/rgb-salient-object-detection-on-hrsod?p=bilateral-reference-for-high-resolution) [![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/bilateral-reference-for-high-resolution/rgb-salient-object-detection-on-uhrsd)](https://paperswithcode.com/sota/rgb-salient-object-detection-on-uhrsd?p=bilateral-reference-for-high-resolution) [![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/bilateral-reference-for-high-resolution/salient-object-detection-on-duts-te)](https://paperswithcode.com/sota/salient-object-detection-on-duts-te?p=bilateral-reference-for-high-resolution) [![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/bilateral-reference-for-high-resolution/salient-object-detection-on-dut-omron)](https://paperswithcode.com/sota/salient-object-detection-on-dut-omron?p=bilateral-reference-for-high-resolution)
<details><summary>Figure of Comparison on HRSOD Papers with Codes (by the time of this work):</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>
<br />
#### Try our online demos for inference:
+ **Inference and evaluation** of your given weights: [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1MaEiBfJ4xIaZZn0DqKrhydHB8X97hNXl)
+ **Online Inference with GUI** with adjustable resolutions: [![Hugging Face Spaces](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Spaces-blue)](https://huggingface.co/spaces/ZhengPeng7/BiRefNet_demo)
+ Online **Multiple Images Inference** on Colab: [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/14Dqg7oeBkFEtchaHLNpig2BcdkZEogba)
<img src="https://drive.google.com/thumbnail?id=12XmDhKtO1o2fEvBu4OE4ULVB2BK0ecWi&sz=w1620" />
## Model Zoo
> For more general use of our BiRefNet, I extended the original academic one to more general ones for better real-life application.
>
> Datasets and datasets are suggested to be downloaded from official pages. But you can also download the packaged ones: [DIS](https://drive.google.com/drive/folders/1hZW6tAGPJwo9mPS7qGGGdpxuvuXiyoMJ), [HRSOD](https://drive.google.com/drive/folders/18_hAE3QM4cwAzEAKXuSNtKjmgFXTQXZN), [COD](https://drive.google.com/drive/folders/1EyHmKWsXfaCR9O0BiZEc3roZbRcs4ECO), [Backbones](https://drive.google.com/drive/folders/1cmce_emsS8A5ha5XT2c_CZiJzlLM81ms).
>
> Find performances (almost all metrics) of all models in the `exp-TASK_SETTINGS` folders in [[**stuff**](https://drive.google.com/drive/folders/1s2Xe0cjq-2ctnJBR24563yMSCOu4CcxM)].
<details><summary>Models in the original paper, for <b>comparison on benchmarks</b>:</summary>
| Task | Training Sets | Backbone | Download |
| :---: | :-------------------------: | :-----------: | :----------------------------------------------------------: |
| DIS | DIS5K-TR | swin_v1_large | [google-drive](https://drive.google.com/file/d/1J90LucvDQaS3R_-9E7QUh1mgJ8eQvccb/view) |
| COD | COD10K-TR, CAMO-TR | swin_v1_large | [google-drive](https://drive.google.com/file/d/1tM5M72k7a8aKF-dYy-QXaqvfEhbFaWkC/view) |
| HRSOD | DUTS-TR | swin_v1_large | [google-drive](https://drive.google.com/file/d/1f7L0Pb1Y3RkOMbqLCW_zO31dik9AiUFa/view) |
| HRSOD | HRSOD-TR | swin_v1_large | google-drive |
| HRSOD | UHRSD-TR | swin_v1_large | google-drive |
| HRSOD | DUTS-TR, HRSOD-TR | swin_v1_large | [google-drive](https://drive.google.com/file/d/1WJooyTkhoDLllaqwbpur_9Hle0XTHEs_/view) |
| HRSOD | DUTS-TR, UHRSD-TR | swin_v1_large | [google-drive](https://drive.google.com/file/d/1Pu1mv3ORobJatIuUoEuZaWDl2ylP3Gw7/view) |
| HRSOD | HRSOD-TR, UHRSD-TR | swin_v1_large | [google-drive](https://drive.google.com/file/d/1xEh7fsgWGaS5c3IffMswasv0_u-aVM9E/view) |
| HRSOD | DUTS-TR, HRSOD-TR, UHRSD-TR | swin_v1_large | [google-drive](https://drive.google.com/file/d/13FaxyyOwyCddfZn2vZo1xG1KNZ3cZ-6B/view) |
</details>
<details><summary>Models trained with customed data (general, matting), for <b>general use in practical application</b>:</summary>
| Task | Training Sets | Backbone | Test Set | Metric (S, wF[, HCE]) | Download |
| :-----------------------: | :----------------------------------------------------------: | :-----------: | :-------: | :-------------------: | :----------------------------------------------------------: |
| **general use** | DIS5K-TR,DIS-TEs, DUTS-TR_TE,HRSOD-TR_TE,UHRSD-TR_TE, HRS10K-TR_TE, TR-P3M-10k, TE-P3M-500-NP, TE-P3M-500-P, TR-humans | swin_v1_large | DIS-VD | 0.911, 0.875, 1069 | [google-drive](https://drive.google.com/file/d/1_IfUnu8Fpfn-nerB89FzdNXQ7zk6FKxc/view) |
| **general use** | DIS5K-TR,DIS-TEs, DUTS-TR_TE,HRSOD-TR_TE,UHRSD-TR_TE, HRS10K-TR_TE, TR-P3M-10k, TE-P3M-500-NP, TE-P3M-500-P, TR-humans | swin_v1_tiny | DIS-VD | 0.882, 0.830, 1175 | [google-drive](https://drive.google.com/file/d/1fzInDWiE2n65tmjaHDSZpqhL0VME6-Yl/view) |
| **general use** | DIS5K-TR, DIS-TEs | swin_v1_large | DIS-VD | 0.907, 0.865, 1059 | [google-drive](https://drive.google.com/file/d/1P6NJzG3Jf1sl7js2q1CPC3yqvBn_O8UJ/view) |
| **matting segmentation** | [P3M-10k](https://github.com/JizhiziLi/P3M), [humans](https://huggingface.co/datasets/schirrmacher/humans) | swin_v1_large | P3M-500-P | 0.983, 0.989 | [google-drive](https://drive.google.com/file/d/1uUeXjEUoD2XF_6YjD_fsct-TJp7TFiqh) |
</details>
<details><summary>Segmentation with box <b>guidance</b>:</summary>
+ Given box guidance: [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1B6aKZ3ekcvKMkSBn0N5mCASLUYMp0whK)
</details>
<details><summary>Model <b>efficiency</b>:</summary>
> Screenshot from the original paper. All tests are conducted on a single A100 GPU.
<img src="https://drive.google.com/thumbnail?id=1mTfSD_qt-rFO1t8DRQcyIa5cgWLf1w2-&sz=h300" /> <img src="https://drive.google.com/thumbnail?id=1F_OURIWILVe4u1rSz-aqt6ur__bAef25&sz=h300" />
</details>
<details><summary><b>ONNX</b> conversion:</summary>
> We converted from `.pth` weights files to `.onnx` files.
> We referred a lot to the [Kazuhito00/BiRefNet-ONNX-Sample](https://github.com/Kazuhito00/BiRefNet-ONNX-Sample), many thanks to @Kazuhito00.
+ Check our [Colab demo for ONNX conversion](https://colab.research.google.com/drive/1z6OruR52LOvDDpnp516F-N4EyPGrp5om) or the [notebook file for local running](https://drive.google.com/file/d/1cgL2qyvOO5q3ySfhytypX46swdQwZLrJ), where you can do the conversion/inference by yourself and find all relevant info.
+ As tested, BiRefNets with SwinL (default backbone) cost `~90%` more time (the inference costs `~165ms` on an A100 GPU) using ONNX files. Meanwhile, BiRefNets with SwinT (lightweight) cost `~75%` more time (the inference costs `~93.8ms` on an A100 GPU) using ONNX files. Input resolution is `1024x1024` as default.
+ The results of the original pth files and the converted onnx files are slightly different, which is acceptable.
+ Pay attention to the compatibility among `onnxruntime-gpu, CUDA, and CUDNN` (we use `torch==2.0.1, cuda=11.8` here).
</details>
## Third-Party Creations
> Concerning edge devices with less computing power, we provide a lightweight version with `swin_v1_tiny` as the backbone, which is x4+ faster and x5+ smaller. The details can be found in [this issue](https://github.com/ZhengPeng7/BiRefNet/issues/11#issuecomment-2041033576) and links there.
We found there've been some 3rd party applications based on our BiRefNet. Many thanks for their contribution to the community!
Choose the one you like to try with clicks instead of codes:
1. **Applications**:
+ Thanks [**lbq779660843/BiRefNet-Tensorrt**](https://github.com/lbq779660843/BiRefNet-Tensorrt) and [**yuanyang1991/birefnet_tensorrt**](https://github.com/yuanyang1991/birefnet_tensorrt): they both provided the project to convert BiRefNet to **TensorRT**, which is faster and better for deployment. Their repos offer solid local establishment (Win and Linux) and [colab demo](https://colab.research.google.com/drive/1r8GkFPyMMO0OkMX6ih5FjZnUCQrl2SHV?usp=sharing), respectively. And @yuanyang1991 kindly offered the comparison among the inference efficiency of naive PyTorch, ONNX, and TensorRT on an RTX 4080S:
| Methods | [Pytorch](https://drive.google.com/file/d/1_IfUnu8Fpfn-nerB89FzdNXQ7zk6FKxc/view) | [ONNX](https://drive.google.com/drive/u/0/folders/1kZM55bwsRdS__bdnsXpkmH6QPyza-9-N) | TensorRT |
|:------------------------------------------------------------------------------------:|:--------------:|:--------------:|:--------------:|
| &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;First Inference Time&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp; | 0.71s | 5.32s | **0.17s** |
| Methods | [Pytorch](https://drive.google.com/file/d/1_IfUnu8Fpfn-nerB89FzdNXQ7zk6FKxc/view) | [ONNX](https://drive.google.com/drive/u/0/folders/1kZM55bwsRdS__bdnsXpkmH6QPyza-9-N) | TensorRT |
|:------------------------------------------------------------------------------------:|:--------------:|:--------------:|:--------------:|
| Avg Inf Time (excluding 1st) | 0.15s | 4.43s | **0.11s** |
+ Thanks [**dimitribarbot/sd-webui-birefnet**](https://github.com/dimitribarbot/sd-webui-birefnet): this project allows to add a BiRefNet section to the original **Stable Diffusion WebUI**'s Extras tab.
<p align="center"><img src="https://drive.google.com/thumbnail?id=159bLXI71FWh4ZsHTvc-wApSN9ytVRmua&sz=w1620" /></p>
+ Thanks [**fal.ai/birefnet**](https://fal.ai/models/birefnet): this project on `fal.ai` encapsulates BiRefNet **online** with more useful options in **UI** and **API** to call the model.
<p align="center"><img src="https://drive.google.com/thumbnail?id=1rNk81YV_Pzb2GykrzfGvX6T7KBXR0wrA&sz=w1620" /></p>
+ Thanks [**ZHO-ZHO-ZHO/ComfyUI-BiRefNet-ZHO**](https://github.com/ZHO-ZHO-ZHO/ComfyUI-BiRefNet-ZHO): this project further improves the **UI** for BiRefNet in ComfyUI, especially for **video data**.
<p align="center"><img src="https://drive.google.com/thumbnail?id=1GOqEreyS7ENzTPN0RqxEjaA76RpMlkYM&sz=w1620" /></p>
<https://github.com/ZhengPeng7/BiRefNet/assets/25921713/3a1c7ab2-9847-4dac-8935-43a2d3cd2671>
+ Thanks [**viperyl/ComfyUI-BiRefNet**](https://github.com/viperyl/ComfyUI-BiRefNet): this project packs BiRefNet as **ComfyUI nodes**, and makes this SOTA model easier use for everyone.
<p align="center"><img src="https://drive.google.com/thumbnail?id=1KfxCQUUa2y9T-aysEaeVVjCUt3Z0zSkL&sz=w1620" /></p>
+ Thanks [**Rishabh**](https://github.com/rishabh063) for offering a demo for the [easier multiple images inference on colab](https://colab.research.google.com/drive/14Dqg7oeBkFEtchaHLNpig2BcdkZEogba).
2. **More Visual Comparisons**
+ Thanks [**twitter.com/ZHOZHO672070**](https://twitter.com/ZHOZHO672070) for the comparison with more background-removal methods in images:
<img src="https://drive.google.com/thumbnail?id=1nvVIFt_Ezs-crPSQxUDqkUBz598fTe63&sz=w1620" />
+ Thanks [**twitter.com/toyxyz3**](https://twitter.com/toyxyz3) for the comparison with more background-removal methods in videos:
<https://github.com/ZhengPeng7/BiRefNet/assets/25921713/40136198-01cc-4106-81f9-81c985f02e31>
<https://github.com/ZhengPeng7/BiRefNet/assets/25921713/1a32860c-0893-49dd-b557-c2e35a83c160>
## Usage
#### Environment Setup
```shell
# PyTorch==2.0.1 is used for faster training with compilation.
conda create -n birefnet python=3.9 -y && conda activate birefnet
pip install -r requirements.txt
```
#### Dataset Preparation
Download combined training / test sets I have organized well from: [DIS](https://drive.google.com/drive/folders/1hZW6tAGPJwo9mPS7qGGGdpxuvuXiyoMJ)--[COD](https://drive.google.com/drive/folders/1EyHmKWsXfaCR9O0BiZEc3roZbRcs4ECO)--[HRSOD](https://drive.google.com/drive/folders/18_hAE3QM4cwAzEAKXuSNtKjmgFXTQXZN) or the single official ones in the `single_ones` folder, or their official pages. You can also find the same ones on my **BaiduDisk**: [DIS](https://pan.baidu.com/s/1O_pQIGAE4DKqL93xOxHpxw?pwd=PSWD)--[COD](https://pan.baidu.com/s/1RnxAzaHSTGBC1N6r_RfeqQ?pwd=PSWD)--[HRSOD](https://pan.baidu.com/s/1_Del53_0lBuG0DKJJAk4UA?pwd=PSWD).
#### Weights Preparation
Download backbone weights from [my google-drive folder](https://drive.google.com/drive/folders/1s2Xe0cjq-2ctnJBR24563yMSCOu4CcxM) or their official pages.
## Run
```shell
# Train & Test & Evaluation
./train_test.sh RUN_NAME GPU_NUMBERS_FOR_TRAINING GPU_NUMBERS_FOR_TEST
# Example: ./train_test.sh tmp-proj 0,1,2,3,4,5,6,7 0
# See train.sh / test.sh for only training / test-evaluation.
# After the evaluation, run `gen_best_ep.py` to select the best ckpt from a specific metric (you choose it from Sm, wFm, HCE (DIS only)).
```
#### Well-trained weights:
Download the `BiRefNet-{TASK}-{EPOCH}.pth` from [[**stuff**](https://drive.google.com/drive/folders/1s2Xe0cjq-2ctnJBR24563yMSCOu4CcxM)]. Info of the corresponding (predicted\_maps/performance/training\_log) weights can be also found in folders like `exp-BiRefNet-{TASK_SETTINGS}` in the same directory.
You can also download the weights from the release of this repo.
The results might be a bit different from those in the original paper, you can see them in the `eval_results-BiRefNet-{TASK_SETTINGS}` folder in each `exp-xx`, we will update them in the following days. Due to the very high cost I used (A100-80G x 8) which many people cannot afford to (including myself....), I re-trained BiRefNet on a single A100-40G only and achieve the performance on the same level (even better). It means you can directly train the model on a single GPU with 36.5G+ memory. BTW, 5.5G GPU memory is needed for inference in 1024x1024. (I personally paid a lot for renting an A100-40G to re-train BiRefNet on the three tasks... T_T. Hope it can help you.)
But if you have more and more powerful GPUs, you can set GPU IDs and increase the batch size in `config.py` to accelerate the training. We have made all this kind of things adaptive in scripts to seamlessly switch between single-card training and multi-card training. Enjoy it :)
#### Some of my messages:
This project was originally built for DIS only. But after the updates one by one, I made it larger and larger with many functions embedded together. Finally, you can **use it for any binary image segmentation tasks**, such as DIS/COD/SOD, medical image segmentation, anomaly segmentation, etc. You can eaily open/close below things (usually in `config.py`):
+ Multi-GPU training: open/close with one variable.
+ Backbone choices: Swin_v1, PVT_v2, ConvNets, ...
+ Weighted losses: BCE, IoU, SSIM, MAE, Reg, ...
+ Adversarial loss for binary segmentation (proposed in my previous work [MCCL](https://arxiv.org/pdf/2302.14485)).
+ Training tricks: multi-scale supervision, freezing backbone, multi-scale input...
+ Data collator: loading all in memory, smooth combination of different datasets for combined training and test.
+ ...
I really hope you enjoy this project and use it in more works to achieve new SOTAs.
### Quantitative Results
<p align="center"><img src="https://drive.google.com/thumbnail?id=1Ymkh8WN16XMTBOS8dmPTg5eAf-NIl2m5&sz=w1620" /></p>
<p align="center"><img src="https://drive.google.com/thumbnail?id=1W0mi0ZiYbqsaGuohNXU8Gh7Zj4M3neFg&sz=w1620" /></p>
### Qualitative Results
<p align="center"><img src="https://drive.google.com/thumbnail?id=1TYZF8pVZc2V0V6g3ik4iAr9iKvJ8BNrf&sz=w1620" /></p>
<p align="center"><img src="https://drive.google.com/thumbnail?id=1ZGHC32CAdT9cwRloPzOCKWCrVQZvUAlJ&sz=w1620" /></p>
### Citation
```
@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},
volume = {3},
pages = {9150038},
year={2024}
}
```
## Contact
Any questions, discussions, or even complaints, feel free to leave issues here or send me e-mails (zhengpeng0108@gmail.com). You can also join the Discord Group (https://discord.gg/d9NN5sgFrq) or QQ Group (https://qm.qq.com/q/y6WPy7WOIK) if you want to talk a lot publicly.
py/BiRefNet/config.py
+131 -61
View File
@@ -1,95 +1,132 @@
import os
import math
from folder_paths import models_dir
class Config():
def __init__(self) -> None:
# PATH settings
# Make up your file system as: SYS_HOME_DIR/codes/dis/BiRefNet, SYS_HOME_DIR/datasets/dis/xx, SYS_HOME_DIR/weights/xx
if os.name == 'nt':
self.sys_home_dir = os.environ['USERPROFILE'] # For windows system
else:
self.sys_home_dir = os.environ['HOME'] # For Linux system
# TASK settings
self.task = ['DIS5K', 'COD', 'HRSOD', 'General', 'Matting'][0]
self.training_set = {
'DIS5K': ['DIS-TR', 'DIS-TR+DIS-TE1+DIS-TE2+DIS-TE3+DIS-TE4'][0],
'COD': 'TR-COD10K+TR-CAMO',
'HRSOD': ['TR-DUTS', 'TR-HRSOD', 'TR-UHRSD', 'TR-DUTS+TR-HRSOD', 'TR-DUTS+TR-UHRSD', 'TR-HRSOD+TR-UHRSD', 'TR-DUTS+TR-HRSOD+TR-UHRSD'][5],
'General': 'DIS-TE1+DIS-TE2+DIS-TE3+DIS-TE4+DIS-TR+TR-HRSOD+TE-HRSOD+TR-HRS10K+TE-HRS10K+TR-UHRSD+TE-UHRSD+TR-P3M-10k+TE-P3M-500-NP+TE-P3M-500-P+TR-humans', # leave DIS-VD for evaluation.
'Matting': 'TR-P3M-10k+TE-P3M-500-NP+TR-humans+TR-Distrinctions-646',
}[self.task]
self.prompt4loc = ['dense', 'sparse'][0]
# Faster-Training settings
self.load_all = False # Turn it on/off by your case. It may consume a lot of CPU memory. And for multi-GPU (N), it would cost N times the CPU memory to load the data.
self.use_fp16 = False # It may cause nan in training.
self.compile = True and (not self.use_fp16) # 1. Trigger CPU memory leak in some extend, which is an inherent problem of PyTorch.
# Machines with > 70GB CPU memory can run the whole training on DIS5K with default setting.
# 2. Higher PyTorch version may fix it: https://github.com/pytorch/pytorch/issues/119607.
# 3. But compile in Pytorch > 2.0.1 seems to bring no acceleration for training.
self.precisionHigh = True
# MODEL settings
self.ms_supervision = True
self.out_ref = self.ms_supervision and True
self.dec_ipt = True
self.dec_ipt_split = True
self.locate_head = False
self.cxt_num = [0, 3][1] # multi-scale skip connections from encoder
self.mul_scl_ipt = ['', 'add', 'cat'][2]
self.refine = ['', 'itself', 'RefUNet', 'Refiner', 'RefinerPVTInChannels4'][0]
self.progressive_ref = self.refine and True
self.ender = self.progressive_ref and False
self.scale = self.progressive_ref and 2
self.dec_att = ['', 'ASPP', 'ASPPDeformable'][2]
self.squeeze_block = ['', 'BasicDecBlk_x1', 'ResBlk_x4', 'ASPP_x3', 'ASPPDeformable_x3'][1]
self.dec_blk = ['BasicDecBlk', 'ResBlk', 'HierarAttDecBlk'][0]
self.auxiliary_classification = False
self.refine_iteration = 1
self.freeze_bb = False
self.precisionHigh = True
self.compile = True
self.load_all = True
self.verbose_eval = True
self.dec_blk = ['BasicDecBlk', 'ResBlk'][0]
# TRAINING settings
self.batch_size = 4
self.finetune_last_epochs = [
('IoU', 0),
{
'DIS5K': ('IoU', -30),
'COD': ('IoU', -20),
'HRSOD': ('IoU', -20),
'General': ('MAE', -10),
'Matting': ('MAE', -10),
}[self.task]
][1] # choose 0 to skip
self.lr = (1e-4 if 'DIS5K' in self.task else 1e-5) * math.sqrt(self.batch_size / 4) # DIS needs high lr to converge faster. Adapt the lr linearly
self.size = 1024
self.batch_size = 2
self.IoU_finetune_last_epochs = [0, -40][1] # choose 0 to skip
if self.dec_blk == 'HierarAttDecBlk':
self.batch_size = 2 ** [0, 1, 2, 3, 4][2]
self.model = [
'BiRefNet',
][0]
self.num_workers = max(4, self.batch_size) # will be decrease to min(it, batch_size) at the initialization of the data_loader
# Components
self.lat_blk = ['BasicLatBlk'][0]
self.dec_channels_inter = ['fixed', 'adap'][0]
# Backbone
# Backbone settings
self.bb = [
'vgg16', 'vgg16bn', 'resnet50', # 0, 1, 2
'pvt_v2_b2', 'pvt_v2_b5', # 3-bs10, 4-bs5
'swin_v1_b', 'swin_v1_l' # 5-bs9, 6-bs6
'swin_v1_t', 'swin_v1_s', # 3, 4
'swin_v1_b', 'swin_v1_l', # 5-bs9, 6-bs4
'pvt_v2_b0', 'pvt_v2_b1', # 7, 8
'pvt_v2_b2', 'pvt_v2_b5', # 9-bs10, 10-bs5
][6]
self.lateral_channels_in_collection = {
'vgg16': [512, 256, 128, 64], 'vgg16bn': [512, 256, 128, 64], 'resnet50': [1024, 512, 256, 64],
'pvt_v2_b2': [512, 320, 128, 64], 'pvt_v2_b5': [512, 320, 128, 64],
'swin_v1_b': [1024, 512, 256, 128], 'swin_v1_l': [1536, 768, 384, 192],
'swin_v1_t': [768, 384, 192, 96], 'swin_v1_s': [768, 384, 192, 96],
'pvt_v2_b0': [256, 160, 64, 32], 'pvt_v2_b1': [512, 320, 128, 64],
}[self.bb]
if self.mul_scl_ipt == 'cat':
self.lateral_channels_in_collection = [channel * 2 for channel in self.lateral_channels_in_collection]
self.cxt = self.lateral_channels_in_collection[1:][::-1][-self.cxt_num:] if self.cxt_num else []
self.sys_home_dir = models_dir
self.weights_root_dir = os.path.join(self.sys_home_dir, "BiRefNet")
self.weights = {
'pvt_v2_b2': os.path.join(self.weights_root_dir, 'pvt_v2_b2.pth'),
'pvt_v2_b5': os.path.join(self.weights_root_dir, ['pvt_v2_b5.pth', 'pvt_v2_b5_22k.pth'][0]),
'swin_v1_b': os.path.join(self.weights_root_dir, ['swin_base_patch4_window12_384_22kto1k.pth', 'swin_base_patch4_window12_384_22k.pth'][0]),
'swin_v1_l': os.path.join(self.weights_root_dir, ['swin_large_patch4_window12_384_22kto1k.pth', 'swin_large_patch4_window12_384_22k.pth'][0]),
}
# Training
self.num_workers = 5 # will be decrease to min(it, batch_size) at the initialization of the data_loader
self.optimizer = ['Adam', 'AdamW'][0]
self.lr = 1e-5 * math.sqrt(self.batch_size / 5) # adapt the lr linearly
self.lr_decay_epochs = [1e4] # Set to negative N to decay the lr in the last N-th epoch.
self.lr_decay_rate = 0.5
self.only_S_MAE = False
self.SDPA_enabled = False # Bug. Slower and errors occur in multi-GPUs
# MODEL settings - inactive
self.lat_blk = ['BasicLatBlk'][0]
self.dec_channels_inter = ['fixed', 'adap'][0]
self.refine = ['', 'itself', 'RefUNet', 'Refiner', 'RefinerPVTInChannels4'][0]
self.progressive_ref = self.refine and True
self.ender = self.progressive_ref and False
self.scale = self.progressive_ref and 2
self.auxiliary_classification = False # Only for DIS5K, where class labels are saved in `dataset.py`.
self.refine_iteration = 1
self.freeze_bb = False
self.model = [
'BiRefNet',
][0]
# Data
self.data_root_dir = os.path.join(self.sys_home_dir, 'datasets/dis')
self.dataset = ['DIS5K', 'COD', 'SOD'][0]
# TRAINING settings - inactive
self.preproc_methods = ['flip', 'enhance', 'rotate', 'pepper', 'crop'][:4]
self.optimizer = ['Adam', 'AdamW'][1]
self.lr_decay_epochs = [1e5] # Set to negative N to decay the lr in the last N-th epoch.
self.lr_decay_rate = 0.5
# Loss
self.lambdas_pix_last = {
# not 0 means opening this loss
# original rate -- 1 : 30 : 1.5 : 0.2, bce x 30
'bce': 30 * 1, # high performance
'iou': 0.5 * 1, # 0 / 255
'iou_patch': 0.5 * 0, # 0 / 255, win_size = (64, 64)
'mse': 150 * 0, # can smooth the saliency map
'triplet': 3 * 0,
'reg': 100 * 0,
'ssim': 10 * 1, # help contours,
'cnt': 5 * 0, # help contours
}
if self.task not in ['Matting']:
self.lambdas_pix_last = {
# not 0 means opening this loss
# original rate -- 1 : 30 : 1.5 : 0.2, bce x 30
'bce': 30 * 1, # high performance
'iou': 0.5 * 1, # 0 / 255
'iou_patch': 0.5 * 0, # 0 / 255, win_size = (64, 64)
'mae': 30 * 0,
'mse': 30 * 0, # can smooth the saliency map
'triplet': 3 * 0,
'reg': 100 * 0,
'ssim': 10 * 1, # help contours,
'cnt': 5 * 0, # help contours
'structure': 5 * 0, # structure loss from codes of MVANet. A little improvement on DIS-TE[1,2,3], a bit more decrease on DIS-TE4.
}
else:
self.lambdas_pix_last = {
# not 0 means opening this loss
# original rate -- 1 : 30 : 1.5 : 0.2, bce x 30
'bce': 30 * 0, # high performance
'iou': 0.5 * 0, # 0 / 255
'iou_patch': 0.5 * 0, # 0 / 255, win_size = (64, 64)
'mae': 100 * 1,
'mse': 30 * 0, # can smooth the saliency map
'triplet': 3 * 0,
'reg': 100 * 0,
'ssim': 10 * 1, # help contours,
'cnt': 5 * 0, # help contours
'structure': 5 * 0, # structure loss from codes of MVANet. A little improvement on DIS-TE[1,2,3], a bit more decrease on DIS-TE4.
}
self.lambdas_cls = {
'ce': 5.0
}
@@ -97,8 +134,41 @@ class Config():
self.lambda_adv_g = 10. * 0 # turn to 0 to avoid adv training
self.lambda_adv_d = 3. * (self.lambda_adv_g > 0)
# PATH settings - inactive
self.data_root_dir = os.path.join(self.sys_home_dir, 'datasets/dis')
self.weights_root_dir = os.path.join(self.sys_home_dir, 'weights')
self.weights = {
'pvt_v2_b2': os.path.join(self.weights_root_dir, 'pvt_v2_b2.pth'),
'pvt_v2_b5': os.path.join(self.weights_root_dir, ['pvt_v2_b5.pth', 'pvt_v2_b5_22k.pth'][0]),
'swin_v1_b': os.path.join(self.weights_root_dir, ['swin_base_patch4_window12_384_22kto1k.pth', 'swin_base_patch4_window12_384_22k.pth'][0]),
'swin_v1_l': os.path.join(self.weights_root_dir, ['swin_large_patch4_window12_384_22kto1k.pth', 'swin_large_patch4_window12_384_22k.pth'][0]),
'swin_v1_t': os.path.join(self.weights_root_dir, ['swin_tiny_patch4_window7_224_22kto1k_finetune.pth'][0]),
'swin_v1_s': os.path.join(self.weights_root_dir, ['swin_small_patch4_window7_224_22kto1k_finetune.pth'][0]),
'pvt_v2_b0': os.path.join(self.weights_root_dir, ['pvt_v2_b0.pth'][0]),
'pvt_v2_b1': os.path.join(self.weights_root_dir, ['pvt_v2_b1.pth'][0]),
}
# Callbacks - inactive
self.verbose_eval = True
self.only_S_MAE = False
self.SDPA_enabled = False # Bugs. Slower and errors occur in multi-GPUs
# others
self.device = [0, 'cpu'][0] # .to(0) = .to('cuda:0')
self.device = [0, 'cpu'][0] # .to(0) == .to('cuda:0')
self.batch_size_valid = 1
self.rand_seed = 7
run_sh_file = [f for f in os.listdir('.') if 'train.sh' == f] + [os.path.join('..', f) for f in os.listdir('..') if 'train.sh' == f]
if run_sh_file:
with open(run_sh_file[0], 'r') as f:
lines = f.readlines()
self.save_last = int([l.strip() for l in lines if '"{}")'.format(self.task) in l and 'val_last=' in l][0].split('val_last=')[-1].split()[0])
def print_task(self) -> None:
# Return task for choosing settings in shell scripts.
print(self.task)
if __name__ == '__main__':
config = Config()
config.print_task()
py/BiRefNet/dataset.py
+53 -75
View File
@@ -5,19 +5,34 @@ from PIL import Image
from torch.utils import data
from torchvision import transforms
from .preproc import preproc
from .config import Config
from glob import glob
from BiRefNet.image_proc import preproc
from BiRefNet.config import Config
from BiRefNet.utils import path_to_image
Image.MAX_IMAGE_PIXELS = None # remove DecompressionBombWarning
config = Config()
_class_labels_TR_sorted = 'Airplane, Ant, Antenna, Archery, Axe, BabyCarriage, Bag, BalanceBeam, Balcony, Balloon, Basket, BasketballHoop, Beatle, Bed, Bee, Bench, Bicycle, BicycleFrame, BicycleStand, Boat, Bonsai, BoomLift, Bridge, BunkBed, Butterfly, Button, Cable, CableLift, Cage, Camcorder, Cannon, Canoe, Car, CarParkDropArm, Carriage, Cart, Caterpillar, CeilingLamp, Centipede, Chair, Clip, Clock, Clothes, CoatHanger, Comb, ConcretePumpTruck, Crack, Crane, Cup, DentalChair, Desk, DeskChair, Diagram, DishRack, DoorHandle, Dragonfish, Dragonfly, Drum, Earphone, Easel, ElectricIron, Excavator, Eyeglasses, Fan, Fence, Fencing, FerrisWheel, FireExtinguisher, Fishing, Flag, FloorLamp, Forklift, GasStation, Gate, Gear, Goal, Golf, GymEquipment, Hammock, Handcart, Handcraft, Handrail, HangGlider, Harp, Harvester, Headset, Helicopter, Helmet, Hook, HorizontalBar, Hydrovalve, IroningTable, Jewelry, Key, KidsPlayground, Kitchenware, Kite, Knife, Ladder, LaundryRack, Lightning, Lobster, Locust, Machine, MachineGun, MagazineRack, Mantis, Medal, MemorialArchway, Microphone, Missile, MobileHolder, Monitor, Mosquito, Motorcycle, MovingTrolley, Mower, MusicPlayer, MusicStand, ObservationTower, Octopus, OilWell, OlympicLogo, OperatingTable, OutdoorFitnessEquipment, Parachute, Pavilion, Piano, Pipe, PlowHarrow, PoleVault, Punchbag, Rack, Racket, Rifle, Ring, Robot, RockClimbing, Rope, Sailboat, Satellite, Scaffold, Scale, Scissor, Scooter, Sculpture, Seadragon, Seahorse, Seal, SewingMachine, Ship, Shoe, ShoppingCart, ShoppingTrolley, Shower, Shrimp, Signboard, Skateboarding, Skeleton, Skiing, Spade, SpeedBoat, Spider, Spoon, Stair, Stand, Stationary, SteeringWheel, Stethoscope, Stool, Stove, StreetLamp, SweetStand, Swing, Sword, TV, Table, TableChair, TableLamp, TableTennis, Tank, Tapeline, Teapot, Telescope, Tent, TobaccoPipe, Toy, Tractor, TrafficLight, TrafficSign, Trampoline, TransmissionTower, Tree, Tricycle, TrimmerCover, Tripod, Trombone, Truck, Trumpet, Tuba, UAV, Umbrella, UnevenBars, UtilityPole, VacuumCleaner, Violin, Wakesurfing, Watch, WaterTower, WateringPot, Well, WellLid, Wheel, Wheelchair, WindTurbine, Windmill, WineGlass, WireWhisk, Yacht'
_class_labels_TR_sorted = (
'Airplane, Ant, Antenna, Archery, Axe, BabyCarriage, Bag, BalanceBeam, Balcony, Balloon, Basket, BasketballHoop, Beatle, Bed, Bee, Bench, Bicycle, '
'BicycleFrame, BicycleStand, Boat, Bonsai, BoomLift, Bridge, BunkBed, Butterfly, Button, Cable, CableLift, Cage, Camcorder, Cannon, Canoe, Car, '
'CarParkDropArm, Carriage, Cart, Caterpillar, CeilingLamp, Centipede, Chair, Clip, Clock, Clothes, CoatHanger, Comb, ConcretePumpTruck, Crack, Crane, '
'Cup, DentalChair, Desk, DeskChair, Diagram, DishRack, DoorHandle, Dragonfish, Dragonfly, Drum, Earphone, Easel, ElectricIron, Excavator, Eyeglasses, '
'Fan, Fence, Fencing, FerrisWheel, FireExtinguisher, Fishing, Flag, FloorLamp, Forklift, GasStation, Gate, Gear, Goal, Golf, GymEquipment, Hammock, '
'Handcart, Handcraft, Handrail, HangGlider, Harp, Harvester, Headset, Helicopter, Helmet, Hook, HorizontalBar, Hydrovalve, IroningTable, Jewelry, Key, '
'KidsPlayground, Kitchenware, Kite, Knife, Ladder, LaundryRack, Lightning, Lobster, Locust, Machine, MachineGun, MagazineRack, Mantis, Medal, MemorialArchway, '
'Microphone, Missile, MobileHolder, Monitor, Mosquito, Motorcycle, MovingTrolley, Mower, MusicPlayer, MusicStand, ObservationTower, Octopus, OilWell, '
'OlympicLogo, OperatingTable, OutdoorFitnessEquipment, Parachute, Pavilion, Piano, Pipe, PlowHarrow, PoleVault, Punchbag, Rack, Racket, Rifle, Ring, Robot, '
'RockClimbing, Rope, Sailboat, Satellite, Scaffold, Scale, Scissor, Scooter, Sculpture, Seadragon, Seahorse, Seal, SewingMachine, Ship, Shoe, ShoppingCart, '
'ShoppingTrolley, Shower, Shrimp, Signboard, Skateboarding, Skeleton, Skiing, Spade, SpeedBoat, Spider, Spoon, Stair, Stand, Stationary, SteeringWheel, '
'Stethoscope, Stool, Stove, StreetLamp, SweetStand, Swing, Sword, TV, Table, TableChair, TableLamp, TableTennis, Tank, Tapeline, Teapot, Telescope, Tent, '
'TobaccoPipe, Toy, Tractor, TrafficLight, TrafficSign, Trampoline, TransmissionTower, Tree, Tricycle, TrimmerCover, Tripod, Trombone, Truck, Trumpet, Tuba, '
'UAV, Umbrella, UnevenBars, UtilityPole, VacuumCleaner, Violin, Wakesurfing, Watch, WaterTower, WateringPot, Well, WellLid, Wheel, Wheelchair, WindTurbine, Windmill, WineGlass, WireWhisk, Yacht'
)
class_labels_TR_sorted = _class_labels_TR_sorted.split(', ')
class MyData(data.Dataset):
def __init__(self, data_root, image_size, is_train=True):
def __init__(self, datasets, image_size, is_train=True):
self.size_train = image_size
self.size_test = image_size
self.keep_size = not config.size
@@ -25,7 +40,8 @@ class MyData(data.Dataset):
self.is_train = is_train
self.load_all = config.load_all
self.device = config.device
self.dataset = data_root.replace('\\', '/').split('/')[-1]
valid_extensions = ['.png', '.jpg', '.PNG', '.JPG', '.JPEG']
if self.is_train and config.auxiliary_classification:
self.cls_name2id = {_name: _id for _id, _name in enumerate(class_labels_TR_sorted)}
self.transform_image = transforms.Compose([
@@ -37,42 +53,55 @@ class MyData(data.Dataset):
transforms.Resize(self.data_size),
transforms.ToTensor(),
][self.load_all or self.keep_size:])
## 'im' and 'gt' need modifying
image_root = os.path.join(data_root, 'im')
self.image_paths = [os.path.join(image_root, p) for p in os.listdir(image_root)]
self.label_paths = [p.replace('/im/', '/gt/').replace('.jpg', '.png') for p in self.image_paths]
dataset_root = os.path.join(config.data_root_dir, config.task)
# datasets can be a list of different datasets for training on combined sets.
self.image_paths = []
for dataset in datasets.split('+'):
image_root = os.path.join(dataset_root, dataset, 'im')
self.image_paths += [os.path.join(image_root, p) for p in os.listdir(image_root) if any(p.endswith(ext) for ext in valid_extensions)]
self.label_paths = []
for p in self.image_paths:
for ext in valid_extensions:
## 'im' and 'gt' may need modifying
p_gt = p.replace('/im/', '/gt/')[:-(len(p.split('.')[-1])+1)] + ext
file_exists = False
if os.path.exists(p_gt):
self.label_paths.append(p_gt)
file_exists = True
break
if not file_exists:
print('Not exists:', p_gt)
if len(self.label_paths) != len(self.image_paths):
raise ValueError(f"There are different numbers of images ({len(self.label_paths)}) and labels ({len(self.image_paths)})")
if self.load_all:
self.images_loaded, self.labels_loaded = [], []
self.class_labels_loaded = []
# for image_path, label_path in zip(self.image_paths, self.label_paths):
for image_path, label_path in tqdm(zip(self.image_paths, self.label_paths), total=len(self.image_paths)):
_image = cv2.imread(image_path)
_label = cv2.imread(label_path, cv2.IMREAD_GRAYSCALE)
if not self.keep_size:
_image_rs = cv2.resize(_image, (config.size, config.size), interpolation=cv2.INTER_LINEAR)
_label_rs = cv2.resize(_label, (config.size, config.size), interpolation=cv2.INTER_LINEAR)
self.images_loaded.append(
Image.fromarray(cv2.cvtColor(_image_rs, cv2.COLOR_BGR2RGB)).convert('RGB')
)
self.labels_loaded.append(
Image.fromarray(_label_rs).convert('L')
)
_image = path_to_image(image_path, size=(config.size, config.size), color_type='rgb')
_label = path_to_image(label_path, size=(config.size, config.size), color_type='gray')
self.images_loaded.append(_image)
self.labels_loaded.append(_label)
self.class_labels_loaded.append(
self.cls_name2id[label_path.split('/')[-1].split('#')[3]] if self.is_train and config.auxiliary_classification else -1
)
def __getitem__(self, index):
if self.load_all:
image = self.images_loaded[index]
label = self.labels_loaded[index]
class_label = self.class_labels_loaded[index] if self.is_train and config.auxiliary_classification else -1
else:
image = Image.open(self.image_paths[index]).convert('RGB')
image = path_to_image(self.image_paths[index], size=(config.size, config.size), color_type='rgb')
label = path_to_image(self.label_paths[index], size=(config.size, config.size), color_type='gray')
class_label = self.cls_name2id[self.label_paths[index].split('/')[-1].split('#')[3]] if self.is_train and config.auxiliary_classification else -1
# loading image and label
if self.is_train:
image, label = preproc(image, image, preproc_methods=config.preproc_methods)
image, label = preproc(image, label, preproc_methods=config.preproc_methods)
# else:
# if _label.shape[0] > 2048 or _label.shape[1] > 2048:
# _image = cv2.resize(_image, (2048, 2048), interpolation=cv2.INTER_LINEAR)
@@ -87,54 +116,3 @@ class MyData(data.Dataset):
def __len__(self):
return len(self.image_paths)
class YouData(data.Dataset):
def __init__(self, data_root, image_size, is_train=True):
self.size_train = image_size
self.size_test = image_size
self.keep_size = not config.size
self.data_size = (config.size, config.size)
self.is_train = is_train
self.load_all = config.load_all
self.device = config.device
self.dataset = data_root.replace('\\', '/').split('/')[-1]
if self.is_train and config.auxiliary_classification:
self.cls_name2id = {_name: _id for _id, _name in enumerate(class_labels_TR_sorted)}
self.transform_image = transforms.Compose([
transforms.Resize(self.data_size),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
][self.load_all or self.keep_size:])
## 'im' and 'gt' need modifying
self.image_paths = glob(data_root + "/*")
self.img_sizes = []
if self.load_all:
self.images_loaded, self.labels_loaded = [], []
for image_path in tqdm(self.image_paths, total=len(self.image_paths)):
_image = cv2.imread(image_path)
if not self.keep_size:
_image_rs = cv2.resize(_image, (config.size, config.size), interpolation=cv2.INTER_LINEAR)
self.images_loaded.append(
Image.fromarray(cv2.cvtColor(_image_rs, cv2.COLOR_BGR2RGB)).convert('RGB')
)
self.img_sizes.append(_image.shape[:2])
def __getitem__(self, index):
if self.load_all:
image = self.images_loaded[index]
else:
image = Image.open(self.image_paths[index]).convert('RGB')
# loading image and label
if self.is_train:
image, _ = preproc(image, image, preproc_methods=config.preproc_methods)
image = self.transform_image(image)
size = self.img_sizes[index]
return image, size
def __len__(self):
return len(self.image_paths)
py/BiRefNet/eval_existingOnes.py
+146
View File
@@ -0,0 +1,146 @@
import os
import argparse
from glob import glob
import prettytable as pt
from BiRefNet.evaluation.evaluate import evaluator
from BiRefNet.config import Config
config = Config()
def do_eval(args):
# evaluation for whole dataset
# dataset first in evaluation
for _data_name in args.data_lst.split('+'):
pred_data_dir = sorted(glob(os.path.join(args.pred_root, args.model_lst[0], _data_name)))
if not pred_data_dir:
print('Skip dataset {}.'.format(_data_name))
continue
gt_src = os.path.join(args.gt_root, _data_name)
gt_paths = sorted(glob(os.path.join(gt_src, 'gt', '*')))
print('#' * 20, _data_name, '#' * 20)
filename = os.path.join(args.save_dir, '{}_eval.txt'.format(_data_name))
tb = pt.PrettyTable()
tb.vertical_char = '&'
if config.task == 'DIS5K':
tb.field_names = ["Dataset", "Method", "maxFm", "wFmeasure", 'MAE', "Smeasure", "meanEm", "HCE", "maxEm", "meanFm", "adpEm", "adpFm", 'mBA', 'maxBIoU', 'meanBIoU']
elif config.task == 'COD':
tb.field_names = ["Dataset", "Method", "Smeasure", "wFmeasure", "meanFm", "meanEm", "maxEm", 'MAE', "maxFm", "adpEm", "adpFm", "HCE", 'mBA', 'maxBIoU', 'meanBIoU']
elif config.task == 'HRSOD':
tb.field_names = ["Dataset", "Method", "Smeasure", "maxFm", "meanEm", 'MAE', "maxEm", "meanFm", "wFmeasure", "adpEm", "adpFm", "HCE", 'mBA', 'maxBIoU', 'meanBIoU']
elif config.task == 'General':
tb.field_names = ["Dataset", "Method", "maxFm", "wFmeasure", 'MAE', "Smeasure", "meanEm", "HCE", "maxEm", "meanFm", "adpEm", "adpFm", 'mBA', 'maxBIoU', 'meanBIoU']
elif config.task == 'Matting':
tb.field_names = ["Dataset", "Method", "Smeasure", "maxFm", "meanEm", 'MSE', "maxEm", "meanFm", "wFmeasure", "adpEm", "adpFm", "HCE", 'mBA', 'maxBIoU', 'meanBIoU']
else:
tb.field_names = ["Dataset", "Method", "Smeasure", 'MAE', "maxEm", "meanEm", "maxFm", "meanFm", "wFmeasure", "adpEm", "adpFm", "HCE", 'mBA', 'maxBIoU', 'meanBIoU']
for _model_name in args.model_lst[:]:
print('\t', 'Evaluating model: {}...'.format(_model_name))
pred_paths = [p.replace(args.gt_root, os.path.join(args.pred_root, _model_name)).replace('/gt/', '/') for p in gt_paths]
# print(pred_paths[:1], gt_paths[:1])
em, sm, fm, mae, wfm, hce, mba, biou = evaluator(
gt_paths=gt_paths,
pred_paths=pred_paths,
metrics=args.metrics.split('+'),
verbose=config.verbose_eval
)
if config.task == 'DIS5K':
scores = [
fm['curve'].max().round(3), wfm.round(3), mae.round(3), sm.round(3), em['curve'].mean().round(3), int(hce.round()),
em['curve'].max().round(3), fm['curve'].mean().round(3), em['adp'].round(3), fm['adp'].round(3),
mba.round(3), biou['curve'].max().round(3), biou['curve'].mean().round(3),
]
elif config.task == 'COD':
scores = [
sm.round(3), wfm.round(3), fm['curve'].mean().round(3), em['curve'].mean().round(3), em['curve'].max().round(3), mae.round(3),
fm['curve'].max().round(3), em['adp'].round(3), fm['adp'].round(3), int(hce.round()),
mba.round(3), biou['curve'].max().round(3), biou['curve'].mean().round(3),
]
elif config.task == 'HRSOD':
scores = [
sm.round(3), fm['curve'].max().round(3), em['curve'].mean().round(3), mae.round(3),
em['curve'].max().round(3), fm['curve'].mean().round(3), wfm.round(3), em['adp'].round(3), fm['adp'].round(3), int(hce.round()),
mba.round(3), biou['curve'].max().round(3), biou['curve'].mean().round(3),
]
elif config.task == 'General':
scores = [
fm['curve'].max().round(3), wfm.round(3), mae.round(3), sm.round(3), em['curve'].mean().round(3), int(hce.round()),
em['curve'].max().round(3), fm['curve'].mean().round(3), em['adp'].round(3), fm['adp'].round(3),
mba.round(3), biou['curve'].max().round(3), biou['curve'].mean().round(3),
]
elif config.task == 'Matting':
scores = [
sm.round(3), fm['curve'].max().round(3), em['curve'].mean().round(3), mse.round(3),
em['curve'].max().round(3), fm['curve'].mean().round(3), wfm.round(3), em['adp'].round(3), fm['adp'].round(3), int(hce.round()),
mba.round(3), biou['curve'].max().round(3), biou['curve'].mean().round(3),
]
else:
scores = [
sm.round(3), mae.round(3), em['curve'].max().round(3), em['curve'].mean().round(3),
fm['curve'].max().round(3), fm['curve'].mean().round(3), wfm.round(3),
em['adp'].round(3), fm['adp'].round(3), int(hce.round()),
mba.round(3), biou['curve'].max().round(3), biou['curve'].mean().round(3),
]
for idx_score, score in enumerate(scores):
scores[idx_score] = '.' + format(score, '.3f').split('.')[-1] if score <= 1 else format(score, '<4')
records = [_data_name, _model_name] + scores
tb.add_row(records)
# Write results after every check.
with open(filename, 'w+') as file_to_write:
file_to_write.write(str(tb)+'\n')
print(tb)
if __name__ == '__main__':
# set parameters
parser = argparse.ArgumentParser()
parser.add_argument(
'--gt_root', type=str, help='ground-truth root',
default=os.path.join(config.data_root_dir, config.task))
parser.add_argument(
'--pred_root', type=str, help='prediction root',
default='./e_preds')
parser.add_argument(
'--data_lst', type=str, help='test dataset',
default={
'DIS5K': '+'.join(['DIS-VD', 'DIS-TE1', 'DIS-TE2', 'DIS-TE3', 'DIS-TE4'][:]),
'COD': '+'.join(['TE-COD10K', 'NC4K', 'TE-CAMO', 'CHAMELEON'][:]),
'HRSOD': '+'.join(['DAVIS-S', 'TE-HRSOD', 'TE-UHRSD', 'TE-DUTS', 'DUT-OMRON'][:]),
'General': '+'.join(['DIS-VD'][:]),
'Matting': '+'.join(['TE-P3M-500-P'][:]),
}[config.task])
parser.add_argument(
'--save_dir', type=str, help='candidate competitors',
default='e_results')
parser.add_argument(
'--check_integrity', type=bool, help='whether to check the file integrity',
default=False)
parser.add_argument(
'--metrics', type=str, help='candidate competitors',
default='+'.join(['S', 'MAE', 'E', 'F', 'WF', 'MBA', 'BIoU', 'HCE'][:100 if 'DIS5K' in config.task else -1]))
args = parser.parse_args()
args.metrics = '+'.join(['S', 'MAE', 'E', 'F', 'WF', 'MBA', 'BIoU', 'HCE'][:100 if sum(['DIS-' in _data for _data in args.data_lst.split('+')]) else -1])
os.makedirs(args.save_dir, exist_ok=True)
try:
args.model_lst = [m for m in sorted(os.listdir(args.pred_root), key=lambda x: int(x.split('epoch_')[-1]), reverse=True) if int(m.split('epoch_')[-1]) % 1 == 0]
except:
args.model_lst = [m for m in sorted(os.listdir(args.pred_root))]
# check the integrity of each candidates
if args.check_integrity:
for _data_name in args.data_lst.split('+'):
for _model_name in args.model_lst:
gt_pth = os.path.join(args.gt_root, _data_name)
pred_pth = os.path.join(args.pred_root, _model_name, _data_name)
if not sorted(os.listdir(gt_pth)) == sorted(os.listdir(pred_pth)):
print(len(sorted(os.listdir(gt_pth))), len(sorted(os.listdir(pred_pth))))
print('The {} Dataset of {} Model is not matching to the ground-truth'.format(_data_name, _model_name))
else:
print('>>> skip check the integrity of each candidates')
# start engine
do_eval(args)
py/BiRefNet/evaluation/metrics.py
+763
View File
@@ -0,0 +1,763 @@
import os
from tqdm import tqdm
import cv2
import numpy as np
from scipy.ndimage import convolve, distance_transform_edt as bwdist
from skimage.morphology import skeletonize
from skimage.morphology import disk
from skimage.measure import label
_EPS = np.spacing(1)
_TYPE = np.float64
def evaluator(gt_paths, pred_paths, metrics=['S', 'MAE', 'E', 'F', 'WF', 'MBA', 'BIoU', 'HCE'], verbose=False):
# define measures
if 'E' in metrics:
EM = EMeasure()
if 'S' in metrics:
SM = SMeasure()
if 'F' in metrics:
FM = FMeasure()
if 'MAE' in metrics:
MAE = MAEMeasure()
if 'WF' in metrics:
WFM = WeightedFMeasure()
if 'HCE' in metrics:
HCE = HCEMeasure()
if 'MBA' in metrics:
MBA = MBAMeasure()
if 'BIoU' in metrics:
BIoU = BIoUMeasure()
if isinstance(gt_paths, list) and isinstance(pred_paths, list):
# print(len(gt_paths), len(pred_paths))
assert len(gt_paths) == len(pred_paths)
for idx_sample in tqdm(range(len(gt_paths)), total=len(gt_paths)) if verbose else range(len(gt_paths)):
gt = gt_paths[idx_sample]
pred = pred_paths[idx_sample]
pred = pred[:-4] + '.png'
valid_extensions = ['.png', '.jpg', '.PNG', '.JPG', '.JPEG']
file_exists = False
for ext in valid_extensions:
if os.path.exists(pred[:-4] + ext):
pred = pred[:-4] + ext
file_exists = True
break
if file_exists:
pred_ary = cv2.imread(pred, cv2.IMREAD_GRAYSCALE)
else:
print('Not exists:', pred)
gt_ary = cv2.imread(gt, cv2.IMREAD_GRAYSCALE)
pred_ary = cv2.resize(pred_ary, (gt_ary.shape[1], gt_ary.shape[0]))
if 'E' in metrics:
EM.step(pred=pred_ary, gt=gt_ary)
if 'S' in metrics:
SM.step(pred=pred_ary, gt=gt_ary)
if 'F' in metrics:
FM.step(pred=pred_ary, gt=gt_ary)
if 'MAE' in metrics:
MAE.step(pred=pred_ary, gt=gt_ary)
if 'WF' in metrics:
WFM.step(pred=pred_ary, gt=gt_ary)
if 'HCE' in metrics:
ske_path = gt.replace('/gt/', '/ske/')
if os.path.exists(ske_path):
ske_ary = cv2.imread(ske_path, cv2.IMREAD_GRAYSCALE)
ske_ary = ske_ary > 128
else:
ske_ary = skeletonize(gt_ary > 128)
ske_save_dir = os.path.join(*ske_path.split(os.sep)[:-1])
if ske_path[0] == os.sep:
ske_save_dir = os.sep + ske_save_dir
os.makedirs(ske_save_dir, exist_ok=True)
cv2.imwrite(ske_path, ske_ary.astype(np.uint8) * 255)
HCE.step(pred=pred_ary, gt=gt_ary, gt_ske=ske_ary)
if 'MBA' in metrics:
MBA.step(pred=pred_ary, gt=gt_ary)
if 'BIoU' in metrics:
BIoU.step(pred=pred_ary, gt=gt_ary)
if 'E' in metrics:
em = EM.get_results()['em']
else:
em = {'curve': np.array([np.float64(-1)]), 'adp': np.float64(-1)}
if 'S' in metrics:
sm = SM.get_results()['sm']
else:
sm = np.float64(-1)
if 'F' in metrics:
fm = FM.get_results()['fm']
else:
fm = {'curve': np.array([np.float64(-1)]), 'adp': np.float64(-1)}
if 'MAE' in metrics:
mae = MAE.get_results()['mae']
else:
mae = np.float64(-1)
if 'WF' in metrics:
wfm = WFM.get_results()['wfm']
else:
wfm = np.float64(-1)
if 'HCE' in metrics:
hce = HCE.get_results()['hce']
else:
hce = np.float64(-1)
if 'MBA' in metrics:
mba = MBA.get_results()['mba']
else:
mba = np.float64(-1)
if 'BIoU' in metrics:
biou = BIoU.get_results()['biou']
else:
biou = {'curve': np.array([np.float64(-1)])}
return em, sm, fm, mae, wfm, hce, mba, biou
def _prepare_data(pred: np.ndarray, gt: np.ndarray) -> tuple:
gt = gt > 128
pred = pred / 255
if pred.max() != pred.min():
pred = (pred - pred.min()) / (pred.max() - pred.min())
return pred, gt
def _get_adaptive_threshold(matrix: np.ndarray, max_value: float = 1) -> float:
return min(2 * matrix.mean(), max_value)
class FMeasure(object):
def __init__(self, beta: float = 0.3):
self.beta = beta
self.precisions = []
self.recalls = []
self.adaptive_fms = []
self.changeable_fms = []
def step(self, pred: np.ndarray, gt: np.ndarray):
pred, gt = _prepare_data(pred, gt)
adaptive_fm = self.cal_adaptive_fm(pred=pred, gt=gt)
self.adaptive_fms.append(adaptive_fm)
precisions, recalls, changeable_fms = self.cal_pr(pred=pred, gt=gt)
self.precisions.append(precisions)
self.recalls.append(recalls)
self.changeable_fms.append(changeable_fms)
def cal_adaptive_fm(self, pred: np.ndarray, gt: np.ndarray) -> float:
adaptive_threshold = _get_adaptive_threshold(pred, max_value=1)
binary_predcition = pred >= adaptive_threshold
area_intersection = binary_predcition[gt].sum()
if area_intersection == 0:
adaptive_fm = 0
else:
pre = area_intersection / np.count_nonzero(binary_predcition)
rec = area_intersection / np.count_nonzero(gt)
adaptive_fm = (1 + self.beta) * pre * rec / (self.beta * pre + rec)
return adaptive_fm
def cal_pr(self, pred: np.ndarray, gt: np.ndarray) -> tuple:
pred = (pred * 255).astype(np.uint8)
bins = np.linspace(0, 256, 257)
fg_hist, _ = np.histogram(pred[gt], bins=bins)
bg_hist, _ = np.histogram(pred[~gt], bins=bins)
fg_w_thrs = np.cumsum(np.flip(fg_hist), axis=0)
bg_w_thrs = np.cumsum(np.flip(bg_hist), axis=0)
TPs = fg_w_thrs
Ps = fg_w_thrs + bg_w_thrs
Ps[Ps == 0] = 1
T = max(np.count_nonzero(gt), 1)
precisions = TPs / Ps
recalls = TPs / T
numerator = (1 + self.beta) * precisions * recalls
denominator = np.where(numerator == 0, 1, self.beta * precisions + recalls)
changeable_fms = numerator / denominator
return precisions, recalls, changeable_fms
def get_results(self) -> dict:
adaptive_fm = np.mean(np.array(self.adaptive_fms, _TYPE))
changeable_fm = np.mean(np.array(self.changeable_fms, dtype=_TYPE), axis=0)
precision = np.mean(np.array(self.precisions, dtype=_TYPE), axis=0) # N, 256
recall = np.mean(np.array(self.recalls, dtype=_TYPE), axis=0) # N, 256
return dict(fm=dict(adp=adaptive_fm, curve=changeable_fm),
pr=dict(p=precision, r=recall))
class MAEMeasure(object):
def __init__(self):
self.maes = []
def step(self, pred: np.ndarray, gt: np.ndarray):
pred, gt = _prepare_data(pred, gt)
mae = self.cal_mae(pred, gt)
self.maes.append(mae)
def cal_mae(self, pred: np.ndarray, gt: np.ndarray) -> float:
mae = np.mean(np.abs(pred - gt))
return mae
def get_results(self) -> dict:
mae = np.mean(np.array(self.maes, _TYPE))
return dict(mae=mae)
class SMeasure(object):
def __init__(self, alpha: float = 0.5):
self.sms = []
self.alpha = alpha
def step(self, pred: np.ndarray, gt: np.ndarray):
pred, gt = _prepare_data(pred=pred, gt=gt)
sm = self.cal_sm(pred, gt)
self.sms.append(sm)
def cal_sm(self, pred: np.ndarray, gt: np.ndarray) -> float:
y = np.mean(gt)
if y == 0:
sm = 1 - np.mean(pred)
elif y == 1:
sm = np.mean(pred)
else:
sm = self.alpha * self.object(pred, gt) + (1 - self.alpha) * self.region(pred, gt)
sm = max(0, sm)
return sm
def object(self, pred: np.ndarray, gt: np.ndarray) -> float:
fg = pred * gt
bg = (1 - pred) * (1 - gt)
u = np.mean(gt)
object_score = u * self.s_object(fg, gt) + (1 - u) * self.s_object(bg, 1 - gt)
return object_score
def s_object(self, pred: np.ndarray, gt: np.ndarray) -> float:
x = np.mean(pred[gt == 1])
sigma_x = np.std(pred[gt == 1], ddof=1)
score = 2 * x / (np.power(x, 2) + 1 + sigma_x + _EPS)
return score
def region(self, pred: np.ndarray, gt: np.ndarray) -> float:
x, y = self.centroid(gt)
part_info = self.divide_with_xy(pred, gt, x, y)
w1, w2, w3, w4 = part_info['weight']
pred1, pred2, pred3, pred4 = part_info['pred']
gt1, gt2, gt3, gt4 = part_info['gt']
score1 = self.ssim(pred1, gt1)
score2 = self.ssim(pred2, gt2)
score3 = self.ssim(pred3, gt3)
score4 = self.ssim(pred4, gt4)
return w1 * score1 + w2 * score2 + w3 * score3 + w4 * score4
def centroid(self, matrix: np.ndarray) -> tuple:
h, w = matrix.shape
area_object = np.count_nonzero(matrix)
if area_object == 0:
x = np.round(w / 2)
y = np.round(h / 2)
else:
# More details can be found at: https://www.yuque.com/lart/blog/gpbigm
y, x = np.argwhere(matrix).mean(axis=0).round()
return int(x) + 1, int(y) + 1
def divide_with_xy(self, pred: np.ndarray, gt: np.ndarray, x, y) -> dict:
h, w = gt.shape
area = h * w
gt_LT = gt[0:y, 0:x]
gt_RT = gt[0:y, x:w]
gt_LB = gt[y:h, 0:x]
gt_RB = gt[y:h, x:w]
pred_LT = pred[0:y, 0:x]
pred_RT = pred[0:y, x:w]
pred_LB = pred[y:h, 0:x]
pred_RB = pred[y:h, x:w]
w1 = x * y / area
w2 = y * (w - x) / area
w3 = (h - y) * x / area
w4 = 1 - w1 - w2 - w3
return dict(gt=(gt_LT, gt_RT, gt_LB, gt_RB),
pred=(pred_LT, pred_RT, pred_LB, pred_RB),
weight=(w1, w2, w3, w4))
def ssim(self, pred: np.ndarray, gt: np.ndarray) -> float:
h, w = pred.shape
N = h * w
x = np.mean(pred)
y = np.mean(gt)
sigma_x = np.sum((pred - x) ** 2) / (N - 1)
sigma_y = np.sum((gt - y) ** 2) / (N - 1)
sigma_xy = np.sum((pred - x) * (gt - y)) / (N - 1)
alpha = 4 * x * y * sigma_xy
beta = (x ** 2 + y ** 2) * (sigma_x + sigma_y)
if alpha != 0:
score = alpha / (beta + _EPS)
elif alpha == 0 and beta == 0:
score = 1
else:
score = 0
return score
def get_results(self) -> dict:
sm = np.mean(np.array(self.sms, dtype=_TYPE))
return dict(sm=sm)
class EMeasure(object):
def __init__(self):
self.adaptive_ems = []
self.changeable_ems = []
def step(self, pred: np.ndarray, gt: np.ndarray):
pred, gt = _prepare_data(pred=pred, gt=gt)
self.gt_fg_numel = np.count_nonzero(gt)
self.gt_size = gt.shape[0] * gt.shape[1]
changeable_ems = self.cal_changeable_em(pred, gt)
self.changeable_ems.append(changeable_ems)
adaptive_em = self.cal_adaptive_em(pred, gt)
self.adaptive_ems.append(adaptive_em)
def cal_adaptive_em(self, pred: np.ndarray, gt: np.ndarray) -> float:
adaptive_threshold = _get_adaptive_threshold(pred, max_value=1)
adaptive_em = self.cal_em_with_threshold(pred, gt, threshold=adaptive_threshold)
return adaptive_em
def cal_changeable_em(self, pred: np.ndarray, gt: np.ndarray) -> np.ndarray:
changeable_ems = self.cal_em_with_cumsumhistogram(pred, gt)
return changeable_ems
def cal_em_with_threshold(self, pred: np.ndarray, gt: np.ndarray, threshold: float) -> float:
binarized_pred = pred >= threshold
fg_fg_numel = np.count_nonzero(binarized_pred & gt)
fg_bg_numel = np.count_nonzero(binarized_pred & ~gt)
fg___numel = fg_fg_numel + fg_bg_numel
bg___numel = self.gt_size - fg___numel
if self.gt_fg_numel == 0:
enhanced_matrix_sum = bg___numel
elif self.gt_fg_numel == self.gt_size:
enhanced_matrix_sum = fg___numel
else:
parts_numel, combinations = self.generate_parts_numel_combinations(
fg_fg_numel=fg_fg_numel, fg_bg_numel=fg_bg_numel,
pred_fg_numel=fg___numel, pred_bg_numel=bg___numel,
)
results_parts = []
for i, (part_numel, combination) in enumerate(zip(parts_numel, combinations)):
align_matrix_value = 2 * (combination[0] * combination[1]) / \
(combination[0] ** 2 + combination[1] ** 2 + _EPS)
enhanced_matrix_value = (align_matrix_value + 1) ** 2 / 4
results_parts.append(enhanced_matrix_value * part_numel)
enhanced_matrix_sum = sum(results_parts)
em = enhanced_matrix_sum / (self.gt_size - 1 + _EPS)
return em
def cal_em_with_cumsumhistogram(self, pred: np.ndarray, gt: np.ndarray) -> np.ndarray:
pred = (pred * 255).astype(np.uint8)
bins = np.linspace(0, 256, 257)
fg_fg_hist, _ = np.histogram(pred[gt], bins=bins)
fg_bg_hist, _ = np.histogram(pred[~gt], bins=bins)
fg_fg_numel_w_thrs = np.cumsum(np.flip(fg_fg_hist), axis=0)
fg_bg_numel_w_thrs = np.cumsum(np.flip(fg_bg_hist), axis=0)
fg___numel_w_thrs = fg_fg_numel_w_thrs + fg_bg_numel_w_thrs
bg___numel_w_thrs = self.gt_size - fg___numel_w_thrs
if self.gt_fg_numel == 0:
enhanced_matrix_sum = bg___numel_w_thrs
elif self.gt_fg_numel == self.gt_size:
enhanced_matrix_sum = fg___numel_w_thrs
else:
parts_numel_w_thrs, combinations = self.generate_parts_numel_combinations(
fg_fg_numel=fg_fg_numel_w_thrs, fg_bg_numel=fg_bg_numel_w_thrs,
pred_fg_numel=fg___numel_w_thrs, pred_bg_numel=bg___numel_w_thrs,
)
results_parts = np.empty(shape=(4, 256), dtype=np.float64)
for i, (part_numel, combination) in enumerate(zip(parts_numel_w_thrs, combinations)):
align_matrix_value = 2 * (combination[0] * combination[1]) / \
(combination[0] ** 2 + combination[1] ** 2 + _EPS)
enhanced_matrix_value = (align_matrix_value + 1) ** 2 / 4
results_parts[i] = enhanced_matrix_value * part_numel
enhanced_matrix_sum = results_parts.sum(axis=0)
em = enhanced_matrix_sum / (self.gt_size - 1 + _EPS)
return em
def generate_parts_numel_combinations(self, fg_fg_numel, fg_bg_numel, pred_fg_numel, pred_bg_numel):
bg_fg_numel = self.gt_fg_numel - fg_fg_numel
bg_bg_numel = pred_bg_numel - bg_fg_numel
parts_numel = [fg_fg_numel, fg_bg_numel, bg_fg_numel, bg_bg_numel]
mean_pred_value = pred_fg_numel / self.gt_size
mean_gt_value = self.gt_fg_numel / self.gt_size
demeaned_pred_fg_value = 1 - mean_pred_value
demeaned_pred_bg_value = 0 - mean_pred_value
demeaned_gt_fg_value = 1 - mean_gt_value
demeaned_gt_bg_value = 0 - mean_gt_value
combinations = [
(demeaned_pred_fg_value, demeaned_gt_fg_value),
(demeaned_pred_fg_value, demeaned_gt_bg_value),
(demeaned_pred_bg_value, demeaned_gt_fg_value),
(demeaned_pred_bg_value, demeaned_gt_bg_value)
]
return parts_numel, combinations
def get_results(self) -> dict:
adaptive_em = np.mean(np.array(self.adaptive_ems, dtype=_TYPE))
changeable_em = np.mean(np.array(self.changeable_ems, dtype=_TYPE), axis=0)
return dict(em=dict(adp=adaptive_em, curve=changeable_em))
class WeightedFMeasure(object):
def __init__(self, beta: float = 1):
self.beta = beta
self.weighted_fms = []
def step(self, pred: np.ndarray, gt: np.ndarray):
pred, gt = _prepare_data(pred=pred, gt=gt)
if np.all(~gt):
wfm = 0
else:
wfm = self.cal_wfm(pred, gt)
self.weighted_fms.append(wfm)
def cal_wfm(self, pred: np.ndarray, gt: np.ndarray) -> float:
# [Dst,IDXT] = bwdist(dGT);
Dst, Idxt = bwdist(gt == 0, return_indices=True)
# %Pixel dependency
# E = abs(FG-dGT);
E = np.abs(pred - gt)
Et = np.copy(E)
Et[gt == 0] = Et[Idxt[0][gt == 0], Idxt[1][gt == 0]]
# K = fspecial('gaussian',7,5);
# EA = imfilter(Et,K);
K = self.matlab_style_gauss2D((7, 7), sigma=5)
EA = convolve(Et, weights=K, mode="constant", cval=0)
# MIN_E_EA = E;
# MIN_E_EA(GT & EA<E) = EA(GT & EA<E);
MIN_E_EA = np.where(gt & (EA < E), EA, E)
# %Pixel importance
B = np.where(gt == 0, 2 - np.exp(np.log(0.5) / 5 * Dst), np.ones_like(gt))
Ew = MIN_E_EA * B
TPw = np.sum(gt) - np.sum(Ew[gt == 1])
FPw = np.sum(Ew[gt == 0])
R = 1 - np.mean(Ew[gt == 1])
P = TPw / (TPw + FPw + _EPS)
# % Q = (1+Beta^2)*(R*P)./(eps+R+(Beta.*P));
Q = (1 + self.beta) * R * P / (R + self.beta * P + _EPS)
return Q
def matlab_style_gauss2D(self, shape: tuple = (7, 7), sigma: int = 5) -> np.ndarray:
"""
2D gaussian mask - should give the same result as MATLAB's
fspecial('gaussian',[shape],[sigma])
"""
m, n = [(ss - 1) / 2 for ss in shape]
y, x = np.ogrid[-m: m + 1, -n: n + 1]
h = np.exp(-(x * x + y * y) / (2 * sigma * sigma))
h[h < np.finfo(h.dtype).eps * h.max()] = 0
sumh = h.sum()
if sumh != 0:
h /= sumh
return h
def get_results(self) -> dict:
weighted_fm = np.mean(np.array(self.weighted_fms, dtype=_TYPE))
return dict(wfm=weighted_fm)
class HCEMeasure(object):
def __init__(self):
self.hces = []
def step(self, pred: np.ndarray, gt: np.ndarray, gt_ske):
# pred, gt = _prepare_data(pred, gt)
hce = self.cal_hce(pred, gt, gt_ske)
self.hces.append(hce)
def get_results(self) -> dict:
hce = np.mean(np.array(self.hces, _TYPE))
return dict(hce=hce)
def cal_hce(self, pred: np.ndarray, gt: np.ndarray, gt_ske: np.ndarray, relax=5, epsilon=2.0) -> float:
# Binarize gt
if(len(gt.shape)>2):
gt = gt[:, :, 0]
epsilon_gt = 128#(np.amin(gt)+np.amax(gt))/2.0
gt = (gt>epsilon_gt).astype(np.uint8)
# Binarize pred
if(len(pred.shape)>2):
pred = pred[:, :, 0]
epsilon_pred = 128#(np.amin(pred)+np.amax(pred))/2.0
pred = (pred>epsilon_pred).astype(np.uint8)
Union = np.logical_or(gt, pred)
TP = np.logical_and(gt, pred)
FP = pred - TP
FN = gt - TP
# relax the Union of gt and pred
Union_erode = Union.copy()
Union_erode = cv2.erode(Union_erode.astype(np.uint8), disk(1), iterations=relax)
# --- get the relaxed False Positive regions for computing the human efforts in correcting them ---
FP_ = np.logical_and(FP, Union_erode) # get the relaxed FP
for i in range(0, relax):
FP_ = cv2.dilate(FP_.astype(np.uint8), disk(1))
FP_ = np.logical_and(FP_, 1-np.logical_or(TP, FN))
FP_ = np.logical_and(FP, FP_)
# --- get the relaxed False Negative regions for computing the human efforts in correcting them ---
FN_ = np.logical_and(FN, Union_erode) # preserve the structural components of FN
## recover the FN, where pixels are not close to the TP borders
for i in range(0, relax):
FN_ = cv2.dilate(FN_.astype(np.uint8), disk(1))
FN_ = np.logical_and(FN_, 1-np.logical_or(TP, FP))
FN_ = np.logical_and(FN, FN_)
FN_ = np.logical_or(FN_, np.logical_xor(gt_ske, np.logical_and(TP, gt_ske))) # preserve the structural components of FN
## 2. =============Find exact polygon control points and independent regions==============
## find contours from FP_
ctrs_FP, hier_FP = cv2.findContours(FP_.astype(np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)
## find control points and independent regions for human correction
bdies_FP, indep_cnt_FP = self.filter_bdy_cond(ctrs_FP, FP_, np.logical_or(TP,FN_))
## find contours from FN_
ctrs_FN, hier_FN = cv2.findContours(FN_.astype(np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)
## find control points and independent regions for human correction
bdies_FN, indep_cnt_FN = self.filter_bdy_cond(ctrs_FN, FN_, 1-np.logical_or(np.logical_or(TP, FP_), FN_))
poly_FP, poly_FP_len, poly_FP_point_cnt = self.approximate_RDP(bdies_FP, epsilon=epsilon)
poly_FN, poly_FN_len, poly_FN_point_cnt = self.approximate_RDP(bdies_FN, epsilon=epsilon)
# FP_points+FP_indep+FN_points+FN_indep
return poly_FP_point_cnt+indep_cnt_FP+poly_FN_point_cnt+indep_cnt_FN
def filter_bdy_cond(self, bdy_, mask, cond):
cond = cv2.dilate(cond.astype(np.uint8), disk(1))
labels = label(mask) # find the connected regions
lbls = np.unique(labels) # the indices of the connected regions
indep = np.ones(lbls.shape[0]) # the label of each connected regions
indep[0] = 0 # 0 indicate the background region
boundaries = []
h,w = cond.shape[0:2]
ind_map = np.zeros((h, w))
indep_cnt = 0
for i in range(0, len(bdy_)):
tmp_bdies = []
tmp_bdy = []
for j in range(0, bdy_[i].shape[0]):
r, c = bdy_[i][j,0,1],bdy_[i][j,0,0]
if(np.sum(cond[r, c])==0 or ind_map[r, c]!=0):
if(len(tmp_bdy)>0):
tmp_bdies.append(tmp_bdy)
tmp_bdy = []
continue
tmp_bdy.append([c, r])
ind_map[r, c] = ind_map[r, c] + 1
indep[labels[r, c]] = 0 # indicates part of the boundary of this region needs human correction
if(len(tmp_bdy)>0):
tmp_bdies.append(tmp_bdy)
# check if the first and the last boundaries are connected
# if yes, invert the first boundary and attach it after the last boundary
if(len(tmp_bdies)>1):
first_x, first_y = tmp_bdies[0][0]
last_x, last_y = tmp_bdies[-1][-1]
if((abs(first_x-last_x)==1 and first_y==last_y) or
(first_x==last_x and abs(first_y-last_y)==1) or
(abs(first_x-last_x)==1 and abs(first_y-last_y)==1)
):
tmp_bdies[-1].extend(tmp_bdies[0][::-1])
del tmp_bdies[0]
for k in range(0, len(tmp_bdies)):
tmp_bdies[k] = np.array(tmp_bdies[k])[:, np.newaxis, :]
if(len(tmp_bdies)>0):
boundaries.extend(tmp_bdies)
return boundaries, np.sum(indep)
# this function approximate each boundary by DP algorithm
# https://en.wikipedia.org/wiki/Ramer%E2%80%93Douglas%E2%80%93Peucker_algorithm
def approximate_RDP(self, boundaries, epsilon=1.0):
boundaries_ = []
boundaries_len_ = []
pixel_cnt_ = 0
# polygon approximate of each boundary
for i in range(0, len(boundaries)):
boundaries_.append(cv2.approxPolyDP(boundaries[i], epsilon, False))
# count the control points number of each boundary and the total control points number of all the boundaries
for i in range(0, len(boundaries_)):
boundaries_len_.append(len(boundaries_[i]))
pixel_cnt_ = pixel_cnt_ + len(boundaries_[i])
return boundaries_, boundaries_len_, pixel_cnt_
class MBAMeasure(object):
def __init__(self):
self.bas = []
self.all_h = 0
self.all_w = 0
self.all_max = 0
def step(self, pred: np.ndarray, gt: np.ndarray):
# pred, gt = _prepare_data(pred, gt)
refined = gt.copy()
rmin = cmin = 0
rmax, cmax = gt.shape
self.all_h += rmax
self.all_w += cmax
self.all_max += max(rmax, cmax)
refined_h, refined_w = refined.shape
if refined_h != cmax:
refined = np.array(Image.fromarray(pred).resize((cmax, rmax), Image.BILINEAR))
if not(gt.sum() < 32*32):
if not((cmax==cmin) or (rmax==rmin)):
class_refined_prob = np.array(Image.fromarray(pred).resize((cmax-cmin, rmax-rmin), Image.BILINEAR))
refined[rmin:rmax, cmin:cmax] = class_refined_prob
pred = pred > 128
gt = gt > 128
ba = self.cal_ba(pred, gt)
self.bas.append(ba)
def get_disk_kernel(self, radius):
return cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (radius*2+1, radius*2+1))
def cal_ba(self, pred: np.ndarray, gt: np.ndarray) -> np.ndarray:
"""
Calculate the mean absolute error.
:return: ba
"""
gt = gt.astype(np.uint8)
pred = pred.astype(np.uint8)
h, w = gt.shape
min_radius = 1
max_radius = (w+h)/300
num_steps = 5
pred_acc = [None] * num_steps
for i in range(num_steps):
curr_radius = min_radius + int((max_radius-min_radius)/num_steps*i)
kernel = self.get_disk_kernel(curr_radius)
boundary_region = cv2.morphologyEx(gt, cv2.MORPH_GRADIENT, kernel) > 0
gt_in_bound = gt[boundary_region]
pred_in_bound = pred[boundary_region]
num_edge_pixels = (boundary_region).sum()
num_pred_gd_pix = ((gt_in_bound) * (pred_in_bound) + (1-gt_in_bound) * (1-pred_in_bound)).sum()
pred_acc[i] = num_pred_gd_pix / num_edge_pixels
ba = sum(pred_acc)/num_steps
return ba
def get_results(self) -> dict:
mba = np.mean(np.array(self.bas, _TYPE))
return dict(mba=mba)
class BIoUMeasure(object):
def __init__(self, dilation_ratio=0.02):
self.bious = []
self.dilation_ratio = dilation_ratio
def mask_to_boundary(self, mask):
h, w = mask.shape
img_diag = np.sqrt(h ** 2 + w ** 2)
dilation = int(round(self.dilation_ratio * img_diag))
if dilation < 1:
dilation = 1
# Pad image so mask truncated by the image border is also considered as boundary.
new_mask = cv2.copyMakeBorder(mask, 1, 1, 1, 1, cv2.BORDER_CONSTANT, value=0)
kernel = np.ones((3, 3), dtype=np.uint8)
new_mask_erode = cv2.erode(new_mask, kernel, iterations=dilation)
mask_erode = new_mask_erode[1 : h + 1, 1 : w + 1]
# G_d intersects G in the paper.
return mask - mask_erode
def step(self, pred: np.ndarray, gt: np.ndarray):
pred, gt = _prepare_data(pred, gt)
bious = self.cal_biou(pred=pred, gt=gt)
self.bious.append(bious)
def cal_biou(self, pred, gt):
pred = (pred * 255).astype(np.uint8)
pred = self.mask_to_boundary(pred)
gt = (gt * 255).astype(np.uint8)
gt = self.mask_to_boundary(gt)
gt = gt > 128
bins = np.linspace(0, 256, 257)
fg_hist, _ = np.histogram(pred[gt], bins=bins) # ture positive
bg_hist, _ = np.histogram(pred[~gt], bins=bins) # false positive
fg_w_thrs = np.cumsum(np.flip(fg_hist), axis=0)
bg_w_thrs = np.cumsum(np.flip(bg_hist), axis=0)
TPs = fg_w_thrs
Ps = fg_w_thrs + bg_w_thrs # positives
Ps[Ps == 0] = 1
T = max(np.count_nonzero(gt), 1)
ious = TPs / (T + bg_w_thrs)
return ious
def get_results(self) -> dict:
biou = np.mean(np.array(self.bious, dtype=_TYPE), axis=0)
return dict(biou=dict(curve=biou))
py/BiRefNet/gen_best_ep.py
+86
View File
@@ -0,0 +1,86 @@
import os
from glob import glob
import numpy as np
from BiRefNet..config import Config
config = Config()
eval_txts = sorted(glob('e_results/*_eval.txt'))
print('eval_txts:', [_.split(os.sep)[-1] for _ in eval_txts])
score_panel = {}
sep = '&'
metrics = ['sm', 'wfm', 'hce'] # we used HCE for DIS and wFm for others.
if 'DIS5K' not in config.task:
metrics.remove('hce')
for metric in metrics:
print('Metric:', metric)
current_line_nums = []
for idx_et, eval_txt in enumerate(eval_txts):
with open(eval_txt, 'r') as f:
lines = [l for l in f.readlines()[3:] if '.' in l]
current_line_nums.append(len(lines))
for idx_et, eval_txt in enumerate(eval_txts):
with open(eval_txt, 'r') as f:
lines = [l for l in f.readlines()[3:] if '.' in l]
for idx_line, line in enumerate(lines[:min(current_line_nums)]): # Consist line numbers by the minimal result file.
properties = line.strip().strip(sep).split(sep)
dataset = properties[0].strip()
ckpt = properties[1].strip()
if int(ckpt.split('--epoch_')[-1].strip()) < 0:
continue
targe_idx = {
'sm': [5, 2, 2, 5, 2],
'wfm': [3, 3, 8, 3, 8],
'hce': [7, -1, -1, 7, -1]
}[metric][['DIS5K', 'COD', 'HRSOD', 'General', 'Matting'].index(config.task)]
if metric != 'hce':
score_sm = float(properties[targe_idx].strip())
else:
score_sm = int(properties[targe_idx].strip().strip('.'))
if idx_et == 0:
score_panel[ckpt] = []
score_panel[ckpt].append(score_sm)
metrics_min = ['hce', 'mae']
max_or_min = min if metric in metrics_min else max
score_max = max_or_min(score_panel.values(), key=lambda x: np.sum(x))
good_models = []
for k, v in score_panel.items():
if (np.sum(v) <= np.sum(score_max)) if metric in metrics_min else (np.sum(v) >= np.sum(score_max)):
print(k, v)
good_models.append(k)
# Write
with open(eval_txt, 'r') as f:
lines = f.readlines()
info4good_models = lines[:3]
metric_names = [m.strip() for m in lines[1].strip().strip('&').split('&')[2:]]
testset_mean_values = {metric_name: [] for metric_name in metric_names}
for good_model in good_models:
for idx_et, eval_txt in enumerate(eval_txts):
with open(eval_txt, 'r') as f:
lines = f.readlines()
for line in lines:
if set([good_model]) & set([_.strip() for _ in line.split(sep)]):
info4good_models.append(line)
metric_scores = [float(m.strip()) for m in line.strip().strip('&').split('&')[2:]]
for idx_score, metric_score in enumerate(metric_scores):
testset_mean_values[metric_names[idx_score]].append(metric_score)
if 'DIS5K' in config.task:
testset_mean_values_lst = ['{:<4}'.format(int(np.mean(v_lst[:-1]).round())) if name == 'HCE' else '{:.3f}'.format(np.mean(v_lst[:-1])).lstrip('0') for name, v_lst in testset_mean_values.items()] # [:-1] to remove DIS-VD
sample_line_for_placing_mean_values = info4good_models[-2]
numbers_placed_well = sample_line_for_placing_mean_values.replace(sample_line_for_placing_mean_values.split('&')[1].strip(), 'DIS-TEs').strip().split('&')[3:]
for idx_number, (number_placed_well, testset_mean_value) in enumerate(zip(numbers_placed_well, testset_mean_values_lst)):
numbers_placed_well[idx_number] = number_placed_well.replace(number_placed_well.strip(), testset_mean_value)
testset_mean_line = '&'.join(sample_line_for_placing_mean_values.replace(sample_line_for_placing_mean_values.split('&')[1].strip(), 'DIS-TEs').split('&')[:3] + numbers_placed_well) + '\n'
info4good_models.append(testset_mean_line)
info4good_models.append(lines[-1])
info = ''.join(info4good_models)
print(info)
with open(os.path.join('e_results', 'eval-{}_best_on_{}.txt'.format(config.task, metric)), 'w') as f:
f.write(info + '\n')
py/BiRefNet/image_proc.py
+119
View File
@@ -0,0 +1,119 @@
import random
from PIL import Image, ImageEnhance
import numpy as np
import cv2
def refine_foreground(image, mask, r=90):
if mask.size != image.size:
mask = mask.resize(image.size)
image = np.array(image) / 255.0
mask = np.array(mask) / 255.0
estimated_foreground = FB_blur_fusion_foreground_estimator_2(image, mask, r=r)
image_masked = Image.fromarray((estimated_foreground * 255.0).astype(np.uint8))
return image_masked
def FB_blur_fusion_foreground_estimator_2(image, alpha, r=90):
# Thanks to the source: https://github.com/Photoroom/fast-foreground-estimation
alpha = alpha[:, :, None]
F, blur_B = FB_blur_fusion_foreground_estimator(
image, image, image, alpha, r)
return FB_blur_fusion_foreground_estimator(image, F, blur_B, alpha, r=6)[0]
def FB_blur_fusion_foreground_estimator(image, F, B, alpha, r=90):
if isinstance(image, Image.Image):
image = np.array(image) / 255.0
blurred_alpha = cv2.blur(alpha, (r, r))[:, :, None]
blurred_FA = cv2.blur(F * alpha, (r, r))
blurred_F = blurred_FA / (blurred_alpha + 1e-5)
blurred_B1A = cv2.blur(B * (1 - alpha), (r, r))
blurred_B = blurred_B1A / ((1 - blurred_alpha) + 1e-5)
F = blurred_F + alpha * \
(image - alpha * blurred_F - (1 - alpha) * blurred_B)
F = np.clip(F, 0, 1)
return F, blurred_B
def preproc(image, label, preproc_methods=['flip']):
if 'flip' in preproc_methods:
image, label = cv_random_flip(image, label)
if 'crop' in preproc_methods:
image, label = random_crop(image, label)
if 'rotate' in preproc_methods:
image, label = random_rotate(image, label)
if 'enhance' in preproc_methods:
image = color_enhance(image)
if 'pepper' in preproc_methods:
label = random_pepper(label)
return image, label
def cv_random_flip(img, label):
if random.random() > 0.5:
img = img.transpose(Image.FLIP_LEFT_RIGHT)
label = label.transpose(Image.FLIP_LEFT_RIGHT)
return img, label
def random_crop(image, label):
border = 30
image_width = image.size[0]
image_height = image.size[1]
border = int(min(image_width, image_height) * 0.1)
crop_win_width = np.random.randint(image_width - border, image_width)
crop_win_height = np.random.randint(image_height - border, image_height)
random_region = (
(image_width - crop_win_width) >> 1, (image_height - crop_win_height) >> 1, (image_width + crop_win_width) >> 1,
(image_height + crop_win_height) >> 1)
return image.crop(random_region), label.crop(random_region)
def random_rotate(image, label, angle=15):
mode = Image.BICUBIC
if random.random() > 0.8:
random_angle = np.random.randint(-angle, angle)
image = image.rotate(random_angle, mode)
label = label.rotate(random_angle, mode)
return image, label
def color_enhance(image):
bright_intensity = random.randint(5, 15) / 10.0
image = ImageEnhance.Brightness(image).enhance(bright_intensity)
contrast_intensity = random.randint(5, 15) / 10.0
image = ImageEnhance.Contrast(image).enhance(contrast_intensity)
color_intensity = random.randint(0, 20) / 10.0
image = ImageEnhance.Color(image).enhance(color_intensity)
sharp_intensity = random.randint(0, 30) / 10.0
image = ImageEnhance.Sharpness(image).enhance(sharp_intensity)
return image
def random_gaussian(image, mean=0.1, sigma=0.35):
def gaussianNoisy(im, mean=mean, sigma=sigma):
for _i in range(len(im)):
im[_i] += random.gauss(mean, sigma)
return im
img = np.asarray(image)
width, height = img.shape
img = gaussianNoisy(img[:].flatten(), mean, sigma)
img = img.reshape([width, height])
return Image.fromarray(np.uint8(img))
def random_pepper(img, N=0.0015):
img = np.array(img)
noiseNum = int(N * img.shape[0] * img.shape[1])
for i in range(noiseNum):
randX = random.randint(0, img.shape[0] - 1)
randY = random.randint(0, img.shape[1] - 1)
if random.randint(0, 1) == 0:
img[randX, randY] = 0
else:
img[randX, randY] = 255
return Image.fromarray(img)
py/BiRefNet/inference.py
+105
View File
@@ -0,0 +1,105 @@
import os
import argparse
from glob import glob
from tqdm import tqdm
import cv2
import torch
from BiRefNet.dataset import MyData
from BiRefNet.models.birefnet import BiRefNet
from BiRefNet.utils import save_tensor_img, check_state_dict
from BiRefNet.config import Config
config = Config()
def inference(model, data_loader_test, pred_root, method, testset, device=0):
model_training = model.training
if model_training:
model.eval()
for batch in tqdm(data_loader_test, total=len(data_loader_test)) if 1 or config.verbose_eval else data_loader_test:
inputs = batch[0].to(device)
# gts = batch[1].to(device)
label_paths = batch[-1]
with torch.no_grad():
scaled_preds = model(inputs)[-1].sigmoid()
os.makedirs(os.path.join(pred_root, method, testset), exist_ok=True)
for idx_sample in range(scaled_preds.shape[0]):
res = torch.nn.functional.interpolate(
scaled_preds[idx_sample].unsqueeze(0),
size=cv2.imread(label_paths[idx_sample], cv2.IMREAD_GRAYSCALE).shape[:2],
mode='bilinear',
align_corners=True
)
save_tensor_img(res, os.path.join(os.path.join(pred_root, method, testset), label_paths[idx_sample].replace('\\', '/').split('/')[-1])) # test set dir + file name
if model_training:
model.train()
return None
def main(args):
# Init model
device = config.device
if args.ckpt_folder:
print('Testing with models in {}'.format(args.ckpt_folder))
else:
print('Testing with model {}'.format(args.ckpt))
if config.model == 'BiRefNet':
model = BiRefNet(bb_pretrained=False)
weights_lst = sorted(
glob(os.path.join(args.ckpt_folder, '*.pth')) if args.ckpt_folder else [args.ckpt],
key=lambda x: int(x.split('epoch_')[-1].split('.pth')[0]),
reverse=True
)
for testset in args.testsets.split('+'):
print('>>>> Testset: {}...'.format(testset))
data_loader_test = torch.utils.data.DataLoader(
dataset=MyData(testset, image_size=config.size, is_train=False),
batch_size=config.batch_size_valid, shuffle=False, num_workers=config.num_workers, pin_memory=True
)
for weights in weights_lst:
if int(weights.strip('.pth').split('epoch_')[-1]) % 1 != 0:
continue
print('\tInferencing {}...'.format(weights))
# model.load_state_dict(torch.load(weights, map_location='cpu'))
state_dict = torch.load(weights, map_location='cpu')
state_dict = check_state_dict(state_dict)
model.load_state_dict(state_dict)
model = model.to(device)
inference(
model, data_loader_test=data_loader_test, pred_root=args.pred_root,
method='--'.join([w.rstrip('.pth') for w in weights.split(os.sep)[-2:]]),
testset=testset, device=config.device
)
if __name__ == '__main__':
# Parameter from command line
parser = argparse.ArgumentParser(description='')
parser.add_argument('--ckpt', type=str, help='model folder')
parser.add_argument('--ckpt_folder', default=sorted(glob(os.path.join('ckpt', '*')))[-1], type=str, help='model folder')
parser.add_argument('--pred_root', default='e_preds', type=str, help='Output folder')
parser.add_argument('--testsets',
default={
'DIS5K': 'DIS-VD+DIS-TE1+DIS-TE2+DIS-TE3+DIS-TE4',
'COD': 'TE-COD10K+NC4K+TE-CAMO+CHAMELEON',
'HRSOD': 'DAVIS-S+TE-HRSOD+TE-UHRSD+TE-DUTS+DUT-OMRON',
'General': 'DIS-VD',
'Matting': 'TE-P3M-500-P',
'DIS5K-': 'DIS-VD',
'COD-': 'TE-COD10K',
'SOD-': 'DAVIS-S+TE-HRSOD+TE-UHRSD',
}[config.task + ''],
type=str,
help="Test all sets: , 'DIS-VD+DIS-TE1+DIS-TE2+DIS-TE3+DIS-TE4'")
args = parser.parse_args()
if config.precisionHigh:
torch.set_float32_matmul_precision('high')
main(args)
py/BiRefNet/loss.py
+277
View File
@@ -0,0 +1,277 @@
import torch
from torch import nn
import torch.nn.functional as F
from torch.autograd import Variable
from math import exp
from BiRefNet.config import Config
class Discriminator(nn.Module):
def __init__(self, channels=1, img_size=256):
super(Discriminator, self).__init__()
def discriminator_block(in_filters, out_filters, bn=Config().batch_size > 1):
block = [nn.Conv2d(in_filters, out_filters, 3, 2, 1), nn.LeakyReLU(0.2, inplace=True), nn.Dropout2d(0.25)]
if bn:
block.append(nn.BatchNorm2d(out_filters, 0.8))
return block
self.model = nn.Sequential(
*discriminator_block(channels, 16, bn=False),
*discriminator_block(16, 32),
*discriminator_block(32, 64),
*discriminator_block(64, 128),
)
# The height and width of downsampled image
ds_size = img_size // 2 ** 4
self.adv_layer = nn.Sequential(nn.Linear(128 * ds_size ** 2, 1), nn.Sigmoid())
def forward(self, img):
out = self.model(img)
out = out.view(out.shape[0], -1)
validity = self.adv_layer(out)
return validity
class ContourLoss(torch.nn.Module):
def __init__(self):
super(ContourLoss, self).__init__()
def forward(self, pred, target, weight=10):
'''
target, pred: tensor of shape (B, C, H, W), where target[:,:,region_in_contour] == 1,
target[:,:,region_out_contour] == 0.
weight: scalar, length term weight.
'''
# length term
delta_r = pred[:,:,1:,:] - pred[:,:,:-1,:] # horizontal gradient (B, C, H-1, W)
delta_c = pred[:,:,:,1:] - pred[:,:,:,:-1] # vertical gradient (B, C, H, W-1)
delta_r = delta_r[:,:,1:,:-2]**2 # (B, C, H-2, W-2)
delta_c = delta_c[:,:,:-2,1:]**2 # (B, C, H-2, W-2)
delta_pred = torch.abs(delta_r + delta_c)
epsilon = 1e-8 # where is a parameter to avoid square root is zero in practice.
length = torch.mean(torch.sqrt(delta_pred + epsilon)) # eq.(11) in the paper, mean is used instead of sum.
c_in = torch.ones_like(pred)
c_out = torch.zeros_like(pred)
region_in = torch.mean( pred * (target - c_in )**2 ) # equ.(12) in the paper, mean is used instead of sum.
region_out = torch.mean( (1-pred) * (target - c_out)**2 )
region = region_in + region_out
loss = weight * length + region
return loss
class IoULoss(torch.nn.Module):
def __init__(self):
super(IoULoss, self).__init__()
def forward(self, pred, target):
b = pred.shape[0]
IoU = 0.0
for i in range(0, b):
# compute the IoU of the foreground
Iand1 = torch.sum(target[i, :, :, :] * pred[i, :, :, :])
Ior1 = torch.sum(target[i, :, :, :]) + torch.sum(pred[i, :, :, :]) - Iand1
IoU1 = Iand1 / Ior1
# IoU loss is (1-IoU1)
IoU = IoU + (1-IoU1)
# return IoU/b
return IoU
class StructureLoss(torch.nn.Module):
def __init__(self):
super(StructureLoss, self).__init__()
def forward(self, pred, target):
weit = 1+5*torch.abs(F.avg_pool2d(target, kernel_size=31, stride=1, padding=15)-target)
wbce = F.binary_cross_entropy_with_logits(pred, target, reduction='none')
wbce = (weit*wbce).sum(dim=(2,3))/weit.sum(dim=(2,3))
pred = torch.sigmoid(pred)
inter = ((pred * target) * weit).sum(dim=(2, 3))
union = ((pred + target) * weit).sum(dim=(2, 3))
wiou = 1-(inter+1)/(union-inter+1)
return (wbce+wiou).mean()
class PatchIoULoss(torch.nn.Module):
def __init__(self):
super(PatchIoULoss, self).__init__()
self.iou_loss = IoULoss()
def forward(self, pred, target):
win_y, win_x = 64, 64
iou_loss = 0.
for anchor_y in range(0, target.shape[0], win_y):
for anchor_x in range(0, target.shape[1], win_y):
patch_pred = pred[:, :, anchor_y:anchor_y+win_y, anchor_x:anchor_x+win_x]
patch_target = target[:, :, anchor_y:anchor_y+win_y, anchor_x:anchor_x+win_x]
patch_iou_loss = self.iou_loss(patch_pred, patch_target)
iou_loss += patch_iou_loss
return iou_loss
class ThrReg_loss(torch.nn.Module):
def __init__(self):
super(ThrReg_loss, self).__init__()
def forward(self, pred, gt=None):
return torch.mean(1 - ((pred - 0) ** 2 + (pred - 1) ** 2))
class ClsLoss(nn.Module):
"""
Auxiliary classification loss for each refined class output.
"""
def __init__(self):
super(ClsLoss, self).__init__()
self.config = Config()
self.lambdas_cls = self.config.lambdas_cls
self.criterions_last = {
'ce': nn.CrossEntropyLoss()
}
def forward(self, preds, gt):
loss = 0.
for _, pred_lvl in enumerate(preds):
if pred_lvl is None:
continue
for criterion_name, criterion in self.criterions_last.items():
loss += criterion(pred_lvl, gt) * self.lambdas_cls[criterion_name]
return loss
class PixLoss(nn.Module):
"""
Pixel loss for each refined map output.
"""
def __init__(self):
super(PixLoss, self).__init__()
self.config = Config()
self.lambdas_pix_last = self.config.lambdas_pix_last
self.criterions_last = {}
if 'bce' in self.lambdas_pix_last and self.lambdas_pix_last['bce']:
self.criterions_last['bce'] = nn.BCELoss() if not self.config.use_fp16 else nn.BCEWithLogitsLoss()
if 'iou' in self.lambdas_pix_last and self.lambdas_pix_last['iou']:
self.criterions_last['iou'] = IoULoss()
if 'iou_patch' in self.lambdas_pix_last and self.lambdas_pix_last['iou_patch']:
self.criterions_last['iou_patch'] = PatchIoULoss()
if 'ssim' in self.lambdas_pix_last and self.lambdas_pix_last['ssim']:
self.criterions_last['ssim'] = SSIMLoss()
if 'mae' in self.lambdas_pix_last and self.lambdas_pix_last['mae']:
self.criterions_last['mae'] = nn.L1Loss()
if 'mse' in self.lambdas_pix_last and self.lambdas_pix_last['mse']:
self.criterions_last['mse'] = nn.MSELoss()
if 'reg' in self.lambdas_pix_last and self.lambdas_pix_last['reg']:
self.criterions_last['reg'] = ThrReg_loss()
if 'cnt' in self.lambdas_pix_last and self.lambdas_pix_last['cnt']:
self.criterions_last['cnt'] = ContourLoss()
if 'structure' in self.lambdas_pix_last and self.lambdas_pix_last['structure']:
self.criterions_last['structure'] = StructureLoss()
def forward(self, scaled_preds, gt):
loss = 0.
criterions_embedded_with_sigmoid = ['structure', ] + ['bce'] if self.config.use_fp16 else []
for _, pred_lvl in enumerate(scaled_preds):
if pred_lvl.shape != gt.shape:
pred_lvl = nn.functional.interpolate(pred_lvl, size=gt.shape[2:], mode='bilinear', align_corners=True)
for criterion_name, criterion in self.criterions_last.items():
_loss = criterion(pred_lvl.sigmoid() if criterion_name not in criterions_embedded_with_sigmoid else pred_lvl, gt) * self.lambdas_pix_last[criterion_name]
loss += _loss
# print(criterion_name, _loss.item())
return loss
class SSIMLoss(torch.nn.Module):
def __init__(self, window_size=11, size_average=True):
super(SSIMLoss, self).__init__()
self.window_size = window_size
self.size_average = size_average
self.channel = 1
self.window = create_window(window_size, self.channel)
def forward(self, img1, img2):
(_, channel, _, _) = img1.size()
if channel == self.channel and self.window.data.type() == img1.data.type():
window = self.window
else:
window = create_window(self.window_size, channel)
if img1.is_cuda:
window = window.cuda(img1.get_device())
window = window.type_as(img1)
self.window = window
self.channel = channel
return 1 - _ssim(img1, img2, window, self.window_size, channel, self.size_average)
def gaussian(window_size, sigma):
gauss = torch.Tensor([exp(-(x - window_size//2)**2/float(2*sigma**2)) for x in range(window_size)])
return gauss/gauss.sum()
def create_window(window_size, channel):
_1D_window = gaussian(window_size, 1.5).unsqueeze(1)
_2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0)
window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous())
return window
def _ssim(img1, img2, window, window_size, channel, size_average=True):
mu1 = F.conv2d(img1, window, padding = window_size//2, groups=channel)
mu2 = F.conv2d(img2, window, padding = window_size//2, groups=channel)
mu1_sq = mu1.pow(2)
mu2_sq = mu2.pow(2)
mu1_mu2 = mu1*mu2
sigma1_sq = F.conv2d(img1*img1, window, padding=window_size//2, groups=channel) - mu1_sq
sigma2_sq = F.conv2d(img2*img2, window, padding=window_size//2, groups=channel) - mu2_sq
sigma12 = F.conv2d(img1*img2, window, padding=window_size//2, groups=channel) - mu1_mu2
C1 = 0.01**2
C2 = 0.03**2
ssim_map = ((2*mu1_mu2 + C1)*(2*sigma12 + C2))/((mu1_sq + mu2_sq + C1)*(sigma1_sq + sigma2_sq + C2))
if size_average:
return ssim_map.mean()
else:
return ssim_map.mean(1).mean(1).mean(1)
def SSIM(x, y):
C1 = 0.01 ** 2
C2 = 0.03 ** 2
mu_x = nn.AvgPool2d(3, 1, 1)(x)
mu_y = nn.AvgPool2d(3, 1, 1)(y)
mu_x_mu_y = mu_x * mu_y
mu_x_sq = mu_x.pow(2)
mu_y_sq = mu_y.pow(2)
sigma_x = nn.AvgPool2d(3, 1, 1)(x * x) - mu_x_sq
sigma_y = nn.AvgPool2d(3, 1, 1)(y * y) - mu_y_sq
sigma_xy = nn.AvgPool2d(3, 1, 1)(x * y) - mu_x_mu_y
SSIM_n = (2 * mu_x_mu_y + C1) * (2 * sigma_xy + C2)
SSIM_d = (mu_x_sq + mu_y_sq + C1) * (sigma_x + sigma_y + C2)
SSIM = SSIM_n / SSIM_d
return torch.clamp((1 - SSIM) / 2, 0, 1)
def saliency_structure_consistency(x, y):
ssim = torch.mean(SSIM(x,y))
return ssim
py/BiRefNet/make_a_copy.sh
+18
View File
@@ -0,0 +1,18 @@
#!/bin/bash
# Set dst repo here.
repo=$1
mkdir ../${repo}
mkdir ../${repo}/evaluation
mkdir ../${repo}/models
mkdir ../${repo}/models/backbones
mkdir ../${repo}/models/modules
mkdir ../${repo}/models/refinement
cp ./*.sh ../${repo}
cp ./*.py ../${repo}
cp ./evaluation/*.py ../${repo}/evaluation
cp ./models/*.py ../${repo}/models
cp ./models/backbones/*.py ../${repo}/models/backbones
cp ./models/modules/*.py ../${repo}/models/modules
cp ./models/refinement/*.py ../${repo}/models/refinement
cp -r ./.git* ../${repo}
py/BiRefNet/models/backbones/build_backbone.py
+44
View File
@@ -0,0 +1,44 @@
import torch
import torch.nn as nn
from collections import OrderedDict
from torchvision.models import vgg16, vgg16_bn, VGG16_Weights, VGG16_BN_Weights, resnet50, ResNet50_Weights
from models.backbones.pvt_v2 import pvt_v2_b0, pvt_v2_b1, pvt_v2_b2, pvt_v2_b5
from models.backbones.swin_v1 import swin_v1_t, swin_v1_s, swin_v1_b, swin_v1_l
from config import Config
config = Config()
def build_backbone(bb_name, pretrained=True, params_settings=''):
if bb_name == 'vgg16':
bb_net = list(vgg16(pretrained=VGG16_Weights.DEFAULT if pretrained else None).children())[0]
bb = nn.Sequential(OrderedDict({'conv1': bb_net[:4], 'conv2': bb_net[4:9], 'conv3': bb_net[9:16], 'conv4': bb_net[16:23]}))
elif bb_name == 'vgg16bn':
bb_net = list(vgg16_bn(pretrained=VGG16_BN_Weights.DEFAULT if pretrained else None).children())[0]
bb = nn.Sequential(OrderedDict({'conv1': bb_net[:6], 'conv2': bb_net[6:13], 'conv3': bb_net[13:23], 'conv4': bb_net[23:33]}))
elif bb_name == 'resnet50':
bb_net = list(resnet50(pretrained=ResNet50_Weights.DEFAULT if pretrained else None).children())
bb = nn.Sequential(OrderedDict({'conv1': nn.Sequential(*bb_net[0:3]), 'conv2': bb_net[4], 'conv3': bb_net[5], 'conv4': bb_net[6]}))
else:
bb = eval('{}({})'.format(bb_name, params_settings))
if pretrained:
bb = load_weights(bb, bb_name)
return bb
def load_weights(model, model_name):
save_model = torch.load(config.weights[model_name], map_location='cpu')
model_dict = model.state_dict()
state_dict = {k: v if v.size() == model_dict[k].size() else model_dict[k] for k, v in save_model.items() if k in model_dict.keys()}
# to ignore the weights with mismatched size when I modify the backbone itself.
if not state_dict:
save_model_keys = list(save_model.keys())
sub_item = save_model_keys[0] if len(save_model_keys) == 1 else None
state_dict = {k: v if v.size() == model_dict[k].size() else model_dict[k] for k, v in save_model[sub_item].items() if k in model_dict.keys()}
if not state_dict or not sub_item:
print('Weights are not successully loaded. Check the state dict of weights file.')
return None
else:
print('Found correct weights in the "{}" item of loaded state_dict.'.format(sub_item))
model_dict.update(state_dict)
model.load_state_dict(model_dict)
return model
py/BiRefNet/models/backbones/pvt_v2.py
+435
View File
@@ -0,0 +1,435 @@
import torch
import torch.nn as nn
from functools import partial
from timm.models.layers import DropPath, to_2tuple, trunc_normal_
from timm.models.registry import register_model
import math
from config import Config
config = Config()
class Mlp(nn.Module):
def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
super().__init__()
out_features = out_features or in_features
hidden_features = hidden_features or in_features
self.fc1 = nn.Linear(in_features, hidden_features)
self.dwconv = DWConv(hidden_features)
self.act = act_layer()
self.fc2 = nn.Linear(hidden_features, out_features)
self.drop = nn.Dropout(drop)
self.apply(self._init_weights)
def _init_weights(self, m):
if isinstance(m, nn.Linear):
trunc_normal_(m.weight, std=.02)
if isinstance(m, nn.Linear) and m.bias is not None:
nn.init.constant_(m.bias, 0)
elif isinstance(m, nn.LayerNorm):
nn.init.constant_(m.bias, 0)
nn.init.constant_(m.weight, 1.0)
elif isinstance(m, nn.Conv2d):
fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
fan_out //= m.groups
m.weight.data.normal_(0, math.sqrt(2.0 / fan_out))
if m.bias is not None:
m.bias.data.zero_()
def forward(self, x, H, W):
x = self.fc1(x)
x = self.dwconv(x, H, W)
x = self.act(x)
x = self.drop(x)
x = self.fc2(x)
x = self.drop(x)
return x
class Attention(nn.Module):
def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0., sr_ratio=1):
super().__init__()
assert dim % num_heads == 0, f"dim {dim} should be divided by num_heads {num_heads}."
self.dim = dim
self.num_heads = num_heads
head_dim = dim // num_heads
self.scale = qk_scale or head_dim ** -0.5
self.q = nn.Linear(dim, dim, bias=qkv_bias)
self.kv = nn.Linear(dim, dim * 2, bias=qkv_bias)
self.attn_drop_prob = attn_drop
self.attn_drop = nn.Dropout(attn_drop)
self.proj = nn.Linear(dim, dim)
self.proj_drop = nn.Dropout(proj_drop)
self.sr_ratio = sr_ratio
if sr_ratio > 1:
self.sr = nn.Conv2d(dim, dim, kernel_size=sr_ratio, stride=sr_ratio)
self.norm = nn.LayerNorm(dim)
self.apply(self._init_weights)
def _init_weights(self, m):
if isinstance(m, nn.Linear):
trunc_normal_(m.weight, std=.02)
if isinstance(m, nn.Linear) and m.bias is not None:
nn.init.constant_(m.bias, 0)
elif isinstance(m, nn.LayerNorm):
nn.init.constant_(m.bias, 0)
nn.init.constant_(m.weight, 1.0)
elif isinstance(m, nn.Conv2d):
fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
fan_out //= m.groups
m.weight.data.normal_(0, math.sqrt(2.0 / fan_out))
if m.bias is not None:
m.bias.data.zero_()
def forward(self, x, H, W):
B, N, C = x.shape
q = self.q(x).reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)
if self.sr_ratio > 1:
x_ = x.permute(0, 2, 1).reshape(B, C, H, W)
x_ = self.sr(x_).reshape(B, C, -1).permute(0, 2, 1)
x_ = self.norm(x_)
kv = self.kv(x_).reshape(B, -1, 2, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
else:
kv = self.kv(x).reshape(B, -1, 2, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
k, v = kv[0], kv[1]
if config.SDPA_enabled:
x = torch.nn.functional.scaled_dot_product_attention(
q, k, v,
attn_mask=None, dropout_p=self.attn_drop_prob, is_causal=False
).transpose(1, 2).reshape(B, N, C)
else:
attn = (q @ k.transpose(-2, -1)) * self.scale
attn = attn.softmax(dim=-1)
attn = self.attn_drop(attn)
x = (attn @ v).transpose(1, 2).reshape(B, N, C)
x = self.proj(x)
x = self.proj_drop(x)
return x
class Block(nn.Module):
def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, sr_ratio=1):
super().__init__()
self.norm1 = norm_layer(dim)
self.attn = Attention(
dim,
num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale,
attn_drop=attn_drop, proj_drop=drop, sr_ratio=sr_ratio)
# NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
self.norm2 = norm_layer(dim)
mlp_hidden_dim = int(dim * mlp_ratio)
self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
self.apply(self._init_weights)
def _init_weights(self, m):
if isinstance(m, nn.Linear):
trunc_normal_(m.weight, std=.02)
if isinstance(m, nn.Linear) and m.bias is not None:
nn.init.constant_(m.bias, 0)
elif isinstance(m, nn.LayerNorm):
nn.init.constant_(m.bias, 0)
nn.init.constant_(m.weight, 1.0)
elif isinstance(m, nn.Conv2d):
fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
fan_out //= m.groups
m.weight.data.normal_(0, math.sqrt(2.0 / fan_out))
if m.bias is not None:
m.bias.data.zero_()
def forward(self, x, H, W):
x = x + self.drop_path(self.attn(self.norm1(x), H, W))
x = x + self.drop_path(self.mlp(self.norm2(x), H, W))
return x
class OverlapPatchEmbed(nn.Module):
""" Image to Patch Embedding
"""
def __init__(self, img_size=224, patch_size=7, stride=4, in_channels=3, embed_dim=768):
super().__init__()
img_size = to_2tuple(img_size)
patch_size = to_2tuple(patch_size)
self.img_size = img_size
self.patch_size = patch_size
self.H, self.W = img_size[0] // patch_size[0], img_size[1] // patch_size[1]
self.num_patches = self.H * self.W
self.proj = nn.Conv2d(in_channels, embed_dim, kernel_size=patch_size, stride=stride,
padding=(patch_size[0] // 2, patch_size[1] // 2))
self.norm = nn.LayerNorm(embed_dim)
self.apply(self._init_weights)
def _init_weights(self, m):
if isinstance(m, nn.Linear):
trunc_normal_(m.weight, std=.02)
if isinstance(m, nn.Linear) and m.bias is not None:
nn.init.constant_(m.bias, 0)
elif isinstance(m, nn.LayerNorm):
nn.init.constant_(m.bias, 0)
nn.init.constant_(m.weight, 1.0)
elif isinstance(m, nn.Conv2d):
fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
fan_out //= m.groups
m.weight.data.normal_(0, math.sqrt(2.0 / fan_out))
if m.bias is not None:
m.bias.data.zero_()
def forward(self, x):
x = self.proj(x)
_, _, H, W = x.shape
x = x.flatten(2).transpose(1, 2)
x = self.norm(x)
return x, H, W
class PyramidVisionTransformerImpr(nn.Module):
def __init__(self, img_size=224, patch_size=16, in_channels=3, num_classes=1000, embed_dims=[64, 128, 256, 512],
num_heads=[1, 2, 4, 8], mlp_ratios=[4, 4, 4, 4], qkv_bias=False, qk_scale=None, drop_rate=0.,
attn_drop_rate=0., drop_path_rate=0., norm_layer=nn.LayerNorm,
depths=[3, 4, 6, 3], sr_ratios=[8, 4, 2, 1]):
super().__init__()
self.num_classes = num_classes
self.depths = depths
# patch_embed
self.patch_embed1 = OverlapPatchEmbed(img_size=img_size, patch_size=7, stride=4, in_channels=in_channels,
embed_dim=embed_dims[0])
self.patch_embed2 = OverlapPatchEmbed(img_size=img_size // 4, patch_size=3, stride=2, in_channels=embed_dims[0],
embed_dim=embed_dims[1])
self.patch_embed3 = OverlapPatchEmbed(img_size=img_size // 8, patch_size=3, stride=2, in_channels=embed_dims[1],
embed_dim=embed_dims[2])
self.patch_embed4 = OverlapPatchEmbed(img_size=img_size // 16, patch_size=3, stride=2, in_channels=embed_dims[2],
embed_dim=embed_dims[3])
# transformer encoder
dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] # stochastic depth decay rule
cur = 0
self.block1 = nn.ModuleList([Block(
dim=embed_dims[0], num_heads=num_heads[0], mlp_ratio=mlp_ratios[0], qkv_bias=qkv_bias, qk_scale=qk_scale,
drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[cur + i], norm_layer=norm_layer,
sr_ratio=sr_ratios[0])
for i in range(depths[0])])
self.norm1 = norm_layer(embed_dims[0])
cur += depths[0]
self.block2 = nn.ModuleList([Block(
dim=embed_dims[1], num_heads=num_heads[1], mlp_ratio=mlp_ratios[1], qkv_bias=qkv_bias, qk_scale=qk_scale,
drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[cur + i], norm_layer=norm_layer,
sr_ratio=sr_ratios[1])
for i in range(depths[1])])
self.norm2 = norm_layer(embed_dims[1])
cur += depths[1]
self.block3 = nn.ModuleList([Block(
dim=embed_dims[2], num_heads=num_heads[2], mlp_ratio=mlp_ratios[2], qkv_bias=qkv_bias, qk_scale=qk_scale,
drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[cur + i], norm_layer=norm_layer,
sr_ratio=sr_ratios[2])
for i in range(depths[2])])
self.norm3 = norm_layer(embed_dims[2])
cur += depths[2]
self.block4 = nn.ModuleList([Block(
dim=embed_dims[3], num_heads=num_heads[3], mlp_ratio=mlp_ratios[3], qkv_bias=qkv_bias, qk_scale=qk_scale,
drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[cur + i], norm_layer=norm_layer,
sr_ratio=sr_ratios[3])
for i in range(depths[3])])
self.norm4 = norm_layer(embed_dims[3])
# classification head
# self.head = nn.Linear(embed_dims[3], num_classes) if num_classes > 0 else nn.Identity()
self.apply(self._init_weights)
def _init_weights(self, m):
if isinstance(m, nn.Linear):
trunc_normal_(m.weight, std=.02)
if isinstance(m, nn.Linear) and m.bias is not None:
nn.init.constant_(m.bias, 0)
elif isinstance(m, nn.LayerNorm):
nn.init.constant_(m.bias, 0)
nn.init.constant_(m.weight, 1.0)
elif isinstance(m, nn.Conv2d):
fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
fan_out //= m.groups
m.weight.data.normal_(0, math.sqrt(2.0 / fan_out))
if m.bias is not None:
m.bias.data.zero_()
def init_weights(self, pretrained=None):
if isinstance(pretrained, str):
logger = 1
#load_checkpoint(self, pretrained, map_location='cpu', strict=False, logger=logger)
def reset_drop_path(self, drop_path_rate):
dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(self.depths))]
cur = 0
for i in range(self.depths[0]):
self.block1[i].drop_path.drop_prob = dpr[cur + i]
cur += self.depths[0]
for i in range(self.depths[1]):
self.block2[i].drop_path.drop_prob = dpr[cur + i]
cur += self.depths[1]
for i in range(self.depths[2]):
self.block3[i].drop_path.drop_prob = dpr[cur + i]
cur += self.depths[2]
for i in range(self.depths[3]):
self.block4[i].drop_path.drop_prob = dpr[cur + i]
def freeze_patch_emb(self):
self.patch_embed1.requires_grad = False
@torch.jit.ignore
def no_weight_decay(self):
return {'pos_embed1', 'pos_embed2', 'pos_embed3', 'pos_embed4', 'cls_token'} # has pos_embed may be better
def get_classifier(self):
return self.head
def reset_classifier(self, num_classes, global_pool=''):
self.num_classes = num_classes
self.head = nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity()
def forward_features(self, x):
B = x.shape[0]
outs = []
# stage 1
x, H, W = self.patch_embed1(x)
for i, blk in enumerate(self.block1):
x = blk(x, H, W)
x = self.norm1(x)
x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()
outs.append(x)
# stage 2
x, H, W = self.patch_embed2(x)
for i, blk in enumerate(self.block2):
x = blk(x, H, W)
x = self.norm2(x)
x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()
outs.append(x)
# stage 3
x, H, W = self.patch_embed3(x)
for i, blk in enumerate(self.block3):
x = blk(x, H, W)
x = self.norm3(x)
x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()
outs.append(x)
# stage 4
x, H, W = self.patch_embed4(x)
for i, blk in enumerate(self.block4):
x = blk(x, H, W)
x = self.norm4(x)
x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()
outs.append(x)
return outs
# return x.mean(dim=1)
def forward(self, x):
x = self.forward_features(x)
# x = self.head(x)
return x
class DWConv(nn.Module):
def __init__(self, dim=768):
super(DWConv, self).__init__()
self.dwconv = nn.Conv2d(dim, dim, 3, 1, 1, bias=True, groups=dim)
def forward(self, x, H, W):
B, N, C = x.shape
x = x.transpose(1, 2).view(B, C, H, W).contiguous()
x = self.dwconv(x)
x = x.flatten(2).transpose(1, 2)
return x
def _conv_filter(state_dict, patch_size=16):
""" convert patch embedding weight from manual patchify + linear proj to conv"""
out_dict = {}
for k, v in state_dict.items():
if 'patch_embed.proj.weight' in k:
v = v.reshape((v.shape[0], 3, patch_size, patch_size))
out_dict[k] = v
return out_dict
## @register_model
class pvt_v2_b0(PyramidVisionTransformerImpr):
def __init__(self, **kwargs):
super(pvt_v2_b0, self).__init__(
patch_size=4, embed_dims=[32, 64, 160, 256], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4],
qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[2, 2, 2, 2], sr_ratios=[8, 4, 2, 1],
drop_rate=0.0, drop_path_rate=0.1)
## @register_model
class pvt_v2_b1(PyramidVisionTransformerImpr):
def __init__(self, **kwargs):
super(pvt_v2_b1, self).__init__(
patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4],
qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[2, 2, 2, 2], sr_ratios=[8, 4, 2, 1],
drop_rate=0.0, drop_path_rate=0.1)
## @register_model
class pvt_v2_b2(PyramidVisionTransformerImpr):
def __init__(self, in_channels=3, **kwargs):
super(pvt_v2_b2, self).__init__(
patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4],
qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[3, 4, 6, 3], sr_ratios=[8, 4, 2, 1],
drop_rate=0.0, drop_path_rate=0.1, in_channels=in_channels)
## @register_model
class pvt_v2_b3(PyramidVisionTransformerImpr):
def __init__(self, **kwargs):
super(pvt_v2_b3, self).__init__(
patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4],
qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[3, 4, 18, 3], sr_ratios=[8, 4, 2, 1],
drop_rate=0.0, drop_path_rate=0.1)
## @register_model
class pvt_v2_b4(PyramidVisionTransformerImpr):
def __init__(self, **kwargs):
super(pvt_v2_b4, self).__init__(
patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4],
qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[3, 8, 27, 3], sr_ratios=[8, 4, 2, 1],
drop_rate=0.0, drop_path_rate=0.1)
## @register_model
class pvt_v2_b5(PyramidVisionTransformerImpr):
def __init__(self, **kwargs):
super(pvt_v2_b5, self).__init__(
patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[4, 4, 4, 4],
qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[3, 6, 40, 3], sr_ratios=[8, 4, 2, 1],
drop_rate=0.0, drop_path_rate=0.1)
py/BiRefNet/models/backbones/swin_v1.py
+627
View File
@@ -0,0 +1,627 @@
# --------------------------------------------------------
# Swin Transformer
# Copyright (c) 2021 Microsoft
# Licensed under The MIT License [see LICENSE for details]
# Written by Ze Liu, Yutong Lin, Yixuan Wei
# --------------------------------------------------------
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.checkpoint as checkpoint
import numpy as np
from timm.models.layers import DropPath, to_2tuple, trunc_normal_
from config import Config
config = Config()
class Mlp(nn.Module):
""" Multilayer perceptron."""
def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
super().__init__()
out_features = out_features or in_features
hidden_features = hidden_features or in_features
self.fc1 = nn.Linear(in_features, hidden_features)
self.act = act_layer()
self.fc2 = nn.Linear(hidden_features, out_features)
self.drop = nn.Dropout(drop)
def forward(self, x):
x = self.fc1(x)
x = self.act(x)
x = self.drop(x)
x = self.fc2(x)
x = self.drop(x)
return x
def window_partition(x, window_size):
"""
Args:
x: (B, H, W, C)
window_size (int): window size
Returns:
windows: (num_windows*B, window_size, window_size, C)
"""
B, H, W, C = x.shape
x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)
return windows
def window_reverse(windows, window_size, H, W):
"""
Args:
windows: (num_windows*B, window_size, window_size, C)
window_size (int): Window size
H (int): Height of image
W (int): Width of image
Returns:
x: (B, H, W, C)
"""
B = int(windows.shape[0] / (H * W / window_size / window_size))
x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1)
x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
return x
class WindowAttention(nn.Module):
""" Window based multi-head self attention (W-MSA) module with relative position bias.
It supports both of shifted and non-shifted window.
Args:
dim (int): Number of input channels.
window_size (tuple[int]): The height and width of the window.
num_heads (int): Number of attention heads.
qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set
attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0
proj_drop (float, optional): Dropout ratio of output. Default: 0.0
"""
def __init__(self, dim, window_size, num_heads, qkv_bias=True, qk_scale=None, attn_drop=0., proj_drop=0.):
super().__init__()
self.dim = dim
self.window_size = window_size # Wh, Ww
self.num_heads = num_heads
head_dim = dim // num_heads
self.scale = qk_scale or head_dim ** -0.5
# define a parameter table of relative position bias
self.relative_position_bias_table = nn.Parameter(
torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads)) # 2*Wh-1 * 2*Ww-1, nH
# get pair-wise relative position index for each token inside the window
coords_h = torch.arange(self.window_size[0])
coords_w = torch.arange(self.window_size[1])
coords = torch.stack(torch.meshgrid([coords_h, coords_w], indexing='ij')) # 2, Wh, Ww
coords_flatten = torch.flatten(coords, 1) # 2, Wh*Ww
relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] # 2, Wh*Ww, Wh*Ww
relative_coords = relative_coords.permute(1, 2, 0).contiguous() # Wh*Ww, Wh*Ww, 2
relative_coords[:, :, 0] += self.window_size[0] - 1 # shift to start from 0
relative_coords[:, :, 1] += self.window_size[1] - 1
relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1
relative_position_index = relative_coords.sum(-1) # Wh*Ww, Wh*Ww
self.register_buffer("relative_position_index", relative_position_index)
self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
self.attn_drop_prob = attn_drop
self.attn_drop = nn.Dropout(attn_drop)
self.proj = nn.Linear(dim, dim)
self.proj_drop = nn.Dropout(proj_drop)
trunc_normal_(self.relative_position_bias_table, std=.02)
self.softmax = nn.Softmax(dim=-1)
def forward(self, x, mask=None):
""" Forward function.
Args:
x: input features with shape of (num_windows*B, N, C)
mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None
"""
B_, N, C = x.shape
qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
q, k, v = qkv[0], qkv[1], qkv[2] # make torchscript happy (cannot use tensor as tuple)
q = q * self.scale
if config.SDPA_enabled:
x = torch.nn.functional.scaled_dot_product_attention(
q, k, v,
attn_mask=None, dropout_p=self.attn_drop_prob, is_causal=False
).transpose(1, 2).reshape(B_, N, C)
else:
attn = (q @ k.transpose(-2, -1))
relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1) # Wh*Ww,Wh*Ww,nH
relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous() # nH, Wh*Ww, Wh*Ww
attn = attn + relative_position_bias.unsqueeze(0)
if mask is not None:
nW = mask.shape[0]
attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)
attn = attn.view(-1, self.num_heads, N, N)
attn = self.softmax(attn)
else:
attn = self.softmax(attn)
attn = self.attn_drop(attn)
x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
x = self.proj(x)
x = self.proj_drop(x)
return x
class SwinTransformerBlock(nn.Module):
""" Swin Transformer Block.
Args:
dim (int): Number of input channels.
num_heads (int): Number of attention heads.
window_size (int): Window size.
shift_size (int): Shift size for SW-MSA.
mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
drop (float, optional): Dropout rate. Default: 0.0
attn_drop (float, optional): Attention dropout rate. Default: 0.0
drop_path (float, optional): Stochastic depth rate. Default: 0.0
act_layer (nn.Module, optional): Activation layer. Default: nn.GELU
norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
"""
def __init__(self, dim, num_heads, window_size=7, shift_size=0,
mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., drop_path=0.,
act_layer=nn.GELU, norm_layer=nn.LayerNorm):
super().__init__()
self.dim = dim
self.num_heads = num_heads
self.window_size = window_size
self.shift_size = shift_size
self.mlp_ratio = mlp_ratio
assert 0 <= self.shift_size < self.window_size, "shift_size must in 0-window_size"
self.norm1 = norm_layer(dim)
self.attn = WindowAttention(
dim, window_size=to_2tuple(self.window_size), num_heads=num_heads,
qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
self.norm2 = norm_layer(dim)
mlp_hidden_dim = int(dim * mlp_ratio)
self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
self.H = None
self.W = None
def forward(self, x, mask_matrix):
""" Forward function.
Args:
x: Input feature, tensor size (B, H*W, C).
H, W: Spatial resolution of the input feature.
mask_matrix: Attention mask for cyclic shift.
"""
B, L, C = x.shape
H, W = self.H, self.W
assert L == H * W, "input feature has wrong size"
shortcut = x
x = self.norm1(x)
x = x.view(B, H, W, C)
# pad feature maps to multiples of window size
pad_l = pad_t = 0
pad_r = (self.window_size - W % self.window_size) % self.window_size
pad_b = (self.window_size - H % self.window_size) % self.window_size
x = F.pad(x, (0, 0, pad_l, pad_r, pad_t, pad_b))
_, Hp, Wp, _ = x.shape
# cyclic shift
if self.shift_size > 0:
shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))
attn_mask = mask_matrix
else:
shifted_x = x
attn_mask = None
# partition windows
x_windows = window_partition(shifted_x, self.window_size) # nW*B, window_size, window_size, C
x_windows = x_windows.view(-1, self.window_size * self.window_size, C) # nW*B, window_size*window_size, C
# W-MSA/SW-MSA
attn_windows = self.attn(x_windows, mask=attn_mask) # nW*B, window_size*window_size, C
# merge windows
attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C)
shifted_x = window_reverse(attn_windows, self.window_size, Hp, Wp) # B H' W' C
# reverse cyclic shift
if self.shift_size > 0:
x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2))
else:
x = shifted_x
if pad_r > 0 or pad_b > 0:
x = x[:, :H, :W, :].contiguous()
x = x.view(B, H * W, C)
# FFN
x = shortcut + self.drop_path(x)
x = x + self.drop_path(self.mlp(self.norm2(x)))
return x
class PatchMerging(nn.Module):
""" Patch Merging Layer
Args:
dim (int): Number of input channels.
norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
"""
def __init__(self, dim, norm_layer=nn.LayerNorm):
super().__init__()
self.dim = dim
self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)
self.norm = norm_layer(4 * dim)
def forward(self, x, H, W):
""" Forward function.
Args:
x: Input feature, tensor size (B, H*W, C).
H, W: Spatial resolution of the input feature.
"""
B, L, C = x.shape
assert L == H * W, "input feature has wrong size"
x = x.view(B, H, W, C)
# padding
pad_input = (H % 2 == 1) or (W % 2 == 1)
if pad_input:
x = F.pad(x, (0, 0, 0, W % 2, 0, H % 2))
x0 = x[:, 0::2, 0::2, :] # B H/2 W/2 C
x1 = x[:, 1::2, 0::2, :] # B H/2 W/2 C
x2 = x[:, 0::2, 1::2, :] # B H/2 W/2 C
x3 = x[:, 1::2, 1::2, :] # B H/2 W/2 C
x = torch.cat([x0, x1, x2, x3], -1) # B H/2 W/2 4*C
x = x.view(B, -1, 4 * C) # B H/2*W/2 4*C
x = self.norm(x)
x = self.reduction(x)
return x
class BasicLayer(nn.Module):
""" A basic Swin Transformer layer for one stage.
Args:
dim (int): Number of feature channels
depth (int): Depths of this stage.
num_heads (int): Number of attention head.
window_size (int): Local window size. Default: 7.
mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4.
qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
drop (float, optional): Dropout rate. Default: 0.0
attn_drop (float, optional): Attention dropout rate. Default: 0.0
drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0
norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None
use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.
"""
def __init__(self,
dim,
depth,
num_heads,
window_size=7,
mlp_ratio=4.,
qkv_bias=True,
qk_scale=None,
drop=0.,
attn_drop=0.,
drop_path=0.,
norm_layer=nn.LayerNorm,
downsample=None,
use_checkpoint=False):
super().__init__()
self.window_size = window_size
self.shift_size = window_size // 2
self.depth = depth
self.use_checkpoint = use_checkpoint
# build blocks
self.blocks = nn.ModuleList([
SwinTransformerBlock(
dim=dim,
num_heads=num_heads,
window_size=window_size,
shift_size=0 if (i % 2 == 0) else window_size // 2,
mlp_ratio=mlp_ratio,
qkv_bias=qkv_bias,
qk_scale=qk_scale,
drop=drop,
attn_drop=attn_drop,
drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,
norm_layer=norm_layer)
for i in range(depth)])
# patch merging layer
if downsample is not None:
self.downsample = downsample(dim=dim, norm_layer=norm_layer)
else:
self.downsample = None
def forward(self, x, H, W):
""" Forward function.
Args:
x: Input feature, tensor size (B, H*W, C).
H, W: Spatial resolution of the input feature.
"""
# calculate attention mask for SW-MSA
Hp = int(np.ceil(H / self.window_size)) * self.window_size
Wp = int(np.ceil(W / self.window_size)) * self.window_size
img_mask = torch.zeros((1, Hp, Wp, 1), device=x.device) # 1 Hp Wp 1
h_slices = (slice(0, -self.window_size),
slice(-self.window_size, -self.shift_size),
slice(-self.shift_size, None))
w_slices = (slice(0, -self.window_size),
slice(-self.window_size, -self.shift_size),
slice(-self.shift_size, None))
cnt = 0
for h in h_slices:
for w in w_slices:
img_mask[:, h, w, :] = cnt
cnt += 1
mask_windows = window_partition(img_mask, self.window_size) # nW, window_size, window_size, 1
mask_windows = mask_windows.view(-1, self.window_size * self.window_size)
attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))
for blk in self.blocks:
blk.H, blk.W = H, W
if self.use_checkpoint:
x = checkpoint.checkpoint(blk, x, attn_mask)
else:
x = blk(x, attn_mask)
if self.downsample is not None:
x_down = self.downsample(x, H, W)
Wh, Ww = (H + 1) // 2, (W + 1) // 2
return x, H, W, x_down, Wh, Ww
else:
return x, H, W, x, H, W
class PatchEmbed(nn.Module):
""" Image to Patch Embedding
Args:
patch_size (int): Patch token size. Default: 4.
in_channels (int): Number of input image channels. Default: 3.
embed_dim (int): Number of linear projection output channels. Default: 96.
norm_layer (nn.Module, optional): Normalization layer. Default: None
"""
def __init__(self, patch_size=4, in_channels=3, embed_dim=96, norm_layer=None):
super().__init__()
patch_size = to_2tuple(patch_size)
self.patch_size = patch_size
self.in_channels = in_channels
self.embed_dim = embed_dim
self.proj = nn.Conv2d(in_channels, embed_dim, kernel_size=patch_size, stride=patch_size)
if norm_layer is not None:
self.norm = norm_layer(embed_dim)
else:
self.norm = None
def forward(self, x):
"""Forward function."""
# padding
_, _, H, W = x.size()
if W % self.patch_size[1] != 0:
x = F.pad(x, (0, self.patch_size[1] - W % self.patch_size[1]))
if H % self.patch_size[0] != 0:
x = F.pad(x, (0, 0, 0, self.patch_size[0] - H % self.patch_size[0]))
x = self.proj(x) # B C Wh Ww
if self.norm is not None:
Wh, Ww = x.size(2), x.size(3)
x = x.flatten(2).transpose(1, 2)
x = self.norm(x)
x = x.transpose(1, 2).view(-1, self.embed_dim, Wh, Ww)
return x
class SwinTransformer(nn.Module):
""" Swin Transformer backbone.
A PyTorch impl of : `Swin Transformer: Hierarchical Vision Transformer using Shifted Windows` -
https://arxiv.org/pdf/2103.14030
Args:
pretrain_img_size (int): Input image size for training the pretrained model,
used in absolute postion embedding. Default 224.
patch_size (int | tuple(int)): Patch size. Default: 4.
in_channels (int): Number of input image channels. Default: 3.
embed_dim (int): Number of linear projection output channels. Default: 96.
depths (tuple[int]): Depths of each Swin Transformer stage.
num_heads (tuple[int]): Number of attention head of each stage.
window_size (int): Window size. Default: 7.
mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4.
qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True
qk_scale (float): Override default qk scale of head_dim ** -0.5 if set.
drop_rate (float): Dropout rate.
attn_drop_rate (float): Attention dropout rate. Default: 0.
drop_path_rate (float): Stochastic depth rate. Default: 0.2.
norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm.
ape (bool): If True, add absolute position embedding to the patch embedding. Default: False.
patch_norm (bool): If True, add normalization after patch embedding. Default: True.
out_indices (Sequence[int]): Output from which stages.
frozen_stages (int): Stages to be frozen (stop grad and set eval mode).
-1 means not freezing any parameters.
use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.
"""
def __init__(self,
pretrain_img_size=224,
patch_size=4,
in_channels=3,
embed_dim=96,
depths=[2, 2, 6, 2],
num_heads=[3, 6, 12, 24],
window_size=7,
mlp_ratio=4.,
qkv_bias=True,
qk_scale=None,
drop_rate=0.,
attn_drop_rate=0.,
drop_path_rate=0.2,
norm_layer=nn.LayerNorm,
ape=False,
patch_norm=True,
out_indices=(0, 1, 2, 3),
frozen_stages=-1,
use_checkpoint=False):
super().__init__()
self.pretrain_img_size = pretrain_img_size
self.num_layers = len(depths)
self.embed_dim = embed_dim
self.ape = ape
self.patch_norm = patch_norm
self.out_indices = out_indices
self.frozen_stages = frozen_stages
# split image into non-overlapping patches
self.patch_embed = PatchEmbed(
patch_size=patch_size, in_channels=in_channels, embed_dim=embed_dim,
norm_layer=norm_layer if self.patch_norm else None)
# absolute position embedding
if self.ape:
pretrain_img_size = to_2tuple(pretrain_img_size)
patch_size = to_2tuple(patch_size)
patches_resolution = [pretrain_img_size[0] // patch_size[0], pretrain_img_size[1] // patch_size[1]]
self.absolute_pos_embed = nn.Parameter(torch.zeros(1, embed_dim, patches_resolution[0], patches_resolution[1]))
trunc_normal_(self.absolute_pos_embed, std=.02)
self.pos_drop = nn.Dropout(p=drop_rate)
# stochastic depth
dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] # stochastic depth decay rule
# build layers
self.layers = nn.ModuleList()
for i_layer in range(self.num_layers):
layer = BasicLayer(
dim=int(embed_dim * 2 ** i_layer),
depth=depths[i_layer],
num_heads=num_heads[i_layer],
window_size=window_size,
mlp_ratio=mlp_ratio,
qkv_bias=qkv_bias,
qk_scale=qk_scale,
drop=drop_rate,
attn_drop=attn_drop_rate,
drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])],
norm_layer=norm_layer,
downsample=PatchMerging if (i_layer < self.num_layers - 1) else None,
use_checkpoint=use_checkpoint)
self.layers.append(layer)
num_features = [int(embed_dim * 2 ** i) for i in range(self.num_layers)]
self.num_features = num_features
# add a norm layer for each output
for i_layer in out_indices:
layer = norm_layer(num_features[i_layer])
layer_name = f'norm{i_layer}'
self.add_module(layer_name, layer)
self._freeze_stages()
def _freeze_stages(self):
if self.frozen_stages >= 0:
self.patch_embed.eval()
for param in self.patch_embed.parameters():
param.requires_grad = False
if self.frozen_stages >= 1 and self.ape:
self.absolute_pos_embed.requires_grad = False
if self.frozen_stages >= 2:
self.pos_drop.eval()
for i in range(0, self.frozen_stages - 1):
m = self.layers[i]
m.eval()
for param in m.parameters():
param.requires_grad = False
def forward(self, x):
"""Forward function."""
x = self.patch_embed(x)
Wh, Ww = x.size(2), x.size(3)
if self.ape:
# interpolate the position embedding to the corresponding size
absolute_pos_embed = F.interpolate(self.absolute_pos_embed, size=(Wh, Ww), mode='bicubic')
x = (x + absolute_pos_embed) # B Wh*Ww C
outs = []#x.contiguous()]
x = x.flatten(2).transpose(1, 2)
x = self.pos_drop(x)
for i in range(self.num_layers):
layer = self.layers[i]
x_out, H, W, x, Wh, Ww = layer(x, Wh, Ww)
if i in self.out_indices:
norm_layer = getattr(self, f'norm{i}')
x_out = norm_layer(x_out)
out = x_out.view(-1, H, W, self.num_features[i]).permute(0, 3, 1, 2).contiguous()
outs.append(out)
return tuple(outs)
def train(self, mode=True):
"""Convert the model into training mode while keep layers freezed."""
super(SwinTransformer, self).train(mode)
self._freeze_stages()
def swin_v1_t():
model = SwinTransformer(embed_dim=96, depths=[2, 2, 6, 2], num_heads=[3, 6, 12, 24], window_size=7)
return model
def swin_v1_s():
model = SwinTransformer(embed_dim=96, depths=[2, 2, 18, 2], num_heads=[3, 6, 12, 24], window_size=7)
return model
def swin_v1_b():
model = SwinTransformer(embed_dim=128, depths=[2, 2, 18, 2], num_heads=[4, 8, 16, 32], window_size=12)
return model
def swin_v1_l():
model = SwinTransformer(embed_dim=192, depths=[2, 2, 18, 2], num_heads=[6, 12, 24, 48], window_size=12)
return model
py/BiRefNet/models/birefnet.py
+286
View File
@@ -0,0 +1,286 @@
import torch
import torch.nn as nn
import torch.nn.functional as F
from kornia.filters import laplacian
from huggingface_hub import PyTorchModelHubMixin
from BiRefNet.config import Config
from BiRefNet.dataset import class_labels_TR_sorted
from BiRefNet.models.backbones.build_backbone import build_backbone
from BiRefNet.models.modules.decoder_blocks import BasicDecBlk, ResBlk
from BiRefNet.models.modules.lateral_blocks import BasicLatBlk
from BiRefNet.models.modules.aspp import ASPP, ASPPDeformable
from BiRefNet.models.refinement.refiner import Refiner, RefinerPVTInChannels4, RefUNet
from BiRefNet.models.refinement.stem_layer import StemLayer
class BiRefNet(
nn.Module,
PyTorchModelHubMixin,
library_name="birefnet",
repo_url="https://github.com/ZhengPeng7/BiRefNet",
tags=['Image Segmentation', 'Background Removal', 'Mask Generation', 'Dichotomous Image Segmentation', 'Camouflaged Object Detection', 'Salient Object Detection']
):
def __init__(self, bb_pretrained=True):
super(BiRefNet, self).__init__()
self.config = Config()
self.epoch = 1
self.bb = build_backbone(self.config.bb, pretrained=bb_pretrained)
channels = self.config.lateral_channels_in_collection
if self.config.auxiliary_classification:
self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
self.cls_head = nn.Sequential(
nn.Linear(channels[0], len(class_labels_TR_sorted))
)
if self.config.squeeze_block:
self.squeeze_module = nn.Sequential(*[
eval(self.config.squeeze_block.split('_x')[0])(channels[0]+sum(self.config.cxt), channels[0])
for _ in range(eval(self.config.squeeze_block.split('_x')[1]))
])
self.decoder = Decoder(channels)
if self.config.ender:
self.dec_end = nn.Sequential(
nn.Conv2d(1, 16, 3, 1, 1),
nn.Conv2d(16, 1, 3, 1, 1),
nn.ReLU(inplace=True),
)
# refine patch-level segmentation
if self.config.refine:
if self.config.refine == 'itself':
self.stem_layer = StemLayer(in_channels=3+1, inter_channels=48, out_channels=3, norm_layer='BN' if self.config.batch_size > 1 else 'LN')
else:
self.refiner = eval('{}({})'.format(self.config.refine, 'in_channels=3+1'))
if self.config.freeze_bb:
# Freeze the backbone...
print(self.named_parameters())
for key, value in self.named_parameters():
if 'bb.' in key and 'refiner.' not in key:
value.requires_grad = False
def forward_enc(self, x):
if self.config.bb in ['vgg16', 'vgg16bn', 'resnet50']:
x1 = self.bb.conv1(x); x2 = self.bb.conv2(x1); x3 = self.bb.conv3(x2); x4 = self.bb.conv4(x3)
else:
x1, x2, x3, x4 = self.bb(x)
if self.config.mul_scl_ipt == 'cat':
B, C, H, W = x.shape
x1_, x2_, x3_, x4_ = self.bb(F.interpolate(x, size=(H//2, W//2), mode='bilinear', align_corners=True))
x1 = torch.cat([x1, F.interpolate(x1_, size=x1.shape[2:], mode='bilinear', align_corners=True)], dim=1)
x2 = torch.cat([x2, F.interpolate(x2_, size=x2.shape[2:], mode='bilinear', align_corners=True)], dim=1)
x3 = torch.cat([x3, F.interpolate(x3_, size=x3.shape[2:], mode='bilinear', align_corners=True)], dim=1)
x4 = torch.cat([x4, F.interpolate(x4_, size=x4.shape[2:], mode='bilinear', align_corners=True)], dim=1)
elif self.config.mul_scl_ipt == 'add':
B, C, H, W = x.shape
x1_, x2_, x3_, x4_ = self.bb(F.interpolate(x, size=(H//2, W//2), mode='bilinear', align_corners=True))
x1 = x1 + F.interpolate(x1_, size=x1.shape[2:], mode='bilinear', align_corners=True)
x2 = x2 + F.interpolate(x2_, size=x2.shape[2:], mode='bilinear', align_corners=True)
x3 = x3 + F.interpolate(x3_, size=x3.shape[2:], mode='bilinear', align_corners=True)
x4 = x4 + F.interpolate(x4_, size=x4.shape[2:], mode='bilinear', align_corners=True)
class_preds = self.cls_head(self.avgpool(x4).view(x4.shape[0], -1)) if self.training and self.config.auxiliary_classification else None
if self.config.cxt:
x4 = torch.cat(
(
*[
F.interpolate(x1, size=x4.shape[2:], mode='bilinear', align_corners=True),
F.interpolate(x2, size=x4.shape[2:], mode='bilinear', align_corners=True),
F.interpolate(x3, size=x4.shape[2:], mode='bilinear', align_corners=True),
][-len(self.config.cxt):],
x4
),
dim=1
)
return (x1, x2, x3, x4), class_preds
def forward_ori(self, x):
########## Encoder ##########
(x1, x2, x3, x4), class_preds = self.forward_enc(x)
if self.config.squeeze_block:
x4 = self.squeeze_module(x4)
########## Decoder ##########
features = [x, x1, x2, x3, x4]
if self.training and self.config.out_ref:
features.append(laplacian(torch.mean(x, dim=1).unsqueeze(1), kernel_size=5))
scaled_preds = self.decoder(features)
return scaled_preds, class_preds
def forward(self, x):
scaled_preds, class_preds = self.forward_ori(x)
class_preds_lst = [class_preds]
return [scaled_preds, class_preds_lst] if self.training else scaled_preds
class Decoder(nn.Module):
def __init__(self, channels):
super(Decoder, self).__init__()
self.config = Config()
DecoderBlock = eval(self.config.dec_blk)
LateralBlock = eval(self.config.lat_blk)
if self.config.dec_ipt:
self.split = self.config.dec_ipt_split
N_dec_ipt = 64
DBlock = SimpleConvs
ic = 64
ipt_cha_opt = 1
self.ipt_blk5 = DBlock(2**10*3 if self.split else 3, [N_dec_ipt, channels[0]//8][ipt_cha_opt], inter_channels=ic)
self.ipt_blk4 = DBlock(2**8*3 if self.split else 3, [N_dec_ipt, channels[0]//8][ipt_cha_opt], inter_channels=ic)
self.ipt_blk3 = DBlock(2**6*3 if self.split else 3, [N_dec_ipt, channels[1]//8][ipt_cha_opt], inter_channels=ic)
self.ipt_blk2 = DBlock(2**4*3 if self.split else 3, [N_dec_ipt, channels[2]//8][ipt_cha_opt], inter_channels=ic)
self.ipt_blk1 = DBlock(2**0*3 if self.split else 3, [N_dec_ipt, channels[3]//8][ipt_cha_opt], inter_channels=ic)
else:
self.split = None
self.decoder_block4 = DecoderBlock(channels[0]+([N_dec_ipt, channels[0]//8][ipt_cha_opt] if self.config.dec_ipt else 0), channels[1])
self.decoder_block3 = DecoderBlock(channels[1]+([N_dec_ipt, channels[0]//8][ipt_cha_opt] if self.config.dec_ipt else 0), channels[2])
self.decoder_block2 = DecoderBlock(channels[2]+([N_dec_ipt, channels[1]//8][ipt_cha_opt] if self.config.dec_ipt else 0), channels[3])
self.decoder_block1 = DecoderBlock(channels[3]+([N_dec_ipt, channels[2]//8][ipt_cha_opt] if self.config.dec_ipt else 0), channels[3]//2)
self.conv_out1 = nn.Sequential(nn.Conv2d(channels[3]//2+([N_dec_ipt, channels[3]//8][ipt_cha_opt] if self.config.dec_ipt else 0), 1, 1, 1, 0))
self.lateral_block4 = LateralBlock(channels[1], channels[1])
self.lateral_block3 = LateralBlock(channels[2], channels[2])
self.lateral_block2 = LateralBlock(channels[3], channels[3])
if self.config.ms_supervision:
self.conv_ms_spvn_4 = nn.Conv2d(channels[1], 1, 1, 1, 0)
self.conv_ms_spvn_3 = nn.Conv2d(channels[2], 1, 1, 1, 0)
self.conv_ms_spvn_2 = nn.Conv2d(channels[3], 1, 1, 1, 0)
if self.config.out_ref:
_N = 16
self.gdt_convs_4 = nn.Sequential(nn.Conv2d(channels[1], _N, 3, 1, 1), nn.BatchNorm2d(_N) if self.config.batch_size > 1 else nn.Identity(), nn.ReLU(inplace=True))
self.gdt_convs_3 = nn.Sequential(nn.Conv2d(channels[2], _N, 3, 1, 1), nn.BatchNorm2d(_N) if self.config.batch_size > 1 else nn.Identity(), nn.ReLU(inplace=True))
self.gdt_convs_2 = nn.Sequential(nn.Conv2d(channels[3], _N, 3, 1, 1), nn.BatchNorm2d(_N) if self.config.batch_size > 1 else nn.Identity(), nn.ReLU(inplace=True))
self.gdt_convs_pred_4 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0))
self.gdt_convs_pred_3 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0))
self.gdt_convs_pred_2 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0))
self.gdt_convs_attn_4 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0))
self.gdt_convs_attn_3 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0))
self.gdt_convs_attn_2 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0))
def get_patches_batch(self, x, p):
_size_h, _size_w = p.shape[2:]
patches_batch = []
for idx in range(x.shape[0]):
columns_x = torch.split(x[idx], split_size_or_sections=_size_w, dim=-1)
patches_x = []
for column_x in columns_x:
patches_x += [p.unsqueeze(0) for p in torch.split(column_x, split_size_or_sections=_size_h, dim=-2)]
patch_sample = torch.cat(patches_x, dim=1)
patches_batch.append(patch_sample)
return torch.cat(patches_batch, dim=0)
def forward(self, features):
if self.training and self.config.out_ref:
outs_gdt_pred = []
outs_gdt_label = []
x, x1, x2, x3, x4, gdt_gt = features
else:
x, x1, x2, x3, x4 = features
outs = []
if self.config.dec_ipt:
patches_batch = self.get_patches_batch(x, x4) if self.split else x
x4 = torch.cat((x4, self.ipt_blk5(F.interpolate(patches_batch, size=x4.shape[2:], mode='bilinear', align_corners=True))), 1)
p4 = self.decoder_block4(x4)
m4 = self.conv_ms_spvn_4(p4) if self.config.ms_supervision and self.training else None
if self.config.out_ref:
p4_gdt = self.gdt_convs_4(p4)
if self.training:
# >> GT:
m4_dia = m4
gdt_label_main_4 = gdt_gt * F.interpolate(m4_dia, size=gdt_gt.shape[2:], mode='bilinear', align_corners=True)
outs_gdt_label.append(gdt_label_main_4)
# >> Pred:
gdt_pred_4 = self.gdt_convs_pred_4(p4_gdt)
outs_gdt_pred.append(gdt_pred_4)
gdt_attn_4 = self.gdt_convs_attn_4(p4_gdt).sigmoid()
# >> Finally:
p4 = p4 * gdt_attn_4
_p4 = F.interpolate(p4, size=x3.shape[2:], mode='bilinear', align_corners=True)
_p3 = _p4 + self.lateral_block4(x3)
if self.config.dec_ipt:
patches_batch = self.get_patches_batch(x, _p3) if self.split else x
_p3 = torch.cat((_p3, self.ipt_blk4(F.interpolate(patches_batch, size=x3.shape[2:], mode='bilinear', align_corners=True))), 1)
p3 = self.decoder_block3(_p3)
m3 = self.conv_ms_spvn_3(p3) if self.config.ms_supervision and self.training else None
if self.config.out_ref:
p3_gdt = self.gdt_convs_3(p3)
if self.training:
# >> GT:
# m3 --dilation--> m3_dia
# G_3^gt * m3_dia --> G_3^m, which is the label of gradient
m3_dia = m3
gdt_label_main_3 = gdt_gt * F.interpolate(m3_dia, size=gdt_gt.shape[2:], mode='bilinear', align_corners=True)
outs_gdt_label.append(gdt_label_main_3)
# >> Pred:
# p3 --conv--BN--> F_3^G, where F_3^G predicts the \hat{G_3} with xx
# F_3^G --sigmoid--> A_3^G
gdt_pred_3 = self.gdt_convs_pred_3(p3_gdt)
outs_gdt_pred.append(gdt_pred_3)
gdt_attn_3 = self.gdt_convs_attn_3(p3_gdt).sigmoid()
# >> Finally:
# p3 = p3 * A_3^G
p3 = p3 * gdt_attn_3
_p3 = F.interpolate(p3, size=x2.shape[2:], mode='bilinear', align_corners=True)
_p2 = _p3 + self.lateral_block3(x2)
if self.config.dec_ipt:
patches_batch = self.get_patches_batch(x, _p2) if self.split else x
_p2 = torch.cat((_p2, self.ipt_blk3(F.interpolate(patches_batch, size=x2.shape[2:], mode='bilinear', align_corners=True))), 1)
p2 = self.decoder_block2(_p2)
m2 = self.conv_ms_spvn_2(p2) if self.config.ms_supervision and self.training else None
if self.config.out_ref:
p2_gdt = self.gdt_convs_2(p2)
if self.training:
# >> GT:
m2_dia = m2
gdt_label_main_2 = gdt_gt * F.interpolate(m2_dia, size=gdt_gt.shape[2:], mode='bilinear', align_corners=True)
outs_gdt_label.append(gdt_label_main_2)
# >> Pred:
gdt_pred_2 = self.gdt_convs_pred_2(p2_gdt)
outs_gdt_pred.append(gdt_pred_2)
gdt_attn_2 = self.gdt_convs_attn_2(p2_gdt).sigmoid()
# >> Finally:
p2 = p2 * gdt_attn_2
_p2 = F.interpolate(p2, size=x1.shape[2:], mode='bilinear', align_corners=True)
_p1 = _p2 + self.lateral_block2(x1)
if self.config.dec_ipt:
patches_batch = self.get_patches_batch(x, _p1) if self.split else x
_p1 = torch.cat((_p1, self.ipt_blk2(F.interpolate(patches_batch, size=x1.shape[2:], mode='bilinear', align_corners=True))), 1)
_p1 = self.decoder_block1(_p1)
_p1 = F.interpolate(_p1, size=x.shape[2:], mode='bilinear', align_corners=True)
if self.config.dec_ipt:
patches_batch = self.get_patches_batch(x, _p1) if self.split else x
_p1 = torch.cat((_p1, self.ipt_blk1(F.interpolate(patches_batch, size=x.shape[2:], mode='bilinear', align_corners=True))), 1)
p1_out = self.conv_out1(_p1)
if self.config.ms_supervision and self.training:
outs.append(m4)
outs.append(m3)
outs.append(m2)
outs.append(p1_out)
return outs if not (self.config.out_ref and self.training) else ([outs_gdt_pred, outs_gdt_label], outs)
class SimpleConvs(nn.Module):
def __init__(
self, in_channels: int, out_channels: int, inter_channels=64
) -> None:
super().__init__()
self.conv1 = nn.Conv2d(in_channels, inter_channels, 3, 1, 1)
self.conv_out = nn.Conv2d(inter_channels, out_channels, 3, 1, 1)
def forward(self, x):
return self.conv_out(self.conv1(x))
py/BiRefNet/models/modules/aspp.py
+120
View File
@@ -0,0 +1,120 @@
import torch
import torch.nn as nn
import torch.nn.functional as F
from BiRefNet.models.modules.deform_conv import DeformableConv2d
from BiRefNet.config import Config
config = Config()
class _ASPPModule(nn.Module):
def __init__(self, in_channels, planes, kernel_size, padding, dilation):
super(_ASPPModule, self).__init__()
self.atrous_conv = nn.Conv2d(in_channels, planes, kernel_size=kernel_size,
stride=1, padding=padding, dilation=dilation, bias=False)
self.bn = nn.BatchNorm2d(planes) if config.batch_size > 1 else nn.Identity()
self.relu = nn.ReLU(inplace=True)
def forward(self, x):
x = self.atrous_conv(x)
x = self.bn(x)
return self.relu(x)
class ASPP(nn.Module):
def __init__(self, in_channels=64, out_channels=None, output_stride=16):
super(ASPP, self).__init__()
self.down_scale = 1
if out_channels is None:
out_channels = in_channels
self.in_channelster = 256 // self.down_scale
if output_stride == 16:
dilations = [1, 6, 12, 18]
elif output_stride == 8:
dilations = [1, 12, 24, 36]
else:
raise NotImplementedError
self.aspp1 = _ASPPModule(in_channels, self.in_channelster, 1, padding=0, dilation=dilations[0])
self.aspp2 = _ASPPModule(in_channels, self.in_channelster, 3, padding=dilations[1], dilation=dilations[1])
self.aspp3 = _ASPPModule(in_channels, self.in_channelster, 3, padding=dilations[2], dilation=dilations[2])
self.aspp4 = _ASPPModule(in_channels, self.in_channelster, 3, padding=dilations[3], dilation=dilations[3])
self.global_avg_pool = nn.Sequential(nn.AdaptiveAvgPool2d((1, 1)),
nn.Conv2d(in_channels, self.in_channelster, 1, stride=1, bias=False),
nn.BatchNorm2d(self.in_channelster) if config.batch_size > 1 else nn.Identity(),
nn.ReLU(inplace=True))
self.conv1 = nn.Conv2d(self.in_channelster * 5, out_channels, 1, bias=False)
self.bn1 = nn.BatchNorm2d(out_channels) if config.batch_size > 1 else nn.Identity()
self.relu = nn.ReLU(inplace=True)
self.dropout = nn.Dropout(0.5)
def forward(self, x):
x1 = self.aspp1(x)
x2 = self.aspp2(x)
x3 = self.aspp3(x)
x4 = self.aspp4(x)
x5 = self.global_avg_pool(x)
x5 = F.interpolate(x5, size=x1.size()[2:], mode='bilinear', align_corners=True)
x = torch.cat((x1, x2, x3, x4, x5), dim=1)
x = self.conv1(x)
x = self.bn1(x)
x = self.relu(x)
return self.dropout(x)
##################### Deformable
class _ASPPModuleDeformable(nn.Module):
def __init__(self, in_channels, planes, kernel_size, padding):
super(_ASPPModuleDeformable, self).__init__()
self.atrous_conv = DeformableConv2d(in_channels, planes, kernel_size=kernel_size,
stride=1, padding=padding, bias=False)
self.bn = nn.BatchNorm2d(planes) if config.batch_size > 1 else nn.Identity()
self.relu = nn.ReLU(inplace=True)
def forward(self, x):
x = self.atrous_conv(x)
x = self.bn(x)
return self.relu(x)
class ASPPDeformable(nn.Module):
def __init__(self, in_channels, out_channels=None, parallel_block_sizes=[1, 3, 7]):
super(ASPPDeformable, self).__init__()
self.down_scale = 1
if out_channels is None:
out_channels = in_channels
self.in_channelster = 256 // self.down_scale
self.aspp1 = _ASPPModuleDeformable(in_channels, self.in_channelster, 1, padding=0)
self.aspp_deforms = nn.ModuleList([
_ASPPModuleDeformable(in_channels, self.in_channelster, conv_size, padding=int(conv_size//2)) for conv_size in parallel_block_sizes
])
self.global_avg_pool = nn.Sequential(nn.AdaptiveAvgPool2d((1, 1)),
nn.Conv2d(in_channels, self.in_channelster, 1, stride=1, bias=False),
nn.BatchNorm2d(self.in_channelster) if config.batch_size > 1 else nn.Identity(),
nn.ReLU(inplace=True))
self.conv1 = nn.Conv2d(self.in_channelster * (2 + len(self.aspp_deforms)), out_channels, 1, bias=False)
self.bn1 = nn.BatchNorm2d(out_channels) if config.batch_size > 1 else nn.Identity()
self.relu = nn.ReLU(inplace=True)
self.dropout = nn.Dropout(0.5)
def forward(self, x):
x1 = self.aspp1(x)
x_aspp_deforms = [aspp_deform(x) for aspp_deform in self.aspp_deforms]
x5 = self.global_avg_pool(x)
x5 = F.interpolate(x5, size=x1.size()[2:], mode='bilinear', align_corners=True)
x = torch.cat((x1, *x_aspp_deforms, x5), dim=1)
x = self.conv1(x)
x = self.bn1(x)
x = self.relu(x)
return self.dropout(x)
py/BiRefNet/models/modules/decoder_blocks.py
+66
View File
@@ -0,0 +1,66 @@
import torch
import torch.nn as nn
from BiRefNet.models.modules.aspp import ASPP, ASPPDeformable
from BiRefNet.config import Config
config = Config()
class BasicDecBlk(nn.Module):
def __init__(self, in_channels=64, out_channels=64, inter_channels=64):
super(BasicDecBlk, self).__init__()
inter_channels = in_channels // 4 if config.dec_channels_inter == 'adap' else 64
self.conv_in = nn.Conv2d(in_channels, inter_channels, 3, 1, padding=1)
self.relu_in = nn.ReLU(inplace=True)
if config.dec_att == 'ASPP':
self.dec_att = ASPP(in_channels=inter_channels)
elif config.dec_att == 'ASPPDeformable':
self.dec_att = ASPPDeformable(in_channels=inter_channels)
self.conv_out = nn.Conv2d(inter_channels, out_channels, 3, 1, padding=1)
self.bn_in = nn.BatchNorm2d(inter_channels) if config.batch_size > 1 else nn.Identity()
self.bn_out = nn.BatchNorm2d(out_channels) if config.batch_size > 1 else nn.Identity()
def forward(self, x):
x = self.conv_in(x)
x = self.bn_in(x)
x = self.relu_in(x)
if hasattr(self, 'dec_att'):
x = self.dec_att(x)
x = self.conv_out(x)
x = self.bn_out(x)
return x
class ResBlk(nn.Module):
def __init__(self, in_channels=64, out_channels=None, inter_channels=64):
super(ResBlk, self).__init__()
if out_channels is None:
out_channels = in_channels
inter_channels = in_channels // 4 if config.dec_channels_inter == 'adap' else 64
self.conv_in = nn.Conv2d(in_channels, inter_channels, 3, 1, padding=1)
self.bn_in = nn.BatchNorm2d(inter_channels) if config.batch_size > 1 else nn.Identity()
self.relu_in = nn.ReLU(inplace=True)
if config.dec_att == 'ASPP':
self.dec_att = ASPP(in_channels=inter_channels)
elif config.dec_att == 'ASPPDeformable':
self.dec_att = ASPPDeformable(in_channels=inter_channels)
self.conv_out = nn.Conv2d(inter_channels, out_channels, 3, 1, padding=1)
self.bn_out = nn.BatchNorm2d(out_channels) if config.batch_size > 1 else nn.Identity()
self.conv_resi = nn.Conv2d(in_channels, out_channels, 1, 1, 0)
def forward(self, x):
_x = self.conv_resi(x)
x = self.conv_in(x)
x = self.bn_in(x)
x = self.relu_in(x)
if hasattr(self, 'dec_att'):
x = self.dec_att(x)
x = self.conv_out(x)
x = self.bn_out(x)
return x + _x
py/BiRefNet/modules/deform_conv.py → py/BiRefNet/models/modules/deform_conv.py
View File
py/BiRefNet/models/modules/lateral_blocks.py
+21
View File
@@ -0,0 +1,21 @@
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from functools import partial
from BiRefNet.config import Config
config = Config()
class BasicLatBlk(nn.Module):
def __init__(self, in_channels=64, out_channels=64, inter_channels=64):
super(BasicLatBlk, self).__init__()
inter_channels = in_channels // 4 if config.dec_channels_inter == 'adap' else 64
self.conv = nn.Conv2d(in_channels, out_channels, 1, 1, 0)
def forward(self, x):
x = self.conv(x)
return x
py/BiRefNet/modules/mlp.py → py/BiRefNet/models/modules/mlp.py
View File
py/BiRefNet/models/modules/prompt_encoder.py
+222
View File
@@ -0,0 +1,222 @@
import numpy as np
import torch
import torch.nn as nn
from typing import Any, Optional, Tuple, Type
class PromptEncoder(nn.Module):
def __init__(
self,
embed_dim=256,
image_embedding_size=1024,
input_image_size=(1024, 1024),
mask_in_chans=16,
activation=nn.GELU
) -> None:
super().__init__()
"""
Codes are partially from SAM: https://github.com/facebookresearch/segment-anything/blob/6fdee8f2727f4506cfbbe553e23b895e27956588/segment_anything/modeling/prompt_encoder.py.
Arguments:
embed_dim (int): The prompts' embedding dimension
image_embedding_size (tuple(int, int)): The spatial size of the
image embedding, as (H, W).
input_image_size (int): The padded size of the image as input
to the image encoder, as (H, W).
mask_in_chans (int): The number of hidden channels used for
encoding input masks.
activation (nn.Module): The activation to use when encoding
input masks.
"""
super().__init__()
self.embed_dim = embed_dim
self.input_image_size = input_image_size
self.image_embedding_size = image_embedding_size
self.pe_layer = PositionEmbeddingRandom(embed_dim // 2)
self.num_point_embeddings: int = 4 # pos/neg point + 2 box corners
point_embeddings = [nn.Embedding(1, embed_dim) for i in range(self.num_point_embeddings)]
self.point_embeddings = nn.ModuleList(point_embeddings)
self.not_a_point_embed = nn.Embedding(1, embed_dim)
self.mask_input_size = (4 * image_embedding_size[0], 4 * image_embedding_size[1])
self.mask_downscaling = nn.Sequential(
nn.Conv2d(1, mask_in_chans // 4, kernel_size=2, stride=2),
LayerNorm2d(mask_in_chans // 4),
activation(),
nn.Conv2d(mask_in_chans // 4, mask_in_chans, kernel_size=2, stride=2),
LayerNorm2d(mask_in_chans),
activation(),
nn.Conv2d(mask_in_chans, embed_dim, kernel_size=1),
)
self.no_mask_embed = nn.Embedding(1, embed_dim)
def get_dense_pe(self) -> torch.Tensor:
"""
Returns the positional encoding used to encode point prompts,
applied to a dense set of points the shape of the image encoding.
Returns:
torch.Tensor: Positional encoding with shape
1x(embed_dim)x(embedding_h)x(embedding_w)
"""
return self.pe_layer(self.image_embedding_size).unsqueeze(0)
def _embed_points(
self,
points: torch.Tensor,
labels: torch.Tensor,
pad: bool,
) -> torch.Tensor:
"""Embeds point prompts."""
points = points + 0.5 # Shift to center of pixel
if pad:
padding_point = torch.zeros((points.shape[0], 1, 2), device=points.device)
padding_label = -torch.ones((labels.shape[0], 1), device=labels.device)
points = torch.cat([points, padding_point], dim=1)
labels = torch.cat([labels, padding_label], dim=1)
point_embedding = self.pe_layer.forward_with_coords(points, self.input_image_size)
point_embedding[labels == -1] = 0.0
point_embedding[labels == -1] += self.not_a_point_embed.weight
point_embedding[labels == 0] += self.point_embeddings[0].weight
point_embedding[labels == 1] += self.point_embeddings[1].weight
return point_embedding
def _embed_boxes(self, boxes: torch.Tensor) -> torch.Tensor:
"""Embeds box prompts."""
boxes = boxes + 0.5 # Shift to center of pixel
coords = boxes.reshape(-1, 2, 2)
corner_embedding = self.pe_layer.forward_with_coords(coords, self.input_image_size)
corner_embedding[:, 0, :] += self.point_embeddings[2].weight
corner_embedding[:, 1, :] += self.point_embeddings[3].weight
return corner_embedding
def _embed_masks(self, masks: torch.Tensor) -> torch.Tensor:
"""Embeds mask inputs."""
mask_embedding = self.mask_downscaling(masks)
return mask_embedding
def _get_batch_size(
self,
points: Optional[Tuple[torch.Tensor, torch.Tensor]],
boxes: Optional[torch.Tensor],
masks: Optional[torch.Tensor],
) -> int:
"""
Gets the batch size of the output given the batch size of the input prompts.
"""
if points is not None:
return points[0].shape[0]
elif boxes is not None:
return boxes.shape[0]
elif masks is not None:
return masks.shape[0]
else:
return 1
def _get_device(self) -> torch.device:
return self.point_embeddings[0].weight.device
def forward(
self,
points: Optional[Tuple[torch.Tensor, torch.Tensor]],
boxes: Optional[torch.Tensor],
masks: Optional[torch.Tensor],
) -> Tuple[torch.Tensor, torch.Tensor]:
"""
Embeds different types of prompts, returning both sparse and dense
embeddings.
Arguments:
points (tuple(torch.Tensor, torch.Tensor) or none): point coordinates
and labels to embed.
boxes (torch.Tensor or none): boxes to embed
masks (torch.Tensor or none): masks to embed
Returns:
torch.Tensor: sparse embeddings for the points and boxes, with shape
BxNx(embed_dim), where N is determined by the number of input points
and boxes.
torch.Tensor: dense embeddings for the masks, in the shape
Bx(embed_dim)x(embed_H)x(embed_W)
"""
bs = self._get_batch_size(points, boxes, masks)
sparse_embeddings = torch.empty((bs, 0, self.embed_dim), device=self._get_device())
if points is not None:
coords, labels = points
point_embeddings = self._embed_points(coords, labels, pad=(boxes is None))
sparse_embeddings = torch.cat([sparse_embeddings, point_embeddings], dim=1)
if boxes is not None:
box_embeddings = self._embed_boxes(boxes)
sparse_embeddings = torch.cat([sparse_embeddings, box_embeddings], dim=1)
if masks is not None:
dense_embeddings = self._embed_masks(masks)
else:
dense_embeddings = self.no_mask_embed.weight.reshape(1, -1, 1, 1).expand(
bs, -1, self.image_embedding_size[0], self.image_embedding_size[1]
)
return sparse_embeddings, dense_embeddings
class PositionEmbeddingRandom(nn.Module):
"""
Positional encoding using random spatial frequencies.
"""
def __init__(self, num_pos_feats: int = 64, scale: Optional[float] = None) -> None:
super().__init__()
if scale is None or scale <= 0.0:
scale = 1.0
self.register_buffer(
"positional_encoding_gaussian_matrix",
scale * torch.randn((2, num_pos_feats)),
)
def _pe_encoding(self, coords: torch.Tensor) -> torch.Tensor:
"""Positionally encode points that are normalized to [0,1]."""
# assuming coords are in [0, 1]^2 square and have d_1 x ... x d_n x 2 shape
coords = 2 * coords - 1
coords = coords @ self.positional_encoding_gaussian_matrix
coords = 2 * np.pi * coords
# outputs d_1 x ... x d_n x C shape
return torch.cat([torch.sin(coords), torch.cos(coords)], dim=-1)
def forward(self, size: Tuple[int, int]) -> torch.Tensor:
"""Generate positional encoding for a grid of the specified size."""
h, w = size
device: Any = self.positional_encoding_gaussian_matrix.device
grid = torch.ones((h, w), device=device, dtype=torch.float32)
y_embed = grid.cumsum(dim=0) - 0.5
x_embed = grid.cumsum(dim=1) - 0.5
y_embed = y_embed / h
x_embed = x_embed / w
pe = self._pe_encoding(torch.stack([x_embed, y_embed], dim=-1))
return pe.permute(2, 0, 1) # C x H x W
def forward_with_coords(
self, coords_input: torch.Tensor, image_size: Tuple[int, int]
) -> torch.Tensor:
"""Positionally encode points that are not normalized to [0,1]."""
coords = coords_input.clone()
coords[:, :, 0] = coords[:, :, 0] / image_size[1]
coords[:, :, 1] = coords[:, :, 1] / image_size[0]
return self._pe_encoding(coords.to(torch.float)) # B x N x C
class LayerNorm2d(nn.Module):
def __init__(self, num_channels: int, eps: float = 1e-6) -> None:
super().__init__()
self.weight = nn.Parameter(torch.ones(num_channels))
self.bias = nn.Parameter(torch.zeros(num_channels))
self.eps = eps
def forward(self, x: torch.Tensor) -> torch.Tensor:
u = x.mean(1, keepdim=True)
s = (x - u).pow(2).mean(1, keepdim=True)
x = (x - u) / torch.sqrt(s + self.eps)
x = self.weight[:, None, None] * x + self.bias[:, None, None]
return x
py/BiRefNet/modules/utils.py → py/BiRefNet/models/modules/utils.py
View File
py/BiRefNet/models/refinement/refiner.py
+252
View File
@@ -0,0 +1,252 @@
import torch
import torch.nn as nn
from collections import OrderedDict
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision.models import vgg16, vgg16_bn
from torchvision.models import resnet50
from BiRefNet.config import Config
from BiRefNet.dataset import class_labels_TR_sorted
from BiRefNet.models.backbones.build_backbone import build_backbone
from BiRefNet.models.modules.decoder_blocks import BasicDecBlk
from BiRefNet.models.modules.lateral_blocks import BasicLatBlk
from BiRefNet.models.refinement.stem_layer import StemLayer
class RefinerPVTInChannels4(nn.Module):
def __init__(self, in_channels=3+1):
super(RefinerPVTInChannels4, self).__init__()
self.config = Config()
self.epoch = 1
self.bb = build_backbone(self.config.bb, params_settings='in_channels=4')
lateral_channels_in_collection = {
'vgg16': [512, 256, 128, 64], 'vgg16bn': [512, 256, 128, 64], 'resnet50': [1024, 512, 256, 64],
'pvt_v2_b2': [512, 320, 128, 64], 'pvt_v2_b5': [512, 320, 128, 64],
'swin_v1_b': [1024, 512, 256, 128], 'swin_v1_l': [1536, 768, 384, 192],
}
channels = lateral_channels_in_collection[self.config.bb]
self.squeeze_module = BasicDecBlk(channels[0], channels[0])
self.decoder = Decoder(channels)
if 0:
for key, value in self.named_parameters():
if 'bb.' in key:
value.requires_grad = False
def forward(self, x):
if isinstance(x, list):
x = torch.cat(x, dim=1)
########## Encoder ##########
if self.config.bb in ['vgg16', 'vgg16bn', 'resnet50']:
x1 = self.bb.conv1(x)
x2 = self.bb.conv2(x1)
x3 = self.bb.conv3(x2)
x4 = self.bb.conv4(x3)
else:
x1, x2, x3, x4 = self.bb(x)
x4 = self.squeeze_module(x4)
########## Decoder ##########
features = [x, x1, x2, x3, x4]
scaled_preds = self.decoder(features)
return scaled_preds
class Refiner(nn.Module):
def __init__(self, in_channels=3+1):
super(Refiner, self).__init__()
self.config = Config()
self.epoch = 1
self.stem_layer = StemLayer(in_channels=in_channels, inter_channels=48, out_channels=3, norm_layer='BN' if self.config.batch_size > 1 else 'LN')
self.bb = build_backbone(self.config.bb)
lateral_channels_in_collection = {
'vgg16': [512, 256, 128, 64], 'vgg16bn': [512, 256, 128, 64], 'resnet50': [1024, 512, 256, 64],
'pvt_v2_b2': [512, 320, 128, 64], 'pvt_v2_b5': [512, 320, 128, 64],
'swin_v1_b': [1024, 512, 256, 128], 'swin_v1_l': [1536, 768, 384, 192],
}
channels = lateral_channels_in_collection[self.config.bb]
self.squeeze_module = BasicDecBlk(channels[0], channels[0])
self.decoder = Decoder(channels)
if 0:
for key, value in self.named_parameters():
if 'bb.' in key:
value.requires_grad = False
def forward(self, x):
if isinstance(x, list):
x = torch.cat(x, dim=1)
x = self.stem_layer(x)
########## Encoder ##########
if self.config.bb in ['vgg16', 'vgg16bn', 'resnet50']:
x1 = self.bb.conv1(x)
x2 = self.bb.conv2(x1)
x3 = self.bb.conv3(x2)
x4 = self.bb.conv4(x3)
else:
x1, x2, x3, x4 = self.bb(x)
x4 = self.squeeze_module(x4)
########## Decoder ##########
features = [x, x1, x2, x3, x4]
scaled_preds = self.decoder(features)
return scaled_preds
class Decoder(nn.Module):
def __init__(self, channels):
super(Decoder, self).__init__()
self.config = Config()
DecoderBlock = eval('BasicDecBlk')
LateralBlock = eval('BasicLatBlk')
self.decoder_block4 = DecoderBlock(channels[0], channels[1])
self.decoder_block3 = DecoderBlock(channels[1], channels[2])
self.decoder_block2 = DecoderBlock(channels[2], channels[3])
self.decoder_block1 = DecoderBlock(channels[3], channels[3]//2)
self.lateral_block4 = LateralBlock(channels[1], channels[1])
self.lateral_block3 = LateralBlock(channels[2], channels[2])
self.lateral_block2 = LateralBlock(channels[3], channels[3])
if self.config.ms_supervision:
self.conv_ms_spvn_4 = nn.Conv2d(channels[1], 1, 1, 1, 0)
self.conv_ms_spvn_3 = nn.Conv2d(channels[2], 1, 1, 1, 0)
self.conv_ms_spvn_2 = nn.Conv2d(channels[3], 1, 1, 1, 0)
self.conv_out1 = nn.Sequential(nn.Conv2d(channels[3]//2, 1, 1, 1, 0))
def forward(self, features):
x, x1, x2, x3, x4 = features
outs = []
p4 = self.decoder_block4(x4)
_p4 = F.interpolate(p4, size=x3.shape[2:], mode='bilinear', align_corners=True)
_p3 = _p4 + self.lateral_block4(x3)
p3 = self.decoder_block3(_p3)
_p3 = F.interpolate(p3, size=x2.shape[2:], mode='bilinear', align_corners=True)
_p2 = _p3 + self.lateral_block3(x2)
p2 = self.decoder_block2(_p2)
_p2 = F.interpolate(p2, size=x1.shape[2:], mode='bilinear', align_corners=True)
_p1 = _p2 + self.lateral_block2(x1)
_p1 = self.decoder_block1(_p1)
_p1 = F.interpolate(_p1, size=x.shape[2:], mode='bilinear', align_corners=True)
p1_out = self.conv_out1(_p1)
if self.config.ms_supervision:
outs.append(self.conv_ms_spvn_4(p4))
outs.append(self.conv_ms_spvn_3(p3))
outs.append(self.conv_ms_spvn_2(p2))
outs.append(p1_out)
return outs
class RefUNet(nn.Module):
# Refinement
def __init__(self, in_channels=3+1):
super(RefUNet, self).__init__()
self.encoder_1 = nn.Sequential(
nn.Conv2d(in_channels, 64, 3, 1, 1),
nn.Conv2d(64, 64, 3, 1, 1),
nn.BatchNorm2d(64),
nn.ReLU(inplace=True)
)
self.encoder_2 = nn.Sequential(
nn.MaxPool2d(2, 2, ceil_mode=True),
nn.Conv2d(64, 64, 3, 1, 1),
nn.BatchNorm2d(64),
nn.ReLU(inplace=True)
)
self.encoder_3 = nn.Sequential(
nn.MaxPool2d(2, 2, ceil_mode=True),
nn.Conv2d(64, 64, 3, 1, 1),
nn.BatchNorm2d(64),
nn.ReLU(inplace=True)
)
self.encoder_4 = nn.Sequential(
nn.MaxPool2d(2, 2, ceil_mode=True),
nn.Conv2d(64, 64, 3, 1, 1),
nn.BatchNorm2d(64),
nn.ReLU(inplace=True)
)
self.pool4 = nn.MaxPool2d(2, 2, ceil_mode=True)
#####
self.decoder_5 = nn.Sequential(
nn.Conv2d(64, 64, 3, 1, 1),
nn.BatchNorm2d(64),
nn.ReLU(inplace=True)
)
#####
self.decoder_4 = nn.Sequential(
nn.Conv2d(128, 64, 3, 1, 1),
nn.BatchNorm2d(64),
nn.ReLU(inplace=True)
)
self.decoder_3 = nn.Sequential(
nn.Conv2d(128, 64, 3, 1, 1),
nn.BatchNorm2d(64),
nn.ReLU(inplace=True)
)
self.decoder_2 = nn.Sequential(
nn.Conv2d(128, 64, 3, 1, 1),
nn.BatchNorm2d(64),
nn.ReLU(inplace=True)
)
self.decoder_1 = nn.Sequential(
nn.Conv2d(128, 64, 3, 1, 1),
nn.BatchNorm2d(64),
nn.ReLU(inplace=True)
)
self.conv_d0 = nn.Conv2d(64, 1, 3, 1, 1)
self.upscore2 = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
def forward(self, x):
outs = []
if isinstance(x, list):
x = torch.cat(x, dim=1)
hx = x
hx1 = self.encoder_1(hx)
hx2 = self.encoder_2(hx1)
hx3 = self.encoder_3(hx2)
hx4 = self.encoder_4(hx3)
hx = self.decoder_5(self.pool4(hx4))
hx = torch.cat((self.upscore2(hx), hx4), 1)
d4 = self.decoder_4(hx)
hx = torch.cat((self.upscore2(d4), hx3), 1)
d3 = self.decoder_3(hx)
hx = torch.cat((self.upscore2(d3), hx2), 1)
d2 = self.decoder_2(hx)
hx = torch.cat((self.upscore2(d2), hx1), 1)
d1 = self.decoder_1(hx)
x = self.conv_d0(d1)
outs.append(x)
return outs
py/BiRefNet/refinement/stem_layer.py → py/BiRefNet/models/refinement/stem_layer.py
+1 -1
View File
@@ -1,5 +1,5 @@
import torch.nn as nn
from BiRefNet.modules.utils import build_act_layer, build_norm_layer
from BiRefNet.models.modules.utils import build_act_layer, build_norm_layer
class StemLayer(nn.Module):
py/BiRefNet/requirements.txt
+15
View File
@@ -0,0 +1,15 @@
--extra-index-url https://download.pytorch.org/whl/cu118
torch==2.0.1
--extra-index-url https://download.pytorch.org/whl/cu118
torchvision==0.15.2
numpy<2
opencv-python
timm
scipy
scikit-image
kornia
tqdm
prettytable
huggingface_hub
py/BiRefNet/rm_cache.sh
+20
View File
@@ -0,0 +1,20 @@
#!/bin/bash
rm -rf __pycache__ */__pycache__
# Val
rm -r tmp*
# Train
rm slurm*
rm -r ckpt
rm nohup.out*
# Eval
rm -r evaluation/eval-*
rm -r tmp*
rm -r e_logs/
# System
rm core-*-python-*
clear
py/BiRefNet/sub.sh
+17
View File
@@ -0,0 +1,17 @@
#!/bin/sh
# Example: ./sub.sh tmp_proj 0,1,2,3 3 --> Use 0,1,2,3 for training, release GPUs, use GPU:3 for inference.
# module load gcc/11.2.0 cuda/11.8 cudnn/8.6.0_cu11x && cpu_core_num=6
module load compilers/cuda/11.8 compilers/gcc/12.2.0 cudnn/8.4.0.27_cuda11.x && cpu_core_num=32
export PYTHONUNBUFFERED=1
method=${1:-"BSL"}
devices=${2:-0}
gpu_num=$(($(echo ${devices%%,} | grep -o "," | wc -l)+1))
sbatch --nodes=1 -p vip_gpu_ailab -A ai4bio \
--gres=gpu:${gpu_num} --ntasks-per-node=1 --cpus-per-task=$((gpu_num*cpu_core_num)) \
./train_test.sh ${method} ${devices}
hostname
py/BiRefNet/test.sh
+29
View File
@@ -0,0 +1,29 @@
devices=${1:-0}
pred_root=${2:-e_preds}
# Inference
CUDA_VISIBLE_DEVICES=${devices} python inference.py --pred_root ${pred_root}
echo Inference finished at $(date)
# Evaluation
log_dir=e_logs && mkdir ${log_dir}
task=$(python3 config.py)
case "${task}" in
"DIS5K") testsets='DIS-VD,DIS-TE1,DIS-TE2,DIS-TE3,DIS-TE4' ;;
"COD") testsets='CHAMELEON,NC4K,TE-CAMO,TE-COD10K' ;;
"HRSOD") testsets='DAVIS-S,TE-HRSOD,TE-UHRSD,DUT-OMRON,TE-DUTS' ;;
"General") testsets='DIS-VD' ;;
"Matting") testsets='TE-P3M-500-P' ;;
esac
testsets=(`echo ${testsets} | tr ',' ' '`) && testsets=${testsets[@]}
for testset in ${testsets}; do
python eval_existingOnes.py --pred_root ${pred_root} --data_lst ${testset} > ${log_dir}/eval_${testset}.out
# nohup python eval_existingOnes.py --pred_root ${pred_root} --data_lst ${testset} > ${log_dir}/eval_${testset}.out 2>&1 &
done
echo Evaluation started at $(date)
py/BiRefNet/train.py
+333
View File
@@ -0,0 +1,333 @@
import os
import datetime
import argparse
import torch
import torch.nn as nn
import torch.optim as optim
from torch.autograd import Variable
from BiRefNet.config import Config
from BiRefNet.loss import PixLoss, ClsLoss
from BiRefNet.dataset import MyData
from BiRefNet.models.birefnet import BiRefNet
from BiRefNet.utils import Logger, AverageMeter, set_seed, check_state_dict
from torch.utils.data.distributed import DistributedSampler
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.distributed import init_process_group, destroy_process_group, get_rank
from torch.cuda import amp
parser = argparse.ArgumentParser(description='')
parser.add_argument('--resume', default=None, type=str, help='path to latest checkpoint')
parser.add_argument('--epochs', default=120, type=int)
parser.add_argument('--trainset', default='DIS5K', type=str, help="Options: 'DIS5K'")
parser.add_argument('--ckpt_dir', default=None, help='Temporary folder')
parser.add_argument('--testsets', default='DIS-VD+DIS-TE1+DIS-TE2+DIS-TE3+DIS-TE4', type=str)
parser.add_argument('--dist', default=False, type=lambda x: x == 'True')
args = parser.parse_args()
config = Config()
if config.rand_seed:
set_seed(config.rand_seed)
if config.use_fp16:
# Half Precision
scaler = amp.GradScaler(enabled=config.use_fp16)
# DDP
to_be_distributed = args.dist
if to_be_distributed:
init_process_group(backend="nccl", timeout=datetime.timedelta(seconds=3600*10))
device = int(os.environ["LOCAL_RANK"])
else:
device = config.device
epoch_st = 1
# make dir for ckpt
os.makedirs(args.ckpt_dir, exist_ok=True)
# Init log file
logger = Logger(os.path.join(args.ckpt_dir, "log.txt"))
logger_loss_idx = 1
# log model and optimizer params
# logger.info("Model details:"); logger.info(model)
logger.info("datasets: load_all={}, compile={}.".format(config.load_all, config.compile))
logger.info("Other hyperparameters:"); logger.info(args)
print('batch size:', config.batch_size)
if os.path.exists(os.path.join(config.data_root_dir, config.task, args.testsets.strip('+').split('+')[0])):
args.testsets = args.testsets.strip('+').split('+')
else:
args.testsets = []
# Init model
def prepare_dataloader(dataset: torch.utils.data.Dataset, batch_size: int, to_be_distributed=False, is_train=True):
if to_be_distributed:
return torch.utils.data.DataLoader(
dataset=dataset, batch_size=batch_size, num_workers=min(config.num_workers, batch_size), pin_memory=True,
shuffle=False, sampler=DistributedSampler(dataset), drop_last=True
)
else:
return torch.utils.data.DataLoader(
dataset=dataset, batch_size=batch_size, num_workers=min(config.num_workers, batch_size, 0), pin_memory=True,
shuffle=is_train, drop_last=True
)
def init_data_loaders(to_be_distributed):
# Prepare dataset
train_loader = prepare_dataloader(
MyData(datasets=config.training_set, image_size=config.size, is_train=True),
config.batch_size, to_be_distributed=to_be_distributed, is_train=True
)
print(len(train_loader), "batches of train dataloader {} have been created.".format(config.training_set))
test_loaders = {}
for testset in args.testsets:
_data_loader_test = prepare_dataloader(
MyData(datasets=testset, image_size=config.size, is_train=False),
config.batch_size_valid, is_train=False
)
print(len(_data_loader_test), "batches of valid dataloader {} have been created.".format(testset))
test_loaders[testset] = _data_loader_test
return train_loader, test_loaders
def init_models_optimizers(epochs, to_be_distributed):
model = BiRefNet(bb_pretrained=True)
if args.resume:
if os.path.isfile(args.resume):
logger.info("=> loading checkpoint '{}'".format(args.resume))
state_dict = torch.load(args.resume, map_location='cpu')
state_dict = check_state_dict(state_dict)
model.load_state_dict(state_dict)
global epoch_st
epoch_st = int(args.resume.rstrip('.pth').split('epoch_')[-1]) + 1
else:
logger.info("=> no checkpoint found at '{}'".format(args.resume))
if to_be_distributed:
model = model.to(device)
model = DDP(model, device_ids=[device])
else:
model = model.to(device)
if config.compile:
model = torch.compile(model, mode=['default', 'reduce-overhead', 'max-autotune'][0])
if config.precisionHigh:
torch.set_float32_matmul_precision('high')
# Setting optimizer
if config.optimizer == 'AdamW':
optimizer = optim.AdamW(params=model.parameters(), lr=config.lr, weight_decay=1e-2)
elif config.optimizer == 'Adam':
optimizer = optim.Adam(params=model.parameters(), lr=config.lr, weight_decay=0)
lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer,
milestones=[lde if lde > 0 else epochs + lde + 1 for lde in config.lr_decay_epochs],
gamma=config.lr_decay_rate
)
logger.info("Optimizer details:"); logger.info(optimizer)
logger.info("Scheduler details:"); logger.info(lr_scheduler)
return model, optimizer, lr_scheduler
class Trainer:
def __init__(
self, data_loaders, model_opt_lrsch,
):
self.model, self.optimizer, self.lr_scheduler = model_opt_lrsch
self.train_loader, self.test_loaders = data_loaders
if config.out_ref:
self.criterion_gdt = nn.BCELoss() if not config.use_fp16 else nn.BCEWithLogitsLoss()
# Setting Losses
self.pix_loss = PixLoss()
self.cls_loss = ClsLoss()
# Others
self.loss_log = AverageMeter()
if config.lambda_adv_g:
self.optimizer_d, self.lr_scheduler_d, self.disc, self.adv_criterion = self._load_adv_components()
self.disc_update_for_odd = 0
def _load_adv_components(self):
# AIL
from loss import Discriminator
disc = Discriminator(channels=3, img_size=config.size)
if to_be_distributed:
disc = disc.to(device)
disc = DDP(disc, device_ids=[device], broadcast_buffers=False)
else:
disc = disc.to(device)
if config.compile:
disc = torch.compile(disc, mode=['default', 'reduce-overhead', 'max-autotune'][0])
adv_criterion = nn.BCELoss() if not config.use_fp16 else nn.BCEWithLogitsLoss()
if config.optimizer == 'AdamW':
optimizer_d = optim.AdamW(params=disc.parameters(), lr=config.lr, weight_decay=1e-2)
elif config.optimizer == 'Adam':
optimizer_d = optim.Adam(params=disc.parameters(), lr=config.lr, weight_decay=0)
lr_scheduler_d = torch.optim.lr_scheduler.MultiStepLR(
optimizer_d,
milestones=[lde if lde > 0 else args.epochs + lde + 1 for lde in config.lr_decay_epochs],
gamma=config.lr_decay_rate
)
return optimizer_d, lr_scheduler_d, disc, adv_criterion
def _train_batch(self, batch):
inputs = batch[0].to(device)
gts = batch[1].to(device)
class_labels = batch[2].to(device)
if config.use_fp16:
with amp.autocast(enabled=config.use_fp16):
scaled_preds, class_preds_lst = self.model(inputs)
if config.out_ref:
(outs_gdt_pred, outs_gdt_label), scaled_preds = scaled_preds
for _idx, (_gdt_pred, _gdt_label) in enumerate(zip(outs_gdt_pred, outs_gdt_label)):
_gdt_pred = nn.functional.interpolate(_gdt_pred, size=_gdt_label.shape[2:], mode='bilinear', align_corners=True)#.sigmoid()
# _gdt_label = _gdt_label.sigmoid()
loss_gdt = self.criterion_gdt(_gdt_pred, _gdt_label) if _idx == 0 else self.criterion_gdt(_gdt_pred, _gdt_label) + loss_gdt
# self.loss_dict['loss_gdt'] = loss_gdt.item()
if None in class_preds_lst:
loss_cls = 0.
else:
loss_cls = self.cls_loss(class_preds_lst, class_labels) * 1.0
self.loss_dict['loss_cls'] = loss_cls.item()
# Loss
loss_pix = self.pix_loss(scaled_preds, torch.clamp(gts, 0, 1)) * 1.0
self.loss_dict['loss_pix'] = loss_pix.item()
# since there may be several losses for sal, the lambdas for them (lambdas_pix) are inside the loss.py
loss = loss_pix + loss_cls
if config.out_ref:
loss = loss + loss_gdt * 1.0
if config.lambda_adv_g:
# gen
valid = Variable(torch.cuda.FloatTensor(scaled_preds[-1].shape[0], 1).fill_(1.0), requires_grad=False).to(device)
adv_loss_g = self.adv_criterion(self.disc(scaled_preds[-1] * inputs), valid) * config.lambda_adv_g
loss += adv_loss_g
self.loss_dict['loss_adv'] = adv_loss_g.item()
self.disc_update_for_odd += 1
# self.loss_log.update(loss.item(), inputs.size(0))
# self.optimizer.zero_grad()
# loss.backward()
# self.optimizer.step()
self.optimizer.zero_grad()
scaler.scale(loss).backward()
scaler.step(self.optimizer)
scaler.update()
if config.lambda_adv_g and self.disc_update_for_odd % 2 == 0:
# disc
fake = Variable(torch.cuda.FloatTensor(scaled_preds[-1].shape[0], 1).fill_(0.0), requires_grad=False).to(device)
adv_loss_real = self.adv_criterion(self.disc(gts * inputs), valid)
adv_loss_fake = self.adv_criterion(self.disc(scaled_preds[-1].detach() * inputs.detach()), fake)
adv_loss_d = (adv_loss_real + adv_loss_fake) / 2 * config.lambda_adv_d
self.loss_dict['loss_adv_d'] = adv_loss_d.item()
# self.optimizer_d.zero_grad()
# adv_loss_d.backward()
# self.optimizer_d.step()
self.optimizer_d.zero_grad()
scaler.scale(adv_loss_d).backward()
scaler.step(self.optimizer_d)
scaler.update()
else:
scaled_preds, class_preds_lst = self.model(inputs)
if config.out_ref:
(outs_gdt_pred, outs_gdt_label), scaled_preds = scaled_preds
for _idx, (_gdt_pred, _gdt_label) in enumerate(zip(outs_gdt_pred, outs_gdt_label)):
_gdt_pred = nn.functional.interpolate(_gdt_pred, size=_gdt_label.shape[2:], mode='bilinear', align_corners=True).sigmoid()
_gdt_label = _gdt_label.sigmoid()
loss_gdt = self.criterion_gdt(_gdt_pred, _gdt_label) if _idx == 0 else self.criterion_gdt(_gdt_pred, _gdt_label) + loss_gdt
# self.loss_dict['loss_gdt'] = loss_gdt.item()
if None in class_preds_lst:
loss_cls = 0.
else:
loss_cls = self.cls_loss(class_preds_lst, class_labels) * 1.0
self.loss_dict['loss_cls'] = loss_cls.item()
# Loss
loss_pix = self.pix_loss(scaled_preds, torch.clamp(gts, 0, 1)) * 1.0
self.loss_dict['loss_pix'] = loss_pix.item()
# since there may be several losses for sal, the lambdas for them (lambdas_pix) are inside the loss.py
loss = loss_pix + loss_cls
if config.out_ref:
loss = loss + loss_gdt * 1.0
if config.lambda_adv_g:
# gen
valid = Variable(torch.cuda.FloatTensor(scaled_preds[-1].shape[0], 1).fill_(1.0), requires_grad=False).to(device)
adv_loss_g = self.adv_criterion(self.disc(scaled_preds[-1] * inputs), valid) * config.lambda_adv_g
loss += adv_loss_g
self.loss_dict['loss_adv'] = adv_loss_g.item()
self.disc_update_for_odd += 1
self.loss_log.update(loss.item(), inputs.size(0))
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
if config.lambda_adv_g and self.disc_update_for_odd % 2 == 0:
# disc
fake = Variable(torch.cuda.FloatTensor(scaled_preds[-1].shape[0], 1).fill_(0.0), requires_grad=False).to(device)
adv_loss_real = self.adv_criterion(self.disc(gts * inputs), valid)
adv_loss_fake = self.adv_criterion(self.disc(scaled_preds[-1].detach() * inputs.detach()), fake)
adv_loss_d = (adv_loss_real + adv_loss_fake) / 2 * config.lambda_adv_d
self.loss_dict['loss_adv_d'] = adv_loss_d.item()
self.optimizer_d.zero_grad()
adv_loss_d.backward()
self.optimizer_d.step()
def train_epoch(self, epoch):
global logger_loss_idx
self.model.train()
self.loss_dict = {}
if epoch > args.epochs + config.finetune_last_epochs[1]:
for k in self.pix_loss.lambdas_pix_last.keys():
if k.lower() == config.finetune_last_epochs[0].lower():
self.pix_loss.lambdas_pix_last[k] = config.lambdas_pix_last[k] * 0.5
else:
self.pix_loss.lambdas_pix_last[k] = 0
for batch_idx, batch in enumerate(self.train_loader):
self._train_batch(batch)
# Logger
if batch_idx % 20 == 0:
info_progress = 'Epoch[{0}/{1}] Iter[{2}/{3}].'.format(epoch, args.epochs, batch_idx, len(self.train_loader))
info_loss = 'Training Losses'
for loss_name, loss_value in self.loss_dict.items():
info_loss += ', {}: {:.3f}'.format(loss_name, loss_value)
logger.info(' '.join((info_progress, info_loss)))
info_loss = '@==Final== Epoch[{0}/{1}] Training Loss: {loss.avg:.3f} '.format(epoch, args.epochs, loss=self.loss_log)
logger.info(info_loss)
self.lr_scheduler.step()
if config.lambda_adv_g:
self.lr_scheduler_d.step()
return self.loss_log.avg
def main():
trainer = Trainer(
data_loaders=init_data_loaders(to_be_distributed),
model_opt_lrsch=init_models_optimizers(args.epochs, to_be_distributed)
)
for epoch in range(epoch_st, args.epochs+1):
train_loss = trainer.train_epoch(epoch)
# Save checkpoint
# DDP
if epoch >= args.epochs - config.save_last and epoch % config.save_step == 0:
torch.save(
trainer.model.module.state_dict() if to_be_distributed else trainer.model.state_dict(),
os.path.join(args.ckpt_dir, 'epoch_{}.pth'.format(epoch))
)
if to_be_distributed:
destroy_process_group()
if __name__ == '__main__':
main()
py/BiRefNet/train.sh
+42
View File
@@ -0,0 +1,42 @@
#!/bin/bash
# Run script
# Settings of training & test for different tasks.
method="$1"
task=$(python3 config.py)
case "${task}" in
"DIS5K") epochs=600 && val_last=50 && step=5 ;;
"COD") epochs=150 && val_last=50 && step=5 ;;
"HRSOD") epochs=150 && val_last=50 && step=5 ;;
"General") epochs=250 && val_last=20 && step=2 ;;
"Matting") epochs=100 && val_last=20 && step=2 ;;
esac
testsets=NO # Non-existing folder to skip.
# testsets=TE-COD10K # for COD
# Train
devices=$2
nproc_per_node=$(echo ${devices%%,} | grep -o "," | wc -l)
to_be_distributed=`echo ${nproc_per_node} | awk '{if($e > 0) print "True"; else print "False";}'`
echo Training started at $(date)
if [ ${to_be_distributed} == "True" ]
then
# Adapt the nproc_per_node by the number of GPUs. Give 8989 as the default value of master_port.
echo "Multi-GPU mode received..."
CUDA_VISIBLE_DEVICES=${devices} \
torchrun --nproc_per_node $((nproc_per_node+1)) --master_port=${3:-8999} \
train.py --ckpt_dir ckpt/${method} --epochs ${epochs} \
--testsets ${testsets} \
--dist ${to_be_distributed} \
--resume xx/xx-epoch_244.pth
else
echo "Single-GPU mode received..."
CUDA_VISIBLE_DEVICES=${devices} \
python train.py --ckpt_dir ckpt/${method} --epochs ${epochs} \
--testsets ${testsets} \
--dist ${to_be_distributed} \
--resume xx/xx-epoch_244.pth
fi
echo Training finished at $(date)
py/BiRefNet/train_test.sh
+11
View File
@@ -0,0 +1,11 @@
#!/bin/sh
method=${1:-"BSL"}
devices=${2:-"0,1,2,3,4,5,6,7"}
bash train.sh ${method} ${devices}
devices_test=${3:-0}
bash test.sh ${devices_test}
hostname
py/BiRefNet/tutorials/BiRefNet_inference.ipynb
+1575
View File
File diff suppressed because it is too large Load Diff
py/BiRefNet/tutorials/BiRefNet_pth2onnx.ipynb
+312
View File
@@ -0,0 +1,312 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "LTj2A0RUQFNo"
},
"source": [
"# Convert our BiRefNet weights to onnx format.\n",
"\n",
"> This colab file is modified from [Kazuhito00](https://github.com/Kazuhito00)'s nice work.\n",
"\n",
"> Repo: https://github.com/Kazuhito00/BiRefNet-ONNX-Sample \n",
"> Original Colab: https://colab.research.google.com/github/Kazuhito00/BiRefNet-ONNX-Sample/blob/main/Convert2ONNX.ipynb\n",
"\n",
"+ Currently, Colab with 12.7GB RAM / 15GB GPU Mem cannot hold the transformation of BiRefNet in default setting. So, I take BiRefNet with swin_v1_tiny backbone as an example."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Online Colab version: https://colab.research.google.com/drive/1z6OruR52LOvDDpnp516F-N4EyPGrp5om"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "781JHjLJmveh"
},
"outputs": [],
"source": [
"import torch\n",
"\n",
"\n",
"weights_file = 'BiRefNet-general-bb_swin_v1_tiny-epoch_232.pth' # https://github.com/ZhengPeng7/BiRefNet/releases/download/v1/BiRefNet-general-bb_swin_v1_tiny-epoch_232.pth\n",
"device = 'cuda' if torch.cuda.is_available() else 'cpu'"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"with open('config.py') as fp:\n",
" file_lines = fp.read()\n",
"if 'swin_v1_tiny' in weights_file:\n",
" print('Set `swin_v1_tiny` as the backbone.')\n",
" file_lines = file_lines.replace(\n",
" '''\n",
" 'pvt_v2_b2', 'pvt_v2_b5', # 9-bs10, 10-bs5\n",
" ][6]\n",
" ''',\n",
" '''\n",
" 'pvt_v2_b2', 'pvt_v2_b5', # 9-bs10, 10-bs5\n",
" ][3]\n",
" ''',\n",
" )\n",
" with open('config.py', mode=\"w\") as fp:\n",
" fp.write(file_lines)\n",
"else:\n",
" file_lines = file_lines.replace(\n",
" '''\n",
" 'pvt_v2_b2', 'pvt_v2_b5', # 9-bs10, 10-bs5\n",
" ][3]\n",
" ''',\n",
" '''\n",
" 'pvt_v2_b2', 'pvt_v2_b5', # 9-bs10, 10-bs5\n",
" ][6]\n",
" ''',\n",
" )\n",
" with open('config.py', mode=\"w\") as fp:\n",
" fp.write(file_lines)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "7lFgKfPS8Icy"
},
"outputs": [],
"source": [
"from utils import check_state_dict\n",
"from models.birefnet import BiRefNet\n",
"\n",
"\n",
"birefnet = BiRefNet(bb_pretrained=False)\n",
"state_dict = torch.load('./{}'.format(weights_file), map_location=device)\n",
"state_dict = check_state_dict(state_dict)\n",
"birefnet.load_state_dict(state_dict)\n",
"\n",
"torch.set_float32_matmul_precision(['high', 'highest'][0])\n",
"\n",
"birefnet.to(device)\n",
"_ = birefnet.eval()"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "JVgJAdgxQVJW"
},
"source": [
"# Process deform_conv2d in the conversion to ONNX"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "vJiZv0L75kTe"
},
"outputs": [],
"source": [
"from torchvision.ops.deform_conv import DeformConv2d\n",
"import deform_conv2d_onnx_exporter\n",
"\n",
"# register deform_conv2d operator\n",
"deform_conv2d_onnx_exporter.register_deform_conv2d_onnx_op()\n",
"\n",
"def convert_to_onnx(net, file_name='output.onnx', input_shape=(1024, 1024), device=device):\n",
" input = torch.randn(1, 3, input_shape[0], input_shape[1]).to(device)\n",
"\n",
" input_layer_names = ['input_image']\n",
" output_layer_names = ['output_image']\n",
"\n",
" torch.onnx.export(\n",
" net,\n",
" input,\n",
" file_name,\n",
" verbose=False,\n",
" opset_version=17,\n",
" input_names=input_layer_names,\n",
" output_names=output_layer_names,\n",
" )\n",
"convert_to_onnx(birefnet, weights_file.replace('.pth', '.onnx'), input_shape=(1024, 1024), device=device)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "-eU-g40P1zS-"
},
"source": [
"# Load ONNX weights and do the inference."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "LZ4HVqcoDvto"
},
"outputs": [],
"source": [
"from PIL import Image\n",
"from torchvision import transforms\n",
"\n",
"\n",
"transform_image = transforms.Compose([\n",
" transforms.Resize((1024, 1024)),\n",
" transforms.ToTensor(),\n",
" transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n",
"])\n",
"\n",
"imagepath = './Helicopter-HR.jpg'\n",
"image = Image.open(imagepath)\n",
"input_images = transform_image(image).unsqueeze(0).to(device)\n",
"input_images_numpy = input_images.cpu().numpy()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "rwzdKX1EfYkd"
},
"outputs": [],
"source": [
"import onnxruntime\n",
"import matplotlib.pyplot as plt\n",
"\n",
"\n",
"providers = ['CPUExecutionProvider'] if device == 'cpu' else ['CUDAExecutionProvider']\n",
"onnx_session = onnxruntime.InferenceSession(\n",
" weights_file.replace('.pth', '.onnx'),\n",
" providers=providers\n",
")\n",
"input_name = onnx_session.get_inputs()[0].name\n",
"print(onnxruntime.get_device(), onnx_session.get_providers())"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "DJVtxZUZum4-"
},
"outputs": [],
"source": [
"from time import time\n",
"import matplotlib.pyplot as plt\n",
"\n",
"time_st = time()\n",
"pred_onnx = torch.tensor(\n",
" onnx_session.run(None, {input_name: input_images_numpy if device == 'cpu' else input_images_numpy})[-1]\n",
").squeeze(0).sigmoid().cpu()\n",
"print(time() - time_st)\n",
"\n",
"plt.imshow(pred_onnx.squeeze(), cmap='gray'); plt.show()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"with torch.no_grad():\n",
" preds = birefnet(input_images)[-1].sigmoid().cpu()\n",
"plt.imshow(preds.squeeze(), cmap='gray'); plt.show()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"diff = abs(preds - pred_onnx)\n",
"print('sum(diff):', diff.sum())\n",
"plt.imshow((diff).squeeze(), cmap='gray'); plt.show()"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "qzYHflt92Bjd"
},
"source": [
"# Efficiency Comparison between .pth and .onnx"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "A5IYfT-uzphA",
"outputId": "2999e345-950e-41b3-ddd3-9f58a71a3f21"
},
"outputs": [],
"source": [
"%%timeit\n",
"with torch.no_grad():\n",
" preds = birefnet(input_images)[-1].sigmoid().cpu()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "G0Ul4rfNg1za"
},
"outputs": [],
"source": [
"%%timeit\n",
"pred_onnx = torch.tensor(\n",
" onnx_session.run(None, {input_name: input_images_numpy})[-1]\n",
").squeeze(0).sigmoid().cpu()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {
"accelerator": "GPU",
"colab": {
"gpuType": "T4",
"provenance": []
},
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.14"
}
},
"nbformat": 4,
"nbformat_minor": 4
}
py/BiRefNet/utils.py
+97
View File
@@ -0,0 +1,97 @@
import logging
import os
import torch
from torchvision import transforms
import numpy as np
import random
import cv2
from PIL import Image
def path_to_image(path, size=(1024, 1024), color_type=['rgb', 'gray'][0]):
if color_type.lower() == 'rgb':
image = cv2.imread(path)
elif color_type.lower() == 'gray':
image = cv2.imread(path, cv2.IMREAD_GRAYSCALE)
else:
print('Select the color_type to return, either to RGB or gray image.')
return
if size:
image = cv2.resize(image, size, interpolation=cv2.INTER_LINEAR)
if color_type.lower() == 'rgb':
image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB)).convert('RGB')
else:
image = Image.fromarray(image).convert('L')
return image
def check_state_dict(state_dict, unwanted_prefix='_orig_mod.'):
for k, v in list(state_dict.items()):
if k.startswith(unwanted_prefix):
state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k)
return state_dict
def generate_smoothed_gt(gts):
epsilon = 0.001
new_gts = (1-epsilon)*gts+epsilon/2
return new_gts
class Logger():
def __init__(self, path="log.txt"):
self.logger = logging.getLogger('BiRefNet')
self.file_handler = logging.FileHandler(path, "w")
self.stdout_handler = logging.StreamHandler()
self.stdout_handler.setFormatter(logging.Formatter('%(asctime)s %(levelname)s %(message)s'))
self.file_handler.setFormatter(logging.Formatter('%(asctime)s %(levelname)s %(message)s'))
self.logger.addHandler(self.file_handler)
self.logger.addHandler(self.stdout_handler)
self.logger.setLevel(logging.INFO)
self.logger.propagate = False
def info(self, txt):
self.logger.info(txt)
def close(self):
self.file_handler.close()
self.stdout_handler.close()
class AverageMeter(object):
"""Computes and stores the average and current value"""
def __init__(self):
self.reset()
def reset(self):
self.val = 0.0
self.avg = 0.0
self.sum = 0.0
self.count = 0.0
def update(self, val, n=1):
self.val = val
self.sum += val * n
self.count += n
self.avg = self.sum / self.count
def save_checkpoint(state, path, filename="latest.pth"):
torch.save(state, os.path.join(path, filename))
def save_tensor_img(tenor_im, path):
im = tenor_im.cpu().clone()
im = im.squeeze(0)
tensor2pil = transforms.ToPILImage()
im = tensor2pil(im)
im.save(path)
def set_seed(seed):
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
np.random.seed(seed)
random.seed(seed)
torch.backends.cudnn.deterministic = True
py/BiRefNet/backbones/__init__.py → py/BiRefNet_legacy/__init__.py
View File
py/BiRefNet/modules/__init__.py → py/BiRefNet_legacy/backbones/__init__.py
View File
py/BiRefNet/backbones/build_backbone.py → py/BiRefNet_legacy/backbones/build_backbone.py
+2 -2
View File
@@ -2,8 +2,8 @@ import torch
import torch.nn as nn
from collections import OrderedDict
from torchvision.models import vgg16, vgg16_bn, VGG16_Weights, VGG16_BN_Weights, resnet50, ResNet50_Weights
from BiRefNet.backbones.pvt_v2 import pvt_v2_b2, pvt_v2_b5
from BiRefNet.backbones.swin_v1 import swin_v1_t, swin_v1_s, swin_v1_b, swin_v1_l
from BiRefNet_legacy.backbones.pvt_v2 import pvt_v2_b2, pvt_v2_b5
from BiRefNet_legacy.backbones.swin_v1 import swin_v1_t, swin_v1_s, swin_v1_b, swin_v1_l
from ..config import Config
py/BiRefNet/backbones/pvt_v2.py → py/BiRefNet_legacy/backbones/pvt_v2.py
View File
py/BiRefNet/backbones/swin_v1.py → py/BiRefNet_legacy/backbones/swin_v1.py
View File
py/BiRefNet/baseline.py → py/BiRefNet_legacy/baseline.py
+7 -7
View File
@@ -8,13 +8,13 @@ from torchvision.models import vgg16, vgg16_bn
from torchvision.models import resnet50
from kornia.filters import laplacian
from BiRefNet.backbones.build_backbone import build_backbone
from BiRefNet.modules.decoder_blocks import BasicDecBlk, ResBlk, HierarAttDecBlk
from BiRefNet.modules.lateral_blocks import BasicLatBlk
from BiRefNet.modules.aspp import ASPP, ASPPDeformable
from BiRefNet.modules.ing import *
from BiRefNet.refinement.refiner import Refiner, RefinerPVTInChannels4, RefUNet
from BiRefNet.refinement.stem_layer import StemLayer
from BiRefNet_legacy.backbones.build_backbone import build_backbone
from BiRefNet_legacy.modules.decoder_blocks import BasicDecBlk, ResBlk, HierarAttDecBlk
from BiRefNet_legacy.modules.lateral_blocks import BasicLatBlk
from BiRefNet_legacy.modules.aspp import ASPP, ASPPDeformable
from BiRefNet_legacy.modules.ing import *
from BiRefNet_legacy.refinement.refiner import Refiner, RefinerPVTInChannels4, RefUNet
from BiRefNet_legacy.refinement.stem_layer import StemLayer
from .config import Config
from .dataset import class_labels_TR_sorted
py/BiRefNet_legacy/config.py
+104
View File
@@ -0,0 +1,104 @@
import os
import math
from folder_paths import models_dir
class Config():
def __init__(self) -> None:
self.ms_supervision = True
self.out_ref = self.ms_supervision and True
self.dec_ipt = True
self.dec_ipt_split = True
self.locate_head = False
self.cxt_num = [0, 3][1] # multi-scale skip connections from encoder
self.mul_scl_ipt = ['', 'add', 'cat'][2]
self.refine = ['', 'itself', 'RefUNet', 'Refiner', 'RefinerPVTInChannels4'][0]
self.progressive_ref = self.refine and True
self.ender = self.progressive_ref and False
self.scale = self.progressive_ref and 2
self.dec_att = ['', 'ASPP', 'ASPPDeformable'][2]
self.squeeze_block = ['', 'BasicDecBlk_x1', 'ResBlk_x4', 'ASPP_x3', 'ASPPDeformable_x3'][1]
self.dec_blk = ['BasicDecBlk', 'ResBlk', 'HierarAttDecBlk'][0]
self.auxiliary_classification = False
self.refine_iteration = 1
self.freeze_bb = False
self.precisionHigh = True
self.compile = True
self.load_all = True
self.verbose_eval = True
self.size = 1024
self.batch_size = 2
self.IoU_finetune_last_epochs = [0, -40][1] # choose 0 to skip
if self.dec_blk == 'HierarAttDecBlk':
self.batch_size = 2 ** [0, 1, 2, 3, 4][2]
self.model = [
'BiRefNet',
][0]
# Components
self.lat_blk = ['BasicLatBlk'][0]
self.dec_channels_inter = ['fixed', 'adap'][0]
# Backbone
self.bb = [
'vgg16', 'vgg16bn', 'resnet50', # 0, 1, 2
'pvt_v2_b2', 'pvt_v2_b5', # 3-bs10, 4-bs5
'swin_v1_b', 'swin_v1_l' # 5-bs9, 6-bs6
][6]
self.lateral_channels_in_collection = {
'vgg16': [512, 256, 128, 64], 'vgg16bn': [512, 256, 128, 64], 'resnet50': [1024, 512, 256, 64],
'pvt_v2_b2': [512, 320, 128, 64], 'pvt_v2_b5': [512, 320, 128, 64],
'swin_v1_b': [1024, 512, 256, 128], 'swin_v1_l': [1536, 768, 384, 192],
}[self.bb]
if self.mul_scl_ipt == 'cat':
self.lateral_channels_in_collection = [channel * 2 for channel in self.lateral_channels_in_collection]
self.cxt = self.lateral_channels_in_collection[1:][::-1][-self.cxt_num:] if self.cxt_num else []
self.sys_home_dir = models_dir
self.weights_root_dir = os.path.join(self.sys_home_dir, "BiRefNet")
self.weights = {
'pvt_v2_b2': os.path.join(self.weights_root_dir, 'pvt_v2_b2.pth'),
'pvt_v2_b5': os.path.join(self.weights_root_dir, ['pvt_v2_b5.pth', 'pvt_v2_b5_22k.pth'][0]),
'swin_v1_b': os.path.join(self.weights_root_dir, ['swin_base_patch4_window12_384_22kto1k.pth', 'swin_base_patch4_window12_384_22k.pth'][0]),
'swin_v1_l': os.path.join(self.weights_root_dir, ['swin_large_patch4_window12_384_22kto1k.pth', 'swin_large_patch4_window12_384_22k.pth'][0]),
}
# Training
self.num_workers = 5 # will be decrease to min(it, batch_size) at the initialization of the data_loader
self.optimizer = ['Adam', 'AdamW'][0]
self.lr = 1e-5 * math.sqrt(self.batch_size / 5) # adapt the lr linearly
self.lr_decay_epochs = [1e4] # Set to negative N to decay the lr in the last N-th epoch.
self.lr_decay_rate = 0.5
self.only_S_MAE = False
self.SDPA_enabled = False # Bug. Slower and errors occur in multi-GPUs
# Data
self.data_root_dir = os.path.join(self.sys_home_dir, 'datasets/dis')
self.dataset = ['DIS5K', 'COD', 'SOD'][0]
self.preproc_methods = ['flip', 'enhance', 'rotate', 'pepper', 'crop'][:4]
# Loss
self.lambdas_pix_last = {
# not 0 means opening this loss
# original rate -- 1 : 30 : 1.5 : 0.2, bce x 30
'bce': 30 * 1, # high performance
'iou': 0.5 * 1, # 0 / 255
'iou_patch': 0.5 * 0, # 0 / 255, win_size = (64, 64)
'mse': 150 * 0, # can smooth the saliency map
'triplet': 3 * 0,
'reg': 100 * 0,
'ssim': 10 * 1, # help contours,
'cnt': 5 * 0, # help contours
}
self.lambdas_cls = {
'ce': 5.0
}
# Adv
self.lambda_adv_g = 10. * 0 # turn to 0 to avoid adv training
self.lambda_adv_d = 3. * (self.lambda_adv_g > 0)
# others
self.device = [0, 'cpu'][0] # .to(0) = .to('cuda:0')
self.batch_size_valid = 1
self.rand_seed = 7
py/BiRefNet_legacy/dataset.py
+140
View File
@@ -0,0 +1,140 @@
import os
import cv2
from tqdm import tqdm
from PIL import Image
from torch.utils import data
from torchvision import transforms
from .preproc import preproc
from .config import Config
from glob import glob
Image.MAX_IMAGE_PIXELS = None # remove DecompressionBombWarning
config = Config()
_class_labels_TR_sorted = 'Airplane, Ant, Antenna, Archery, Axe, BabyCarriage, Bag, BalanceBeam, Balcony, Balloon, Basket, BasketballHoop, Beatle, Bed, Bee, Bench, Bicycle, BicycleFrame, BicycleStand, Boat, Bonsai, BoomLift, Bridge, BunkBed, Butterfly, Button, Cable, CableLift, Cage, Camcorder, Cannon, Canoe, Car, CarParkDropArm, Carriage, Cart, Caterpillar, CeilingLamp, Centipede, Chair, Clip, Clock, Clothes, CoatHanger, Comb, ConcretePumpTruck, Crack, Crane, Cup, DentalChair, Desk, DeskChair, Diagram, DishRack, DoorHandle, Dragonfish, Dragonfly, Drum, Earphone, Easel, ElectricIron, Excavator, Eyeglasses, Fan, Fence, Fencing, FerrisWheel, FireExtinguisher, Fishing, Flag, FloorLamp, Forklift, GasStation, Gate, Gear, Goal, Golf, GymEquipment, Hammock, Handcart, Handcraft, Handrail, HangGlider, Harp, Harvester, Headset, Helicopter, Helmet, Hook, HorizontalBar, Hydrovalve, IroningTable, Jewelry, Key, KidsPlayground, Kitchenware, Kite, Knife, Ladder, LaundryRack, Lightning, Lobster, Locust, Machine, MachineGun, MagazineRack, Mantis, Medal, MemorialArchway, Microphone, Missile, MobileHolder, Monitor, Mosquito, Motorcycle, MovingTrolley, Mower, MusicPlayer, MusicStand, ObservationTower, Octopus, OilWell, OlympicLogo, OperatingTable, OutdoorFitnessEquipment, Parachute, Pavilion, Piano, Pipe, PlowHarrow, PoleVault, Punchbag, Rack, Racket, Rifle, Ring, Robot, RockClimbing, Rope, Sailboat, Satellite, Scaffold, Scale, Scissor, Scooter, Sculpture, Seadragon, Seahorse, Seal, SewingMachine, Ship, Shoe, ShoppingCart, ShoppingTrolley, Shower, Shrimp, Signboard, Skateboarding, Skeleton, Skiing, Spade, SpeedBoat, Spider, Spoon, Stair, Stand, Stationary, SteeringWheel, Stethoscope, Stool, Stove, StreetLamp, SweetStand, Swing, Sword, TV, Table, TableChair, TableLamp, TableTennis, Tank, Tapeline, Teapot, Telescope, Tent, TobaccoPipe, Toy, Tractor, TrafficLight, TrafficSign, Trampoline, TransmissionTower, Tree, Tricycle, TrimmerCover, Tripod, Trombone, Truck, Trumpet, Tuba, UAV, Umbrella, UnevenBars, UtilityPole, VacuumCleaner, Violin, Wakesurfing, Watch, WaterTower, WateringPot, Well, WellLid, Wheel, Wheelchair, WindTurbine, Windmill, WineGlass, WireWhisk, Yacht'
class_labels_TR_sorted = _class_labels_TR_sorted.split(', ')
class MyData(data.Dataset):
def __init__(self, data_root, image_size, is_train=True):
self.size_train = image_size
self.size_test = image_size
self.keep_size = not config.size
self.data_size = (config.size, config.size)
self.is_train = is_train
self.load_all = config.load_all
self.device = config.device
self.dataset = data_root.replace('\\', '/').split('/')[-1]
if self.is_train and config.auxiliary_classification:
self.cls_name2id = {_name: _id for _id, _name in enumerate(class_labels_TR_sorted)}
self.transform_image = transforms.Compose([
transforms.Resize(self.data_size),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
][self.load_all or self.keep_size:])
self.transform_label = transforms.Compose([
transforms.Resize(self.data_size),
transforms.ToTensor(),
][self.load_all or self.keep_size:])
## 'im' and 'gt' need modifying
image_root = os.path.join(data_root, 'im')
self.image_paths = [os.path.join(image_root, p) for p in os.listdir(image_root)]
self.label_paths = [p.replace('/im/', '/gt/').replace('.jpg', '.png') for p in self.image_paths]
if self.load_all:
self.images_loaded, self.labels_loaded = [], []
self.class_labels_loaded = []
# for image_path, label_path in zip(self.image_paths, self.label_paths):
for image_path, label_path in tqdm(zip(self.image_paths, self.label_paths), total=len(self.image_paths)):
_image = cv2.imread(image_path)
_label = cv2.imread(label_path, cv2.IMREAD_GRAYSCALE)
if not self.keep_size:
_image_rs = cv2.resize(_image, (config.size, config.size), interpolation=cv2.INTER_LINEAR)
_label_rs = cv2.resize(_label, (config.size, config.size), interpolation=cv2.INTER_LINEAR)
self.images_loaded.append(
Image.fromarray(cv2.cvtColor(_image_rs, cv2.COLOR_BGR2RGB)).convert('RGB')
)
self.labels_loaded.append(
Image.fromarray(_label_rs).convert('L')
)
self.class_labels_loaded.append(
self.cls_name2id[label_path.split('/')[-1].split('#')[3]] if self.is_train and config.auxiliary_classification else -1
)
def __getitem__(self, index):
if self.load_all:
image = self.images_loaded[index]
class_label = self.class_labels_loaded[index] if self.is_train and config.auxiliary_classification else -1
else:
image = Image.open(self.image_paths[index]).convert('RGB')
# loading image and label
if self.is_train:
image, label = preproc(image, image, preproc_methods=config.preproc_methods)
# else:
# if _label.shape[0] > 2048 or _label.shape[1] > 2048:
# _image = cv2.resize(_image, (2048, 2048), interpolation=cv2.INTER_LINEAR)
# _label = cv2.resize(_label, (2048, 2048), interpolation=cv2.INTER_LINEAR)
image, label = self.transform_image(image), self.transform_label(label)
if self.is_train:
return image, label, class_label
else:
return image, label, self.label_paths[index]
def __len__(self):
return len(self.image_paths)
class YouData(data.Dataset):
def __init__(self, data_root, image_size, is_train=True):
self.size_train = image_size
self.size_test = image_size
self.keep_size = not config.size
self.data_size = (config.size, config.size)
self.is_train = is_train
self.load_all = config.load_all
self.device = config.device
self.dataset = data_root.replace('\\', '/').split('/')[-1]
if self.is_train and config.auxiliary_classification:
self.cls_name2id = {_name: _id for _id, _name in enumerate(class_labels_TR_sorted)}
self.transform_image = transforms.Compose([
transforms.Resize(self.data_size),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
][self.load_all or self.keep_size:])
## 'im' and 'gt' need modifying
self.image_paths = glob(data_root + "/*")
self.img_sizes = []
if self.load_all:
self.images_loaded, self.labels_loaded = [], []
for image_path in tqdm(self.image_paths, total=len(self.image_paths)):
_image = cv2.imread(image_path)
if not self.keep_size:
_image_rs = cv2.resize(_image, (config.size, config.size), interpolation=cv2.INTER_LINEAR)
self.images_loaded.append(
Image.fromarray(cv2.cvtColor(_image_rs, cv2.COLOR_BGR2RGB)).convert('RGB')
)
self.img_sizes.append(_image.shape[:2])
def __getitem__(self, index):
if self.load_all:
image = self.images_loaded[index]
else:
image = Image.open(self.image_paths[index]).convert('RGB')
# loading image and label
if self.is_train:
image, _ = preproc(image, image, preproc_methods=config.preproc_methods)
image = self.transform_image(image)
size = self.img_sizes[index]
return image, size
def __len__(self):
return len(self.image_paths)
py/BiRefNet/refinement/__init__.py → py/BiRefNet_legacy/modules/__init__.py
View File
py/BiRefNet/modules/aspp.py → py/BiRefNet_legacy/modules/aspp.py
+1 -1
View File
@@ -1,7 +1,7 @@
import torch
import torch.nn as nn
import torch.nn.functional as F
from BiRefNet.modules.deform_conv import DeformableConv2d
from BiRefNet_legacy.modules.deform_conv import DeformableConv2d
from ..config import Config
py/BiRefNet/modules/attentions.py → py/BiRefNet_legacy/modules/attentions.py
View File
py/BiRefNet/modules/decoder_blocks.py → py/BiRefNet_legacy/modules/decoder_blocks.py
+2 -2
View File
@@ -1,7 +1,7 @@
import torch
import torch.nn as nn
from BiRefNet.modules.aspp import ASPP, ASPPDeformable
from BiRefNet.modules.attentions import PSA, SGE
from BiRefNet_legacy.modules.aspp import ASPP, ASPPDeformable
from BiRefNet_legacy.modules.attentions import PSA, SGE
from ..config import Config
py/BiRefNet_legacy/modules/deform_conv.py
+66
View File
@@ -0,0 +1,66 @@
import torch
import torch.nn as nn
from torchvision.ops import deform_conv2d
class DeformableConv2d(nn.Module):
def __init__(self,
in_channels,
out_channels,
kernel_size=3,
stride=1,
padding=1,
bias=False):
super(DeformableConv2d, self).__init__()
assert type(kernel_size) == tuple or type(kernel_size) == int
kernel_size = kernel_size if type(kernel_size) == tuple else (kernel_size, kernel_size)
self.stride = stride if type(stride) == tuple else (stride, stride)
self.padding = padding
self.offset_conv = nn.Conv2d(in_channels,
2 * kernel_size[0] * kernel_size[1],
kernel_size=kernel_size,
stride=stride,
padding=self.padding,
bias=True)
nn.init.constant_(self.offset_conv.weight, 0.)
nn.init.constant_(self.offset_conv.bias, 0.)
self.modulator_conv = nn.Conv2d(in_channels,
1 * kernel_size[0] * kernel_size[1],
kernel_size=kernel_size,
stride=stride,
padding=self.padding,
bias=True)
nn.init.constant_(self.modulator_conv.weight, 0.)
nn.init.constant_(self.modulator_conv.bias, 0.)
self.regular_conv = nn.Conv2d(in_channels,
out_channels=out_channels,
kernel_size=kernel_size,
stride=stride,
padding=self.padding,
bias=bias)
def forward(self, x):
#h, w = x.shape[2:]
#max_offset = max(h, w)/4.
offset = self.offset_conv(x)#.clamp(-max_offset, max_offset)
modulator = 2. * torch.sigmoid(self.modulator_conv(x))
x = deform_conv2d(
input=x,
offset=offset,
weight=self.regular_conv.weight,
bias=self.regular_conv.bias,
padding=self.padding,
mask=modulator,
stride=self.stride,
)
return x
py/BiRefNet/modules/ing.py → py/BiRefNet_legacy/modules/ing.py
+1 -1
View File
@@ -1,5 +1,5 @@
import torch.nn as nn
from BiRefNet.modules.mlp import MLPLayer
from BiRefNet_legacy.modules.mlp import MLPLayer
class BlockA(nn.Module):
py/BiRefNet/modules/lateral_blocks.py → py/BiRefNet_legacy/modules/lateral_blocks.py
View File
py/BiRefNet_legacy/modules/mlp.py
+118
View File
@@ -0,0 +1,118 @@
import torch
import torch.nn as nn
from functools import partial
from timm.models.layers import DropPath, to_2tuple, trunc_normal_
from timm.models.registry import register_model
import math
class MLPLayer(nn.Module):
def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
super().__init__()
out_features = out_features or in_features
hidden_features = hidden_features or in_features
self.fc1 = nn.Linear(in_features, hidden_features)
self.act = act_layer()
self.fc2 = nn.Linear(hidden_features, out_features)
self.drop = nn.Dropout(drop)
def forward(self, x):
x = self.fc1(x)
x = self.act(x)
x = self.drop(x)
x = self.fc2(x)
x = self.drop(x)
return x
class Attention(nn.Module):
def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0., sr_ratio=1):
super().__init__()
assert dim % num_heads == 0, f"dim {dim} should be divided by num_heads {num_heads}."
self.dim = dim
self.num_heads = num_heads
head_dim = dim // num_heads
self.scale = qk_scale or head_dim ** -0.5
self.q = nn.Linear(dim, dim, bias=qkv_bias)
self.kv = nn.Linear(dim, dim * 2, bias=qkv_bias)
self.attn_drop = nn.Dropout(attn_drop)
self.proj = nn.Linear(dim, dim)
self.proj_drop = nn.Dropout(proj_drop)
self.sr_ratio = sr_ratio
if sr_ratio > 1:
self.sr = nn.Conv2d(dim, dim, kernel_size=sr_ratio, stride=sr_ratio)
self.norm = nn.LayerNorm(dim)
def forward(self, x, H, W):
B, N, C = x.shape
q = self.q(x).reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)
if self.sr_ratio > 1:
x_ = x.permute(0, 2, 1).reshape(B, C, H, W)
x_ = self.sr(x_).reshape(B, C, -1).permute(0, 2, 1)
x_ = self.norm(x_)
kv = self.kv(x_).reshape(B, -1, 2, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
else:
kv = self.kv(x).reshape(B, -1, 2, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
k, v = kv[0], kv[1]
attn = (q @ k.transpose(-2, -1)) * self.scale
attn = attn.softmax(dim=-1)
attn = self.attn_drop(attn)
x = (attn @ v).transpose(1, 2).reshape(B, N, C)
x = self.proj(x)
x = self.proj_drop(x)
return x
class Block(nn.Module):
def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, sr_ratio=1):
super().__init__()
self.norm1 = norm_layer(dim)
self.attn = Attention(
dim,
num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale,
attn_drop=attn_drop, proj_drop=drop, sr_ratio=sr_ratio)
# NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
self.norm2 = norm_layer(dim)
mlp_hidden_dim = int(dim * mlp_ratio)
self.mlp = MLPLayer(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
def forward(self, x, H, W):
x = x + self.drop_path(self.attn(self.norm1(x), H, W))
x = x + self.drop_path(self.mlp(self.norm2(x), H, W))
return x
class OverlapPatchEmbed(nn.Module):
""" Image to Patch Embedding
"""
def __init__(self, img_size=224, patch_size=7, stride=4, in_channels=3, embed_dim=768):
super().__init__()
img_size = to_2tuple(img_size)
patch_size = to_2tuple(patch_size)
self.img_size = img_size
self.patch_size = patch_size
self.H, self.W = img_size[0] // patch_size[0], img_size[1] // patch_size[1]
self.num_patches = self.H * self.W
self.proj = nn.Conv2d(in_channels, embed_dim, kernel_size=patch_size, stride=stride,
padding=(patch_size[0] // 2, patch_size[1] // 2))
self.norm = nn.LayerNorm(embed_dim)
def forward(self, x):
x = self.proj(x)
_, _, H, W = x.shape
x = x.flatten(2).transpose(1, 2)
x = self.norm(x)
return x, H, W
py/BiRefNet_legacy/modules/utils.py
+54
View File
@@ -0,0 +1,54 @@
import torch.nn as nn
def build_act_layer(act_layer):
if act_layer == 'ReLU':
return nn.ReLU(inplace=True)
elif act_layer == 'SiLU':
return nn.SiLU(inplace=True)
elif act_layer == 'GELU':
return nn.GELU()
raise NotImplementedError(f'build_act_layer does not support {act_layer}')
def build_norm_layer(dim,
norm_layer,
in_format='channels_last',
out_format='channels_last',
eps=1e-6):
layers = []
if norm_layer == 'BN':
if in_format == 'channels_last':
layers.append(to_channels_first())
layers.append(nn.BatchNorm2d(dim))
if out_format == 'channels_last':
layers.append(to_channels_last())
elif norm_layer == 'LN':
if in_format == 'channels_first':
layers.append(to_channels_last())
layers.append(nn.LayerNorm(dim, eps=eps))
if out_format == 'channels_first':
layers.append(to_channels_first())
else:
raise NotImplementedError(
f'build_norm_layer does not support {norm_layer}')
return nn.Sequential(*layers)
class to_channels_first(nn.Module):
def __init__(self):
super().__init__()
def forward(self, x):
return x.permute(0, 3, 1, 2)
class to_channels_last(nn.Module):
def __init__(self):
super().__init__()
def forward(self, x):
return x.permute(0, 2, 3, 1)
py/BiRefNet/preproc.py → py/BiRefNet_legacy/preproc.py
View File
py/BiRefNet_legacy/refinement/__init__.py
View File
py/BiRefNet/refinement/refiner.py → py/BiRefNet_legacy/refinement/refiner.py
+5 -5
View File
@@ -9,11 +9,11 @@ from torchvision.models import resnet50
from ..config import Config
from ..dataset import class_labels_TR_sorted
from BiRefNet.backbones.build_backbone import build_backbone
from BiRefNet.modules.decoder_blocks import BasicDecBlk
from BiRefNet.modules.lateral_blocks import BasicLatBlk
from BiRefNet.modules.ing import *
from BiRefNet.refinement.stem_layer import StemLayer
from BiRefNet_legacy.backbones.build_backbone import build_backbone
from BiRefNet_legacy.modules.decoder_blocks import BasicDecBlk
from BiRefNet_legacy.modules.lateral_blocks import BasicLatBlk
from BiRefNet_legacy.modules.ing import *
from BiRefNet_legacy.refinement.stem_layer import StemLayer
class RefinerPVTInChannels4(nn.Module):
py/BiRefNet_legacy/refinement/stem_layer.py
+45
View File
@@ -0,0 +1,45 @@
import torch.nn as nn
from BiRefNet_legacy.modules.utils import build_act_layer, build_norm_layer
class StemLayer(nn.Module):
r""" Stem layer of InternImage
Args:
in_channels (int): number of input channels
out_channels (int): number of output channels
act_layer (str): activation layer
norm_layer (str): normalization layer
"""
def __init__(self,
in_channels=3+1,
inter_channels=48,
out_channels=96,
act_layer='GELU',
norm_layer='BN'):
super().__init__()
self.conv1 = nn.Conv2d(in_channels,
inter_channels,
kernel_size=3,
stride=1,
padding=1)
self.norm1 = build_norm_layer(
inter_channels, norm_layer, 'channels_first', 'channels_first'
)
self.act = build_act_layer(act_layer)
self.conv2 = nn.Conv2d(inter_channels,
out_channels,
kernel_size=3,
stride=1,
padding=1)
self.norm2 = build_norm_layer(
out_channels, norm_layer, 'channels_first', 'channels_first'
)
def forward(self, x):
x = self.conv1(x)
x = self.norm1(x)
x = self.act(x)
x = self.conv2(x)
x = self.norm2(x)
return x
py/birefnet_func.py → py/birefnet_legacy.py
+2 -3
View File
@@ -2,9 +2,8 @@ from .imagefunc import *
import torch.nn as nn
from torchvision import transforms
from .BiRefNet.baseline import BiRefNet
from .BiRefNet.config import Config
from .BiRefNet_legacy.baseline import BiRefNet
from .BiRefNet_legacy.config import Config
class BiRefNet_img_processor:
def __init__(self, config):
py/birefnet_ultra.py
+1 -1
View File
@@ -40,7 +40,7 @@ class BiRefNetUltra:
else:
local_files_only = False
from .birefnet_func import BiRefNetRemoveBackground
from .birefnet_legacy import BiRefNetRemoveBackground
birefnetrmbg = BiRefNetRemoveBackground()
for i in image:
py/birefnet_ultra_v2.py
+153
View File
@@ -0,0 +1,153 @@
import sys
import torch
from torchvision import transforms
from .imagefunc import *
sys.path.append(os.path.join(os.path.dirname(__file__), 'BiRefNet'))
from .BiRefNet.models.birefnet import BiRefNet
def get_models():
model_path = os.path.join(folder_paths.models_dir, 'BiRefNet', 'pth')
model_ext = [".pth"]
model_dict = get_files(model_path, model_ext)
return model_dict
class LS_LoadBiRefNetModel:
def __init__(self):
self.birefnet = None
self.state_dict = None
@classmethod
def INPUT_TYPES(s):
tmp_list = list(get_models().keys())
model_list = []
if 'BiRefNet-general-epoch_244.pth' in tmp_list:
model_list.append('BiRefNet-general-epoch_244.pth')
tmp_list.remove('BiRefNet-general-epoch_244.pth')
model_list.extend(tmp_list)
return {
"required": {
"model": (model_list,),
},
}
RETURN_TYPES = ("BIREFNET_MODEL",)
RETURN_NAMES = ("birefnet_model",)
FUNCTION = "load_birefnet_model"
CATEGORY = '😺dzNodes/LayerMask'
def load_birefnet_model(self, model):
from .BiRefNet.utils import check_state_dict
model_dict = get_models()
self.birefnet = BiRefNet(bb_pretrained=False)
self.state_dict = torch.load(model_dict[model], map_location='cpu')
self.state_dict = check_state_dict(self.state_dict)
self.birefnet.load_state_dict(self.state_dict)
return (self.birefnet,)
class LS_BiRefNetUltraV2:
def __init__(self):
self.NODE_NAME = 'BiRefNetUltraV2'
@classmethod
def INPUT_TYPES(cls):
method_list = ['VITMatte', 'VITMatte(local)', 'PyMatting', 'GuidedFilter', ]
device_list = ['cuda','cpu']
return {
"required": {
"image": ("IMAGE",),
"birefnet_model":("BIREFNET_MODEL",),
"detail_method": (method_list,),
"detail_erode": ("INT", {"default": 4, "min": 1, "max": 255, "step": 1}),
"detail_dilate": ("INT", {"default": 2, "min": 1, "max": 255, "step": 1}),
"black_point": ("FLOAT", {"default": 0.01, "min": 0.01, "max": 0.98, "step": 0.01, "display": "slider"}),
"white_point": ("FLOAT", {"default": 0.99, "min": 0.02, "max": 0.99, "step": 0.01, "display": "slider"}),
"process_detail": ("BOOLEAN", {"default": False}),
"device": (device_list,),
"max_megapixels": ("FLOAT", {"default": 2.0, "min": 1, "max": 999, "step": 0.1}),
},
"optional": {
}
}
RETURN_TYPES = ("IMAGE", "MASK", )
RETURN_NAMES = ("image", "mask", )
FUNCTION = "birefnet_ultra_v2"
CATEGORY = '😺dzNodes/LayerMask'
def birefnet_ultra_v2(self, image, birefnet_model, detail_method, detail_erode, detail_dilate,
black_point, white_point, process_detail, device, max_megapixels):
ret_images = []
ret_masks = []
inference_image_size = (1024, 1024)
if detail_method == 'VITMatte(local)':
local_files_only = True
else:
local_files_only = False
torch.set_float32_matmul_precision(['high', 'highest'][0])
birefnet_model.to('cuda')
birefnet_model.eval()
for i in image:
i = torch.unsqueeze(i, 0)
orig_image = tensor2pil(i).convert('RGB')
transform_image = transforms.Compose([
transforms.Resize(inference_image_size),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
inference_image = transform_image(orig_image).unsqueeze(0).to('cuda')
# Prediction
with torch.no_grad():
preds = birefnet_model(inference_image)[-1].sigmoid().cpu()
pred = preds[0].squeeze()
pred_pil = transforms.ToPILImage()(pred)
_mask = pred_pil.resize(inference_image_size)
print(f"_mask type is = {type(_mask)}, mode = {_mask.mode}, size={_mask.size}")
resize_sampler = Image.BILINEAR
_mask = _mask.resize(orig_image.size, resize_sampler)
brightness_image = ImageEnhance.Brightness(_mask)
_mask = brightness_image.enhance(factor=1.01)
_mask = image2mask(_mask)
detail_range = detail_erode + detail_dilate
if process_detail:
if detail_method == 'GuidedFilter':
_mask = guided_filter_alpha(i, _mask, detail_range // 6 + 1)
_mask = tensor2pil(histogram_remap(_mask, black_point, white_point))
elif detail_method == 'PyMatting':
_mask = tensor2pil(mask_edge_detail(i, _mask, detail_range // 8 + 1, black_point, white_point))
else:
_trimap = generate_VITMatte_trimap(_mask, detail_erode, detail_dilate)
_mask = generate_VITMatte(orig_image, _trimap, local_files_only=local_files_only, device=device, max_megapixels=max_megapixels)
_mask = tensor2pil(histogram_remap(pil2tensor(_mask), black_point, white_point))
else:
_mask = tensor2pil(_mask)
ret_image = RGB2RGBA(orig_image, _mask.convert('L'))
ret_images.append(pil2tensor(ret_image))
ret_masks.append(image2mask(_mask))
log(f"{self.NODE_NAME} Processed {len(ret_masks)} image(s).", message_type='finish')
return (torch.cat(ret_images, dim=0), torch.cat(ret_masks, dim=0),)
NODE_CLASS_MAPPINGS = {
"LayerMask: BiRefNetUltraV2": LS_BiRefNetUltraV2,
"LayerMask: LoadBiRefNetModel": LS_LoadBiRefNetModel
}
NODE_DISPLAY_NAME_MAPPINGS = {
"LayerMask: BiRefNetUltraV2": "LayerMask: BiRefNet Ultra V2",
"LayerMask: LoadBiRefNetModel": "LayerMask: Load BiRefNet Model"
}
py/imagefunc.py
+106 -66
View File
@@ -2085,37 +2085,10 @@ class AnyType(str):
def __ne__(self, __value: object) -> bool:
return False
'''Constant'''
chop_mode = [
'normal',
'multply',
'screen',
'add',
'subtract',
'difference',
'darker',
'lighter',
'color_burn',
'color_dodge',
'linear_burn',
'linear_dodge',
'overlay',
'soft_light',
'hard_light',
'vivid_light',
'pin_light',
'linear_light',
'hard_mix'
]
# Blend Mode from Virtuoso Pack https://github.com/chrisfreilich/virtuoso-nodes
chop_mode_v2 = list(BLEND_MODES.keys())
'''Load File'''
def get_files(model_path: str, file_ext_list:list) -> dict:
file_list = []
for ext in file_ext_list:
@@ -2126,12 +2099,6 @@ def get_files(model_path: str, file_ext_list:list) -> dict:
files_dict[filename] = file_list[i]
return files_dict
default_lut_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.normpath(__file__))), 'lut')
default_font_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.normpath(__file__))), 'font')
resource_dir_ini_file = os.path.join(os.path.dirname(os.path.dirname(os.path.normpath(__file__))), "resource_dir.ini")
api_key_ini_file = os.path.join(os.path.dirname(os.path.dirname(os.path.normpath(__file__))), "api_key.ini")
custom_size_file = os.path.join(os.path.dirname(os.path.dirname(os.path.normpath(__file__))), "custom_size.ini")
# def load_inference_prompt() -> str:
# inference_prompt_file = os.path.join(os.path.dirname(os.path.dirname(os.path.normpath(__file__))), "resource",
# "inference.prompt")
@@ -2144,6 +2111,7 @@ custom_size_file = os.path.join(os.path.dirname(os.path.dirname(os.path.normpath
# return ''.join(ret_value)
def load_custom_size() -> list:
custom_size_file = os.path.join(os.path.dirname(os.path.dirname(os.path.normpath(__file__))), "custom_size.ini")
ret_value = ['1024 x 1024',
'768 x 512',
'512 x 768',
@@ -2166,6 +2134,7 @@ def load_custom_size() -> list:
return ret_value
def get_api_key(api_name:str) -> str:
api_key_ini_file = os.path.join(os.path.dirname(os.path.dirname(os.path.normpath(__file__))), "api_key.ini")
ret_value = ''
try:
with open(api_key_ini_file, 'r') as f:
@@ -2185,43 +2154,87 @@ def get_api_key(api_name:str) -> str:
return ret_value
try:
with open(resource_dir_ini_file, 'r') as f:
ini = f.readlines()
for line in ini:
if line.startswith('LUT_dir='):
_ldir = line[line.find('=') + 1:].rstrip().lstrip()
if os.path.exists(_ldir):
default_lut_dir = _ldir
else:
log(f'Invalid LUT directory, default to be used. check {resource_dir_ini_file}')
elif line.startswith('FONT_dir='):
_fdir = line[line.find('=') + 1:].rstrip().lstrip()
if os.path.exists(_fdir):
default_font_dir = _fdir
else:
log(f'Invalid FONT directory, default to be used. check {resource_dir_ini_file}')
except Exception as e:
# log(f'Warning: {resource_dir_ini_file} ' + repr(e) + f", default directory to be used. ")
log(f'Warning: {resource_dir_ini_file} not found' + f", default directory to be used. ")
def get_resource_dir() -> list:
default_lut_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.normpath(__file__))), 'lut')
default_font_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.normpath(__file__))), 'font')
resource_dir_ini_file = os.path.join(os.path.dirname(os.path.dirname(os.path.normpath(__file__))),
"resource_dir.ini")
try:
with open(resource_dir_ini_file, 'r') as f:
ini = f.readlines()
for line in ini:
if line.startswith('LUT_dir='):
_ldir = line[line.find('=') + 1:].rstrip().lstrip()
if os.path.exists(_ldir):
default_lut_dir = _ldir
# else:
# log(f'Invalid LUT directory, default to be used. check {resource_dir_ini_file}')
elif line.startswith('FONT_dir='):
_fdir = line[line.find('=') + 1:].rstrip().lstrip()
if os.path.exists(_fdir):
default_font_dir = _fdir
# else:
# log(f'Invalid FONT directory, default to be used. check {resource_dir_ini_file}')
except Exception as e:
# log(f'Warning: {resource_dir_ini_file} ' + repr(e) + f", default directory to be used. ")
log(f'Warning: {resource_dir_ini_file} not found' + f", default directory to be used. ")
__lut_file_list = glob.glob(default_lut_dir + '/*.cube')
LUT_DICT = {}
for i in range(len(__lut_file_list)):
_, __filename = os.path.split(__lut_file_list[i])
LUT_DICT[__filename] = __lut_file_list[i]
LUT_LIST = list(LUT_DICT.keys())
log(f'Find {len(LUT_LIST)} LUTs in {default_lut_dir}')
__lut_file_list = glob.glob(default_lut_dir + '/*.cube')
LUT_DICT = {}
for i in range(len(__lut_file_list)):
_, __filename = os.path.split(__lut_file_list[i])
LUT_DICT[__filename] = __lut_file_list[i]
LUT_LIST = list(LUT_DICT.keys())
log(f'Find {len(LUT_LIST)} LUTs in {default_lut_dir}')
__font_file_list = glob.glob(default_font_dir + '/*.ttf')
__font_file_list.extend(glob.glob(default_font_dir + '/*.otf'))
FONT_DICT = {}
for i in range(len(__font_file_list)):
_, __filename = os.path.split(__font_file_list[i])
FONT_DICT[__filename] = __font_file_list[i]
__font_file_list = glob.glob(default_font_dir + '/*.ttf')
__font_file_list.extend(glob.glob(default_font_dir + '/*.otf'))
FONT_DICT = {}
for i in range(len(__font_file_list)):
_, __filename = os.path.split(__font_file_list[i])
FONT_DICT[__filename] = __font_file_list[i]
FONT_LIST = list(FONT_DICT.keys())
log(f'Find {len(FONT_LIST)} Fonts in {default_font_dir}')
return (LUT_DICT, FONT_DICT)
(LUT_DICT, FONT_DICT) = get_resource_dir()
FONT_LIST = list(FONT_DICT.keys())
log(f'Find {len(FONT_LIST)} Fonts in {default_font_dir}')
LUT_LIST = list(LUT_DICT.keys())
# def get_models_dir() -> dict:
# models_dir_ini_file = os.path.join(os.path.dirname(os.path.dirname(os.path.normpath(__file__))), "models_dir.ini")
# MODELS_DIR = {}
# model_dir_list = [
# "birefnet_dir",
# "evf-sam_dir",
# "florence2_dir",
# "lama_dir",
# "rmbg_dir",
# "segformerB2_dir",
# "segformerB3_clothes_dir",
# "segformerB3_fashion_dir",
# "sam2_dir",
# "transparent-background_dir",
# "yolo8_dir",
# "yolo_world_dir"
# ]
# try:
# with open(models_dir_ini_file, 'r') as f:
# ini = f.readlines()
# for line in ini:
# for model_dir in model_dir_list:
# if line.startswith(model_dir):
# path = line[line.find('=') + 1:].rstrip().lstrip()
# if os.path.exists(path):
# MODELS_DIR[model_dir] = path
# log(f'Find {len(MODELS_DIR)} path(s) in {models_dir_ini_file}.')
# except Exception as e:
# log(f'Warning: {models_dir_ini_file} not found' + f', default directory to be used.')
#
# return MODELS_DIR
#
# MODELS_DIR = get_models_dir()
def draw_bounding_boxes(image: Image, bboxes: list, color: str = "#FF0000", line_width: int = 5) -> Image:
"""
@@ -2250,6 +2263,33 @@ def draw_bounding_boxes(image: Image, bboxes: list, color: str = "#FF0000", line
return image
'''Constant'''
chop_mode = [
'normal',
'multply',
'screen',
'add',
'subtract',
'difference',
'darker',
'lighter',
'color_burn',
'color_dodge',
'linear_burn',
'linear_dodge',
'overlay',
'soft_light',
'hard_light',
'vivid_light',
'pin_light',
'linear_light',
'hard_mix'
]
# Blend Mode from Virtuoso Pack https://github.com/chrisfreilich/virtuoso-nodes
chop_mode_v2 = list(BLEND_MODES.keys())
gemini_generate_config = {
"temperature": 0,
"top_p": 1,
py/lama.py
-1
View File
@@ -39,7 +39,6 @@ class LaMa:
l_images = []
l_masks = []
ret_images = []
print(f"########################mask.shape={mask.shape},dim={mask.dim()}")
for l in image:
l_images.append(torch.unsqueeze(l, 0))
pyproject.toml
+1 -1
View File
@@ -1,7 +1,7 @@
[project]
name = "comfyui_layerstyle"
description = "A set of nodes for ComfyUI it generate image like Adobe Photoshop's Layer Style. the Drop Shadow is first completed node, and follow-up work is in progress."
version = "1.0.45"
version = "1.0.46"
license = "MIT"
dependencies = ["numpy", "pillow", "torch", "matplotlib", "Scipy", "scikit_image", "opencv-contrib-python", "pymatting", "segment_anything", "timm", "addict", "yapf", "colour-science", "wget", "mediapipe", "loguru", "typer_config", "fastapi", "rich", "google-generativeai", "diffusers", "omegaconf", "tqdm", "transformers", "kornia", "image-reward", "ultralytics", "blend_modes", "blind-watermark", "qrcode", "pyzbar", "transparent-background", "huggingface_hub", "accelerate", "bitsandbytes", "torchscale", "wandb", "hydra-core", "psd-tools", "inference-cli[yolo-world]", "inference-gpu[yolo-world]", "onnxruntime"]
repair_dependency.bat
+2
View File
@@ -6,6 +6,8 @@ set "python_exec=..\..\..\python_embeded\python.exe"
echo Fixing Dependency Package...
echo Installing with ComfyUI Portable
%python_exec% -s -m pip install "inference-cli[yolo-world]>=0.13.0"
%python_exec% -s -m pip install "inference-gpu[yolo-world]>=0.13.0"
%python_exec% -s -m pip uninstall -y opencv-python opencv-contrib-python opencv-python-headless opencv-contrib-python-headless
for /f "delims=" %%i in (%requirements_txt%) do (
%python_exec% -s -m pip install -i https://pypi.tuna.tsinghua.edu.cn/simple "%%i"
repair_dependency_aki.bat
+3
View File
@@ -6,7 +6,10 @@ set "python_exec=..\..\python\python.exe"
echo Fixing Dependency Package...
echo Installing with ComfyUI Portable
%python_exec% -s -m pip install "inference-cli[yolo-world]>=0.13.0"
%python_exec% -s -m pip install "inference-gpu[yolo-world]>=0.13.0"
%python_exec% -s -m pip uninstall -y opencv-python opencv-contrib-python opencv-python-headless opencv-contrib-python-headless
for /f "delims=" %%i in (%requirements_txt%) do (
%python_exec% -s -m pip install -i https://pypi.tuna.tsinghua.edu.cn/simple "%%i"
)