Move all BiRefNet github codes to the first level directory.

BiRefNet_github/LICENSE

@@ -1,21 +0,0 @@
-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.

BiRefNet_github/README.md

@@ -1,234 +0,0 @@
-# <p align=center>`Bilateral Reference for High-Resolution Dichotomous Image Segmentation`</p>
-
-|            *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.pdf)" (___arXiv 2024___).
-
-> **Authors:**
-> [Peng Zheng](https://scholar.google.com/citations?user=TZRzWOsAAAAJ),
-> [Dehong Gao](https://scholar.google.com/citations?user=0uPb8MMAAAAJ),
-> [Deng-Ping Fan](https://scholar.google.com/citations?user=kakwJ5QAAAAJ),
-> [Li Liu](https://scholar.google.com/citations?user=9cMQrVsAAAAJ),
-> [Jorma Laaksonen](https://scholar.google.com/citations?user=qQP6WXIAAAAJ),
-> [Wanli Ouyang](https://scholar.google.com/citations?user=pw_0Z_UAAAAJ), &
-> [Nicu Sebe](https://scholar.google.com/citations?user=stFCYOAAAAAJ).
-
-[[**arXiv**](https://arxiv.org/abs/2401.03407)] [[**code**](https://github.com/ZhengPeng7/BiRefNet)] [[**stuff**](https://drive.google.com/drive/folders/1s2Xe0cjq-2ctnJBR24563yMSCOu4CcxM)]  [[**中文版**](https://drive.google.com/file/d/1aBnJ_R9lbnC2dm8dqD0-pzP2Cu-U1Xpt/view?usp=drive_link)] 
-
-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 Papers with Codes (by the time of this work):</summary><p>
-<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 Papers with Codes (by the time of this work):</summary><p>
-<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 Papers with Codes (by the time of this work):</summary><p>
-<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#scrollTo=DJ4meUYjia6S)
-+ **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 **Single Image Inference** on Colab: [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/14Dqg7oeBkFEtchaHLNpig2BcdkZEogba?usp=drive_link)
-<img src="https://drive.google.com/thumbnail?id=12XmDhKtO1o2fEvBu4OE4ULVB2BK0ecWi&sz=w1620" />
-
-
-
-## Model Zoo
-
-> For more general use of our BiRefNet, I managed to extend the original adademic one to more general ones for better application in real life.
->
-> Datasets and datasets are suggested to download from official pages. But you can also download the packaged ones: [DIS](https://drive.google.com/drive/folders/1hZW6tAGPJwo9mPS7qGGGdpxuvuXiyoMJ?usp=drive_link), [HRSOD](https://drive.google.com/drive/folders/18_hAE3QM4cwAzEAKXuSNtKjmgFXTQXZN?usp=drive_link), [COD](https://drive.google.com/drive/folders/1EyHmKWsXfaCR9O0BiZEc3roZbRcs4ECO?usp=drive_link), [Backbones](https://drive.google.com/drive/folders/1cmce_emsS8A5ha5XT2c_CZiJzlLM81ms?usp=drive_link).
->
-> 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><p>
-
-| Task  |        Training Sets        |   Backbone    |                           Download                           |
-| :---: | :-------------------------: | :-----------: | :----------------------------------------------------------: |
-|  DIS  |          DIS5K-TR           | swin_v1_large | [google-drive](https://drive.google.com/file/d/1J90LucvDQaS3R_-9E7QUh1mgJ8eQvccb/view?usp=drive_link) |
-|  COD  |     COD10K-TR, CAMO-TR      | swin_v1_large | [google-drive](https://drive.google.com/file/d/1tM5M72k7a8aKF-dYy-QXaqvfEhbFaWkC/view?usp=drive_link) |
-| HRSOD |           DUTS-TR           | swin_v1_large | [google-drive](https://drive.google.com/file/d/1f7L0Pb1Y3RkOMbqLCW_zO31dik9AiUFa/view?usp=drive_link) |
-| 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?usp=drive_link) |
-| HRSOD |      DUTS-TR, UHRSD-TR      | swin_v1_large | [google-drive](https://drive.google.com/file/d/1Pu1mv3ORobJatIuUoEuZaWDl2ylP3Gw7/view?usp=drive_link) |
-| HRSOD |     HRSOD-TR, UHRSD-TR      | swin_v1_large | [google-drive](https://drive.google.com/file/d/1xEh7fsgWGaS5c3IffMswasv0_u-aVM9E/view?usp=drive_link) |
-| HRSOD | DUTS-TR, HRSOD-TR, UHRSD-TR | swin_v1_large | [google-drive](https://drive.google.com/file/d/13FaxyyOwyCddfZn2vZo1xG1KNZ3cZ-6B/view?usp=drive_link) |
-
-</details>
-
-
-
-<details><summary>Models trained with customed data (massive, portrait), 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 | swin_v1_large |  DIS-VD   |  0.889, 0.840, 1152   | [google-drive](https://drive.google.com/file/d/1KRVE-U3OHrUuuFPY4FFdE4eYBeHJSA0H/view?usp=drive_link) |
-|      **general use**      | DIS5K-TR,DIS-TEs, DUTS-TR_TE,HRSOD-TR_TE,UHRSD-TR_TE, HRS10K-TR_TE | swin_v1_tiny  |  DIS-VD   |  0.867, 0.809, 1182   | [Google-drive](https://drive.google.com/file/d/16gDZISjNp7rKi5vsJm6_fbYF8ZBK8AoF/view?usp=drive_link) |
-|      **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?usp=drive_link) |
-| **portrait segmentation** |                           P3M-10k                            | swin_v1_large | P3M-500-P |     0.982, 0.990      | [google-drive](https://drive.google.com/file/d/1vrjPoOGj05iSxb4MMeznX5k67VlyfZX5/view?usp=drive_link) |
-
-</details>
-
-
-
-<details><summary>Segmentation with box <b>guidance</b>:</summary>
-
-​	*In progress...*
-
-</details>
-
-
-
-<details><summary>Model <b>efficiency</b>:</summary><p>
-
-> 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>
-
-
-
-## 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 [**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 offerring a demo for the [easier single image inference on colab](https://colab.research.google.com/drive/14Dqg7oeBkFEtchaHLNpig2BcdkZEogba?usp=drive_link).
-
-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 dis python=3.9 -y && conda activate dis
-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
-# See train.sh / test.sh for only training / test-evaluation.
-# After the evluation, 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.pdf)).
-+ 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=184e84BwLuNu1FytSAQ2EnANZ0RFHKPip&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={arXiv},
-  year={2024}
-}
-```
-
-
-
-## Contact
-
-Any question, discussion or even complaint, feel free to leave issues here or send me e-mails ([email protected]).

BiRefNet_github/eval_existingOnes.py

@@ -1,139 +0,0 @@
-import os
-import argparse
-from glob import glob
-import prettytable as pt
-
-from evaluation.evaluate import evaluator
-from 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"]
-        elif config.task == 'COD':
-            tb.field_names = ["Dataset", "Method", "Smeasure", "wFmeasure", "meanFm", "meanEm", "maxEm", 'MAE', "maxFm", "adpEm", "adpFm", "HCE"]
-        elif config.task == 'HRSOD':
-            tb.field_names = ["Dataset", "Method", "Smeasure", "maxFm", "meanEm", 'MAE', "maxEm", "meanFm", "wFmeasure", "adpEm", "adpFm", "HCE"]
-        elif config.task == 'DIS5K+HRSOD+HRS10K':
-            tb.field_names = ["Dataset", "Method", "maxFm", "wFmeasure", 'MAE', "Smeasure", "meanEm", "HCE", "maxEm", "meanFm", "adpEm", "adpFm"]
-        elif config.task == 'P3M-10k':
-            tb.field_names = ["Dataset", "Method", "Smeasure", "maxFm", "meanEm", 'MAE', "maxEm", "meanFm", "wFmeasure", "adpEm", "adpFm", "HCE"]
-        else:
-            tb.field_names = ["Dataset", "Method", "Smeasure", 'MAE', "maxEm", "meanEm", "maxFm", "meanFm", "wFmeasure", "adpEm", "adpFm", "HCE"]
-        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 = 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),
-                ]
-            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()),
-                ]
-            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()),
-                ]
-            elif config.task == 'DIS5K+HRSOD+HRS10K':
-                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),
-                ]
-            elif config.task == 'P3M-10k':
-                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()),
-                ]
-            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()),
-                ]
-            
-            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'][:]),
-            'DIS5K+HRSOD+HRS10K': '+'.join(['DIS-VD'][:]),
-            'P3M-10k': '+'.join(['TE-P3M-500-P', 'TE-P3M-500-NP'][:]),
-        }[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', 'HCE'][:100 if 'DIS5K' in config.task else -1]))
-    args = parser.parse_args()
-
-    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)

BiRefNet_github/evaluation/evaluate.py

@@ -1,60 +0,0 @@
-import os
-import prettytable as pt
-
-from evaluation.metrics import evaluator
-from config import Config
-
-
-config = Config()
-
-def evaluate(pred_dir, method, testset, only_S_MAE=False, epoch=0):
-    filename = os.path.join('evaluation', 'eval-{}.txt'.format(method))
-    if os.path.exists(filename):
-        id_suffix = 1
-        filename = filename.rstrip('.txt') + '_{}.txt'.format(id_suffix)
-        while os.path.exists(filename):
-            id_suffix += 1
-            filename = filename.replace('_{}.txt'.format(id_suffix-1), '_{}.txt'.format(id_suffix))
-    gt_paths = sorted([
-        os.path.join(config.data_root_dir, config.task, testset, 'gt', p)
-        for p in os.listdir(os.path.join(config.data_root_dir, config.task, testset, 'gt'))
-    ])
-    pred_paths = sorted([os.path.join(pred_dir, method, testset, p) for p in os.listdir(os.path.join(pred_dir, method, testset))])
-    with open(filename, 'a+') as file_to_write:
-        tb = pt.PrettyTable()
-        field_names = [
-            "Dataset", "Method", "maxFm", "wFmeasure", 'MAE', "Smeasure", "meanEm", "maxEm", "meanFm",
-            "adpEm", "adpFm", 'HCE'
-        ]
-        tb.field_names = [name for name in field_names if not only_S_MAE or all(metric not in name for metric in ['Em', 'Fm'])]
-        em, sm, fm, mae, wfm, hce = evaluator(
-            gt_paths=gt_paths[:],
-            pred_paths=pred_paths[:],
-            metrics=['S', 'MAE', 'E', 'F', 'HCE'][:10*(not only_S_MAE) + 2],    # , 'WF'
-            verbose=config.verbose_eval,
-        )
-        e_max, e_mean, e_adp = em['curve'].max(), em['curve'].mean(), em['adp'].mean()
-        f_max, f_mean, f_wfm, f_adp = fm['curve'].max(), fm['curve'].mean(), wfm, fm['adp']
-        tb.add_row(
-            [
-                method+str(epoch), testset, f_max.round(3), f_wfm.round(3), mae.round(3), sm.round(3),
-                e_mean.round(3), e_max.round(3), f_mean.round(3), em['adp'].round(3), f_adp.round(3), hce.round(3)
-            ] if not only_S_MAE else [method, testset, mae.round(3), sm.round(3)]
-        )
-        print(tb)
-        file_to_write.write(str(tb).replace('+', '|')+'\n')
-        file_to_write.close()
-    return {'e_max': e_max, 'e_mean': e_mean, 'e_adp': e_adp, 'sm': sm, 'mae': mae, 'f_max': f_max, 'f_mean': f_mean, 'f_wfm': f_wfm, 'f_adp': f_adp, 'hce': hce}
-
-
-def main():
-    only_S_MAE = False
-    pred_dir = '.'
-    method = 'tmp_val'
-    testsets = 'DIS-VD+DIS-TE1'
-    for testset in testsets.split('+'):
-        res_dct = evaluate(pred_dir, method, testset, only_S_MAE=only_S_MAE)
-
-
-if __name__ == '__main__':
-    main()

BiRefNet_github/evaluation/metrics.py

@@ -1,612 +0,0 @@
-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', '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 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'
-        if os.path.exists(pred):
-            pred_ary = cv2.imread(pred, cv2.IMREAD_GRAYSCALE)
-        else:
-            pred_ary = cv2.imread(pred.replace('.png', '.jpg'), cv2.IMREAD_GRAYSCALE)
-        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 '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)
-
-    return em, sm, fm, mae, wfm, hce
-
-
-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_

BiRefNet_github/evaluation/valid.py

@@ -1,9 +0,0 @@
-from inference import inference
-from evaluation.evaluate import evaluate
-
-
-def valid(model, data_loader_test, pred_dir, method='tmp_val', testset='DIS-VD', only_S_MAE=True, device=0):
-    model.eval()
-    inference(model, data_loader_test, pred_dir, method, testset, device=device)
-    performance_dict = evaluate(pred_dir, method, testset, only_S_MAE=only_S_MAE, epoch=model.epoch)
-    return performance_dict

BiRefNet_github/gen_best_ep.py

@@ -1,85 +0,0 @@
-import os
-from glob import glob
-import numpy as np
-from 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', 'DIS5K+HRSOD+HRS10K', 'P3M-10k'].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')

BiRefNet_github/inference.py

@@ -1,105 +0,0 @@
-import os
-import argparse
-from glob import glob
-from tqdm import tqdm
-import cv2
-import torch
-
-from dataset import MyData
-from models.birefnet import BiRefNet
-from utils import save_tensor_img, check_state_dict
-from 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',
-                            'DIS5K+HRSOD+HRS10K': 'DIS-VD',
-                            'P3M-10k': 'TE-P3M-500-P+TE-P3M-500-NP',
-                            '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)

BiRefNet_github/loss.py

@@ -1,274 +0,0 @@
-import torch
-from torch import nn
-import torch.nn.functional as F
-from torch.autograd import Variable
-from math import exp
-from 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 '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

BiRefNet_github/make_a_copy.sh

@@ -1,18 +0,0 @@
-#!/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}

BiRefNet_github/requirements.txt

@@ -1,15 +0,0 @@
---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

BiRefNet_github/rm_cache.sh

@@ -1,20 +0,0 @@
-#!/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

BiRefNet_github/sub.sh

@@ -1,19 +0,0 @@
-#!/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 compilers/cuda/11.8
-
-export PYTHONUNBUFFERED=1
-export LD_PRELOAD=/home/bingxing2/apps/compilers/gcc/12.2.0/lib64/libstdc++.so.6
-export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:${HOME}/miniconda3/lib:/home/bingxing2/apps/cudnn/8.4.0.27_cuda11.x/lib
-
-method=${1:-"BSL"}
-devices=${2:-0}
-
-sbatch --nodes=1 -p vip_gpu_ailab -A ai4bio \
---ntasks-per-node=1 \
---gres=gpu:$(($(echo ${devices%%,} | grep -o "," | wc -l)+1)) \
---cpus-per-task=32 \
-./train_test.sh ${method} ${devices}
-
-hostname

BiRefNet_github/test.sh

@@ -1,28 +0,0 @@
-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' ;;
-    "DIS5K+HRSOD+HRS10K") testsets='DIS-VD' ;;
-    "P3M-10k") testsets='TE-P3M-500-P,TE-P3M-500-NP' ;;
-esac
-testsets=(`echo ${testsets} | tr ',' ' '`) && testsets=${testsets[@]}
-
-for testset in ${testsets}; do
-    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)

BiRefNet_github/train.py

@@ -1,377 +0,0 @@
-import os
-import datetime
-import argparse
-import torch
-import torch.nn as nn
-import torch.optim as optim
-from torch.autograd import Variable
-
-from config import Config
-from loss import PixLoss, ClsLoss
-from dataset import MyData
-from models.birefnet import BiRefNet
-from utils import Logger, AverageMeter, set_seed, check_state_dict
-from evaluation.valid import valid
-
-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.IoU_finetune_last_epochs:
-            self.pix_loss.lambdas_pix_last['bce'] *= 0
-            self.pix_loss.lambdas_pix_last['ssim'] *= 1
-            self.pix_loss.lambdas_pix_last['iou'] *= 0.5
-
-        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 validate_model(self, epoch):
-        num_image_testset_all = {'DIS-VD': 470, 'DIS-TE1': 500, 'DIS-TE2': 500, 'DIS-TE3': 500, 'DIS-TE4': 500}
-        num_image_testset = {}
-        for testset in args.testsets:
-            if 'DIS-TE' in testset:
-                num_image_testset[testset] = num_image_testset_all[testset]
-        weighted_scores = {'f_max': 0, 'f_mean': 0, 'f_wfm': 0, 'sm': 0, 'e_max': 0, 'e_mean': 0, 'mae': 0}
-        len_all_data_loaders = 0
-        self.model.epoch = epoch
-        for testset, data_loader_test in self.test_loaders.items():
-            print('Validating {}...'.format(testset))
-            performance_dict = valid(
-                self.model,
-                data_loader_test,
-                pred_dir='.',
-                method=args.ckpt_dir.split('/')[-1] if args.ckpt_dir.split('/')[-1].strip('.').strip('/') else 'tmp_val',
-                testset=testset,
-                only_S_MAE=config.only_S_MAE,
-                device=device
-            )
-            print('Test set: {}:'.format(testset))
-            if config.only_S_MAE:
-                print('Smeasure: {:.4f}, MAE: {:.4f}'.format(
-                    performance_dict['sm'], performance_dict['mae']
-                ))
-            else:
-                print('Fmax: {:.4f}, Fwfm: {:.4f}, Smeasure: {:.4f}, Emean: {:.4f}, MAE: {:.4f}'.format(
-                    performance_dict['f_max'], performance_dict['f_wfm'], performance_dict['sm'], performance_dict['e_mean'], performance_dict['mae']
-                ))
-            if '-TE' in testset:
-                for metric in ['sm', 'mae'] if config.only_S_MAE else ['f_max', 'f_mean', 'f_wfm', 'sm', 'e_max', 'e_mean', 'mae']:
-                    weighted_scores[metric] += performance_dict[metric] * len(data_loader_test)
-                len_all_data_loaders += len(data_loader_test)
-        print('Weighted Scores:')
-        for metric, score in weighted_scores.items():
-            if score:
-                print('\t{}: {:.4f}.'.format(metric, score / len_all_data_loaders))
-
-
-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 config.val_step and epoch >= args.epochs - config.save_last and (args.epochs - epoch) % config.val_step == 0:
-            if to_be_distributed:
-                if get_rank() == 0:
-                    print('Validating at rank-{}...'.format(get_rank()))
-                    trainer.validate_model(epoch)
-            else:
-                trainer.validate_model(epoch)
-    if to_be_distributed:
-        destroy_process_group()
-
-if __name__ == '__main__':
-    main()

BiRefNet_github/train_test.sh

@@ -1,11 +0,0 @@
-#!/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

BiRefNet_github/waiting4eval.py

@@ -1,141 +0,0 @@
-# --------------------------------------------------------
-# Make evaluation along with training. Swith time with space/computation.
-# Licensed under The MIT License [see LICENSE for details]
-# Written by Peng Zheng
-# --------------------------------------------------------
-import os
-from glob import glob
-from time import sleep
-import argparse
-import torch
-
-from config import Config
-from models.birefnet import BiRefNet
-from dataset import MyData
-from evaluation.valid import valid
-
-
-parser = argparse.ArgumentParser(description='')
-parser.add_argument('--cuda_idx', default=-1, type=int)
-parser.add_argument('--val_step', default=5*1, type=int)
-parser.add_argument('--program_id', default=0, type=int)
-# id-th one of this program will evaluate  val_step * N + program_id -th epoch model.
-# Test more models, number of programs == number of GPUs: [models[num_all - program_id_1], models[num_all - program_id_max(n, val_step-1)], ...] programs with id>val_step will speed up the evaluation on (val_step - id)%val_step -th epoch models.
-# Test fastest, only sequentially searched val_step*N -th models -- set all program_id as the same.
-parser.add_argument('--testsets', default='DIS-VD+DIS-TE1+DIS-TE2+DIS-TE3+DIS-TE4', type=str)
-args_eval = parser.parse_args()
-
-args_eval.program_id = (args_eval.val_step - args_eval.program_id) % args_eval.val_step
-
-config = Config()
-config.only_S_MAE = True
-device = 'cpu' if args_eval.cuda_idx < 0 else 'cuda:{}'.format(args_eval.cuda_idx)
-ckpt_dir, testsets = glob(os.path.join('ckpt', '*'))[0], args_eval.testsets
-
-
-def validate_model(model, test_loaders, epoch):
-    num_image_testset_all = {'DIS-VD': 470, 'DIS-TE1': 500, 'DIS-TE2': 500, 'DIS-TE3': 500, 'DIS-TE4': 500}
-    num_image_testset = {}
-    for testset in testsets.split('+'):
-        if 'DIS-TE' in testset:
-            num_image_testset[testset] = num_image_testset_all[testset]
-    weighted_scores = {'f_max': 0, 'sm': 0, 'e_max': 0, 'mae': 0}
-    len_all_data_loaders = 0
-    model.epoch = epoch
-    for testset, data_loader_test in test_loaders.items():
-        print('Validating {}...'.format(testset))
-        performance_dict = valid(
-            model,
-            data_loader_test,
-            pred_dir='.',
-            method=ckpt_dir.split('/')[-1] if ckpt_dir.split('/')[-1].strip('.').strip('/') else 'tmp_val',
-            testset=testset,
-            only_S_MAE=config.only_S_MAE,
-            device=device
-        )
-        print('Test set: {}:'.format(testset))
-        if config.only_S_MAE:
-            print('Smeasure: {:.4f}, MAE: {:.4f}'.format(
-                performance_dict['sm'], performance_dict['mae']
-            ))
-        else:
-            print('Fmax: {:.4f}, Fwfm: {:.4f}, Smeasure: {:.4f}, Emean: {:.4f}, MAE: {:.4f}'.format(
-                performance_dict['f_max'], performance_dict['f_wfm'], performance_dict['sm'], performance_dict['e_mean'], performance_dict['mae']
-            ))
-        if '-TE' in testset:
-            for metric in ['sm', 'mae'] if config.only_S_MAE else ['f_max', 'f_wfm', 'sm', 'e_mean', 'mae']:
-                weighted_scores[metric] += performance_dict[metric] * len(data_loader_test)
-            len_all_data_loaders += len(data_loader_test)
-    print('Weighted Scores:')
-    for metric, score in weighted_scores.items():
-        if score:
-            print('\t{}: {:.4f}.'.format(metric, score / len_all_data_loaders))
-
-@torch.no_grad()
-def main():
-    config = Config()
-    # Dataloader
-    test_loaders = {}
-    for testset in testsets.split('+'):
-        dataset = MyData(
-            datasets=testset,
-            image_size=config.size, is_train=False
-        )
-        _data_loader_test = torch.utils.data.DataLoader(
-            dataset=dataset, batch_size=config.batch_size_valid, num_workers=min(config.num_workers, config.batch_size_valid),
-            pin_memory=device != 'cpu', shuffle=False
-        )
-        print(len(_data_loader_test), "batches of valid dataloader {} have been created.".format(testset))
-        test_loaders[testset] = _data_loader_test
-
-    # Model, 3070MiB GPU memory for inference
-    model = BiRefNet(bb_pretrained=False).to(device)
-    models_evaluated = []
-    continous_sleep_time = 0
-    while True:
-        if (
-            (models_evaluated and continous_sleep_time > 60*60*2) or
-            (not models_evaluated and continous_sleep_time > 60*60*24)
-        ):
-            # If no ckpt has been saved, we wait for 24h;
-            # elif some ckpts have been saved, we wait for 2h for new ones;
-            # else: exit this waiting.
-            print('Exiting the waiting for evaluation.')
-            break
-        models_evaluated_record = 'tmp_models_evaluated.txt'
-        if os.path.exists(models_evaluated_record):
-            with open(models_evaluated_record, 'r') as f:
-                models_evaluated_global = f.read().splitlines()
-        else:
-            models_evaluated_global = []
-        models_detected = [
-            m for idx_m, m in enumerate(sorted(
-                glob(os.path.join(ckpt_dir, '*.pth')),
-                key=lambda x: int(x.rstrip('.pth').split('epoch_')[-1]), reverse=True
-            )) if idx_m % args_eval.val_step == args_eval.program_id and m not in models_evaluated + models_evaluated_global
-        ]
-        if models_detected:
-            from time import time
-            time_st = time()
-            # register the evaluated models
-            model_not_evaluated_latest = models_detected[0]
-            with open('tmp_models_evaluated.txt', 'a') as f:
-                f.write(model_not_evaluated_latest + '\n')
-            models_evaluated.append(model_not_evaluated_latest)
-            print('Loading {} for validation...'.format(model_not_evaluated_latest))
-
-            # evaluate the current model
-            state_dict = torch.load(model_not_evaluated_latest, map_location=device)
-            model.load_state_dict(state_dict, strict=False)
-            validate_model(model, test_loaders, int(model_not_evaluated_latest.rstrip('.pth').split('epoch_')[-1]))
-            continous_sleep_time = 0
-            print('Duration of this evaluation:', time() - time_st)
-        else:
-            sleep_interval = 60 * 2
-            sleep(sleep_interval)
-            continous_sleep_time += sleep_interval
-            continue
-
-
-if __name__ == '__main__':
-    main()

config.json

@@ -5,7 +5,7 @@
   ],
   "auto_map": {
     "AutoConfig": "BiRefNet_config.BiRefNetConfig",
-    "AutoModelForImageSegmentation": "BiRefNet_github.models.birefnet.BiRefNet"
+    "AutoModelForImageSegmentation": "birefnet.BiRefNet"
   },
   "custom_pipelines": {
     "image-segmentation": {