ZhengPeng7 commited on
Commit
45e89f1
β€’
1 Parent(s): d73d557

Move all BiRefNet github codes to the first level directory.

Browse files
Files changed (38) hide show
  1. BiRefNet_github/LICENSE +0 -21
  2. BiRefNet_github/README.md +0 -234
  3. BiRefNet_github/eval_existingOnes.py +0 -139
  4. BiRefNet_github/evaluation/evaluate.py +0 -60
  5. BiRefNet_github/evaluation/metrics.py +0 -612
  6. BiRefNet_github/evaluation/valid.py +0 -9
  7. BiRefNet_github/gen_best_ep.py +0 -85
  8. BiRefNet_github/inference.py +0 -105
  9. BiRefNet_github/loss.py +0 -274
  10. BiRefNet_github/make_a_copy.sh +0 -18
  11. BiRefNet_github/requirements.txt +0 -15
  12. BiRefNet_github/rm_cache.sh +0 -20
  13. BiRefNet_github/sub.sh +0 -19
  14. BiRefNet_github/test.sh +0 -28
  15. BiRefNet_github/train.py +0 -377
  16. BiRefNet_github/train_test.sh +0 -11
  17. BiRefNet_github/waiting4eval.py +0 -141
  18. BiRefNet_github/models/birefnet.py β†’ birefnet.py +0 -0
  19. config.json +1 -1
  20. BiRefNet_github/config.py β†’ config.py +0 -0
  21. BiRefNet_github/dataset.py β†’ dataset.py +0 -0
  22. {BiRefNet_github/models β†’ models}/backbones/build_backbone.py +0 -0
  23. {BiRefNet_github/models β†’ models}/backbones/pvt_v2.py +0 -0
  24. {BiRefNet_github/models β†’ models}/backbones/swin_v1.py +0 -0
  25. {BiRefNet_github/models β†’ models}/modules/aspp.py +0 -0
  26. {BiRefNet_github/models β†’ models}/modules/attentions.py +0 -0
  27. {BiRefNet_github/models β†’ models}/modules/decoder_blocks.py +0 -0
  28. {BiRefNet_github/models β†’ models}/modules/deform_conv.py +0 -0
  29. {BiRefNet_github/models β†’ models}/modules/ing.py +0 -0
  30. {BiRefNet_github/models β†’ models}/modules/lateral_blocks.py +0 -0
  31. {BiRefNet_github/models β†’ models}/modules/mlp.py +0 -0
  32. {BiRefNet_github/models β†’ models}/modules/prompt_encoder.py +0 -0
  33. {BiRefNet_github/models β†’ models}/modules/utils.py +0 -0
  34. {BiRefNet_github/models β†’ models}/refinement/refiner.py +0 -0
  35. {BiRefNet_github/models β†’ models}/refinement/stem_layer.py +0 -0
  36. BiRefNet_github/preproc.py β†’ preproc.py +0 -0
  37. BiRefNet_github/train.sh β†’ train.sh +0 -0
  38. BiRefNet_github/utils.py β†’ utils.py +0 -0
BiRefNet_github/LICENSE DELETED
@@ -1,21 +0,0 @@
1
- MIT License
2
-
3
- Copyright (c) 2024 ZhengPeng
4
-
5
- Permission is hereby granted, free of charge, to any person obtaining a copy
6
- of this software and associated documentation files (the "Software"), to deal
7
- in the Software without restriction, including without limitation the rights
8
- to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
- copies of the Software, and to permit persons to whom the Software is
10
- furnished to do so, subject to the following conditions:
11
-
12
- The above copyright notice and this permission notice shall be included in all
13
- copies or substantial portions of the Software.
14
-
15
- THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
- IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
- FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
- AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
- LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
- OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
- SOFTWARE.
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
BiRefNet_github/README.md DELETED
@@ -1,234 +0,0 @@
1
- # <p align=center>`Bilateral Reference for High-Resolution Dichotomous Image Segmentation`</p>
2
-
3
- | *DIS-Sample_1* | *DIS-Sample_2* |
4
- | :------------------------------: | :-------------------------------: |
5
- | <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" /> |
6
-
7
- This repo is the official implementation of "[**Bilateral Reference for High-Resolution Dichotomous Image Segmentation**](https://arxiv.org/pdf/2401.03407.pdf)" (___arXiv 2024___).
8
-
9
- > **Authors:**
10
- > [Peng Zheng](https://scholar.google.com/citations?user=TZRzWOsAAAAJ),
11
- > [Dehong Gao](https://scholar.google.com/citations?user=0uPb8MMAAAAJ),
12
- > [Deng-Ping Fan](https://scholar.google.com/citations?user=kakwJ5QAAAAJ),
13
- > [Li Liu](https://scholar.google.com/citations?user=9cMQrVsAAAAJ),
14
- > [Jorma Laaksonen](https://scholar.google.com/citations?user=qQP6WXIAAAAJ),
15
- > [Wanli Ouyang](https://scholar.google.com/citations?user=pw_0Z_UAAAAJ), &
16
- > [Nicu Sebe](https://scholar.google.com/citations?user=stFCYOAAAAAJ).
17
-
18
- [[**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)]
19
-
20
- Our BiRefNet has achieved SOTA on many similar HR tasks:
21
-
22
- **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)
23
-
24
- <details><summary>Figure of Comparison on Papers with Codes (by the time of this work):</summary><p>
25
- <img src="https://drive.google.com/thumbnail?id=1DLt6CFXdT1QSWDj_6jRkyZINXZ4vmyRp&sz=w1620" />
26
- <img src="https://drive.google.com/thumbnail?id=1gn5GyKFlJbMIkre1JyEdHDSYcrFmcLD0&sz=w1620" />
27
- <img src="https://drive.google.com/thumbnail?id=16CVYYOtafEeZhHqv0am2Daku1n_exMP6&sz=w1620" />
28
- <img src="https://drive.google.com/thumbnail?id=10K45xwPXmaTG4Ex-29ss9payA9yBnyLn&sz=w1620" />
29
- <img src="https://drive.google.com/thumbnail?id=16EuyqKFJOqwMmagvfnbC9hUurL9pYLLB&sz=w1620" />
30
- </details>
31
- <br />
32
-
33
- **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)
34
-
35
- <details><summary>Figure of Comparison on Papers with Codes (by the time of this work):</summary><p>
36
- <img src="https://drive.google.com/thumbnail?id=1DLt6CFXdT1QSWDj_6jRkyZINXZ4vmyRp&sz=w1620" />
37
- <img src="https://drive.google.com/thumbnail?id=1gn5GyKFlJbMIkre1JyEdHDSYcrFmcLD0&sz=w1620" />
38
- <img src="https://drive.google.com/thumbnail?id=16CVYYOtafEeZhHqv0am2Daku1n_exMP6&sz=w1620" />
39
- </details>
40
- <br />
41
-
42
- **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)
43
-
44
- <details><summary>Figure of Comparison on Papers with Codes (by the time of this work):</summary><p>
45
- <img src="https://drive.google.com/thumbnail?id=1hNfQtlTAHT4-AVbk_47852zyRp1NOFLs&sz=w1620" />
46
- <img src="https://drive.google.com/thumbnail?id=1bcVldUAxYkMI3OMTyaP_jNuOugDfYj-d&sz=w1620" />
47
- <img src="https://drive.google.com/thumbnail?id=1p1zgyVz27cGEqQMtOKzm_6zoYK3Sw_Zk&sz=w1620" />
48
- <img src="https://drive.google.com/thumbnail?id=1TubAvcoEbH_mHu3I-AxflnB71nkf35jJ&sz=w1620" />
49
- <img src="https://drive.google.com/thumbnail?id=1A3V9HjVtcMQdnGPwuy-DBVhwKuo0q2lT&sz=w1620" />
50
- </details>
51
- <br />
52
-
53
- #### Try our online demos for inference:
54
-
55
- + **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)
56
- + **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)
57
- + 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)
58
- <img src="https://drive.google.com/thumbnail?id=12XmDhKtO1o2fEvBu4OE4ULVB2BK0ecWi&sz=w1620" />
59
-
60
-
61
-
62
- ## Model Zoo
63
-
64
- > 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.
65
- >
66
- > 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).
67
- >
68
- > Find performances (almost all metrics) of all models in the `exp-TASK_SETTINGS` folders in [[**stuff**](https://drive.google.com/drive/folders/1s2Xe0cjq-2ctnJBR24563yMSCOu4CcxM)].
69
-
70
-
71
-
72
- <details><summary>Models in the original paper, for <b>comparison on benchmarks</b>:</summary><p>
73
-
74
- | Task | Training Sets | Backbone | Download |
75
- | :---: | :-------------------------: | :-----------: | :----------------------------------------------------------: |
76
- | DIS | DIS5K-TR | swin_v1_large | [google-drive](https://drive.google.com/file/d/1J90LucvDQaS3R_-9E7QUh1mgJ8eQvccb/view?usp=drive_link) |
77
- | COD | COD10K-TR, CAMO-TR | swin_v1_large | [google-drive](https://drive.google.com/file/d/1tM5M72k7a8aKF-dYy-QXaqvfEhbFaWkC/view?usp=drive_link) |
78
- | HRSOD | DUTS-TR | swin_v1_large | [google-drive](https://drive.google.com/file/d/1f7L0Pb1Y3RkOMbqLCW_zO31dik9AiUFa/view?usp=drive_link) |
79
- | HRSOD | HRSOD-TR | swin_v1_large | google-drive |
80
- | HRSOD | UHRSD-TR | swin_v1_large | google-drive |
81
- | HRSOD | DUTS-TR, HRSOD-TR | swin_v1_large | [google-drive](https://drive.google.com/file/d/1WJooyTkhoDLllaqwbpur_9Hle0XTHEs_/view?usp=drive_link) |
82
- | HRSOD | DUTS-TR, UHRSD-TR | swin_v1_large | [google-drive](https://drive.google.com/file/d/1Pu1mv3ORobJatIuUoEuZaWDl2ylP3Gw7/view?usp=drive_link) |
83
- | HRSOD | HRSOD-TR, UHRSD-TR | swin_v1_large | [google-drive](https://drive.google.com/file/d/1xEh7fsgWGaS5c3IffMswasv0_u-aVM9E/view?usp=drive_link) |
84
- | HRSOD | DUTS-TR, HRSOD-TR, UHRSD-TR | swin_v1_large | [google-drive](https://drive.google.com/file/d/13FaxyyOwyCddfZn2vZo1xG1KNZ3cZ-6B/view?usp=drive_link) |
85
-
86
- </details>
87
-
88
-
89
-
90
- <details><summary>Models trained with customed data (massive, portrait), for <b>general use in practical application</b>:</summary>
91
-
92
- | Task | Training Sets | Backbone | Test Set | Metric (S, wF[, HCE]) | Download |
93
- | :-----------------------: | :----------------------------------------------------------: | :-----------: | :-------: | :-------------------: | :----------------------------------------------------------: |
94
- | **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) |
95
- | **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) |
96
- | **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) |
97
- | **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) |
98
-
99
- </details>
100
-
101
-
102
-
103
- <details><summary>Segmentation with box <b>guidance</b>:</summary>
104
-
105
- ​ *In progress...*
106
-
107
- </details>
108
-
109
-
110
-
111
- <details><summary>Model <b>efficiency</b>:</summary><p>
112
-
113
- > Screenshot from the original paper. All tests are conducted on a single A100 GPU.
114
-
115
- <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" />
116
-
117
- </details>
118
-
119
-
120
-
121
- ## Third-Party Creations
122
-
123
- > 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.
124
-
125
- We found there've been some 3rd party applications based on our BiRefNet. Many thanks for their contribution to the community!
126
- Choose the one you like to try with clicks instead of codes:
127
- 1. **Applications**:
128
- + 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.
129
- <p align="center"><img src="https://drive.google.com/thumbnail?id=1rNk81YV_Pzb2GykrzfGvX6T7KBXR0wrA&sz=w1620" /></p>
130
-
131
- + 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**.
132
- <p align="center"><img src="https://drive.google.com/thumbnail?id=1GOqEreyS7ENzTPN0RqxEjaA76RpMlkYM&sz=w1620" /></p>
133
-
134
- <https://github.com/ZhengPeng7/BiRefNet/assets/25921713/3a1c7ab2-9847-4dac-8935-43a2d3cd2671>
135
-
136
- + 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.
137
- <p align="center"><img src="https://drive.google.com/thumbnail?id=1KfxCQUUa2y9T-aysEaeVVjCUt3Z0zSkL&sz=w1620" /></p>
138
-
139
- + 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).
140
-
141
- 2. **More Visual Comparisons**
142
- + Thanks [**twitter.com/ZHOZHO672070**](https://twitter.com/ZHOZHO672070) for the comparison with more background-removal methods in images:
143
-
144
- <img src="https://drive.google.com/thumbnail?id=1nvVIFt_Ezs-crPSQxUDqkUBz598fTe63&sz=w1620" />
145
-
146
- + Thanks [**twitter.com/toyxyz3**](https://twitter.com/toyxyz3) for the comparison with more background-removal methods in videos:
147
-
148
- <https://github.com/ZhengPeng7/BiRefNet/assets/25921713/40136198-01cc-4106-81f9-81c985f02e31>
149
-
150
- <https://github.com/ZhengPeng7/BiRefNet/assets/25921713/1a32860c-0893-49dd-b557-c2e35a83c160>
151
-
152
-
153
- ## Usage
154
-
155
- #### Environment Setup
156
-
157
- ```shell
158
- # PyTorch==2.0.1 is used for faster training with compilation.
159
- conda create -n dis python=3.9 -y && conda activate dis
160
- pip install -r requirements.txt
161
- ```
162
-
163
- #### Dataset Preparation
164
-
165
- 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).
166
-
167
- #### Weights Preparation
168
-
169
- Download backbone weights from [my google-drive folder](https://drive.google.com/drive/folders/1s2Xe0cjq-2ctnJBR24563yMSCOu4CcxM) or their official pages.
170
-
171
- #### Run
172
-
173
- ```shell
174
- # Train & Test & Evaluation
175
- ./train_test.sh RUN_NAME GPU_NUMBERS_FOR_TRAINING GPU_NUMBERS_FOR_TEST
176
- # See train.sh / test.sh for only training / test-evaluation.
177
- # 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)).
178
- ```
179
-
180
- #### Well-trained weights:
181
-
182
- 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.
183
-
184
- You can also download the weights from the release of this repo.
185
-
186
- 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.)
187
-
188
- 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 :)
189
-
190
- #### Some of my messages:
191
-
192
- 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`):
193
- + Multi-GPU training: open/close with one variable.
194
- + Backbone choices: Swin_v1, PVT_v2, ConvNets, ...
195
- + Weighted losses: BCE, IoU, SSIM, MAE, Reg, ...
196
- + Adversarial loss for binary segmentation (proposed in my previous work [MCCL](https://arxiv.org/pdf/2302.14485.pdf)).
197
- + Training tricks: multi-scale supervision, freezing backbone, multi-scale input...
198
- + Data collator: loading all in memory, smooth combination of different datasets for combined training and test.
199
- + ...
200
- I really hope you enjoy this project and use it in more works to achieve new SOTAs.
201
-
202
-
203
- ### Quantitative Results
204
-
205
- <p align="center"><img src="https://drive.google.com/thumbnail?id=184e84BwLuNu1FytSAQ2EnANZ0RFHKPip&sz=w1620" /></p>
206
-
207
- <p align="center"><img src="https://drive.google.com/thumbnail?id=1W0mi0ZiYbqsaGuohNXU8Gh7Zj4M3neFg&sz=w1620" /></p>
208
-
209
-
210
-
211
- ### Qualitative Results
212
-
213
- <p align="center"><img src="https://drive.google.com/thumbnail?id=1TYZF8pVZc2V0V6g3ik4iAr9iKvJ8BNrf&sz=w1620" /></p>
214
-
215
- <p align="center"><img src="https://drive.google.com/thumbnail?id=1ZGHC32CAdT9cwRloPzOCKWCrVQZvUAlJ&sz=w1620" /></p>
216
-
217
-
218
-
219
- ### Citation
220
-
221
- ```
222
- @article{zheng2024birefnet,
223
- title={Bilateral Reference for High-Resolution Dichotomous Image Segmentation},
224
- author={Zheng, Peng and Gao, Dehong and Fan, Deng-Ping and Liu, Li and Laaksonen, Jorma and Ouyang, Wanli and Sebe, Nicu},
225
- journal={arXiv},
226
- year={2024}
227
- }
228
- ```
229
-
230
-
231
-
232
- ## Contact
233
-
234
- Any question, discussion or even complaint, feel free to leave issues here or send me e-mails ([email protected]).
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
BiRefNet_github/eval_existingOnes.py DELETED
@@ -1,139 +0,0 @@
1
- import os
2
- import argparse
3
- from glob import glob
4
- import prettytable as pt
5
-
6
- from evaluation.evaluate import evaluator
7
- from config import Config
8
-
9
-
10
- config = Config()
11
-
12
-
13
- def do_eval(args):
14
- # evaluation for whole dataset
15
- # dataset first in evaluation
16
- for _data_name in args.data_lst.split('+'):
17
- pred_data_dir = sorted(glob(os.path.join(args.pred_root, args.model_lst[0], _data_name)))
18
- if not pred_data_dir:
19
- print('Skip dataset {}.'.format(_data_name))
20
- continue
21
- gt_src = os.path.join(args.gt_root, _data_name)
22
- gt_paths = sorted(glob(os.path.join(gt_src, 'gt', '*')))
23
- print('#' * 20, _data_name, '#' * 20)
24
- filename = os.path.join(args.save_dir, '{}_eval.txt'.format(_data_name))
25
- tb = pt.PrettyTable()
26
- tb.vertical_char = '&'
27
- if config.task == 'DIS5K':
28
- tb.field_names = ["Dataset", "Method", "maxFm", "wFmeasure", 'MAE', "Smeasure", "meanEm", "HCE", "maxEm", "meanFm", "adpEm", "adpFm"]
29
- elif config.task == 'COD':
30
- tb.field_names = ["Dataset", "Method", "Smeasure", "wFmeasure", "meanFm", "meanEm", "maxEm", 'MAE', "maxFm", "adpEm", "adpFm", "HCE"]
31
- elif config.task == 'HRSOD':
32
- tb.field_names = ["Dataset", "Method", "Smeasure", "maxFm", "meanEm", 'MAE', "maxEm", "meanFm", "wFmeasure", "adpEm", "adpFm", "HCE"]
33
- elif config.task == 'DIS5K+HRSOD+HRS10K':
34
- tb.field_names = ["Dataset", "Method", "maxFm", "wFmeasure", 'MAE', "Smeasure", "meanEm", "HCE", "maxEm", "meanFm", "adpEm", "adpFm"]
35
- elif config.task == 'P3M-10k':
36
- tb.field_names = ["Dataset", "Method", "Smeasure", "maxFm", "meanEm", 'MAE', "maxEm", "meanFm", "wFmeasure", "adpEm", "adpFm", "HCE"]
37
- else:
38
- tb.field_names = ["Dataset", "Method", "Smeasure", 'MAE', "maxEm", "meanEm", "maxFm", "meanFm", "wFmeasure", "adpEm", "adpFm", "HCE"]
39
- for _model_name in args.model_lst[:]:
40
- print('\t', 'Evaluating model: {}...'.format(_model_name))
41
- pred_paths = [p.replace(args.gt_root, os.path.join(args.pred_root, _model_name)).replace('/gt/', '/') for p in gt_paths]
42
- # print(pred_paths[:1], gt_paths[:1])
43
- em, sm, fm, mae, wfm, hce = evaluator(
44
- gt_paths=gt_paths,
45
- pred_paths=pred_paths,
46
- metrics=args.metrics.split('+'),
47
- verbose=config.verbose_eval
48
- )
49
- if config.task == 'DIS5K':
50
- scores = [
51
- fm['curve'].max().round(3), wfm.round(3), mae.round(3), sm.round(3), em['curve'].mean().round(3), int(hce.round()),
52
- em['curve'].max().round(3), fm['curve'].mean().round(3), em['adp'].round(3), fm['adp'].round(3),
53
- ]
54
- elif config.task == 'COD':
55
- scores = [
56
- sm.round(3), wfm.round(3), fm['curve'].mean().round(3), em['curve'].mean().round(3), em['curve'].max().round(3), mae.round(3),
57
- fm['curve'].max().round(3), em['adp'].round(3), fm['adp'].round(3), int(hce.round()),
58
- ]
59
- elif config.task == 'HRSOD':
60
- scores = [
61
- sm.round(3), fm['curve'].max().round(3), em['curve'].mean().round(3), mae.round(3),
62
- em['curve'].max().round(3), fm['curve'].mean().round(3), wfm.round(3), em['adp'].round(3), fm['adp'].round(3), int(hce.round()),
63
- ]
64
- elif config.task == 'DIS5K+HRSOD+HRS10K':
65
- scores = [
66
- fm['curve'].max().round(3), wfm.round(3), mae.round(3), sm.round(3), em['curve'].mean().round(3), int(hce.round()),
67
- em['curve'].max().round(3), fm['curve'].mean().round(3), em['adp'].round(3), fm['adp'].round(3),
68
- ]
69
- elif config.task == 'P3M-10k':
70
- scores = [
71
- sm.round(3), fm['curve'].max().round(3), em['curve'].mean().round(3), mae.round(3),
72
- em['curve'].max().round(3), fm['curve'].mean().round(3), wfm.round(3), em['adp'].round(3), fm['adp'].round(3), int(hce.round()),
73
- ]
74
- else:
75
- scores = [
76
- sm.round(3), mae.round(3), em['curve'].max().round(3), em['curve'].mean().round(3),
77
- fm['curve'].max().round(3), fm['curve'].mean().round(3), wfm.round(3),
78
- em['adp'].round(3), fm['adp'].round(3), int(hce.round()),
79
- ]
80
-
81
- for idx_score, score in enumerate(scores):
82
- scores[idx_score] = '.' + format(score, '.3f').split('.')[-1] if score <= 1 else format(score, '<4')
83
- records = [_data_name, _model_name] + scores
84
- tb.add_row(records)
85
- # Write results after every check.
86
- with open(filename, 'w+') as file_to_write:
87
- file_to_write.write(str(tb)+'\n')
88
- print(tb)
89
-
90
-
91
- if __name__ == '__main__':
92
- # set parameters
93
- parser = argparse.ArgumentParser()
94
- parser.add_argument(
95
- '--gt_root', type=str, help='ground-truth root',
96
- default=os.path.join(config.data_root_dir, config.task))
97
- parser.add_argument(
98
- '--pred_root', type=str, help='prediction root',
99
- default='./e_preds')
100
- parser.add_argument(
101
- '--data_lst', type=str, help='test dataset',
102
- default={
103
- 'DIS5K': '+'.join(['DIS-VD', 'DIS-TE1', 'DIS-TE2', 'DIS-TE3', 'DIS-TE4'][:]),
104
- 'COD': '+'.join(['TE-COD10K', 'NC4K', 'TE-CAMO', 'CHAMELEON'][:]),
105
- 'HRSOD': '+'.join(['DAVIS-S', 'TE-HRSOD', 'TE-UHRSD', 'TE-DUTS', 'DUT-OMRON'][:]),
106
- 'DIS5K+HRSOD+HRS10K': '+'.join(['DIS-VD'][:]),
107
- 'P3M-10k': '+'.join(['TE-P3M-500-P', 'TE-P3M-500-NP'][:]),
108
- }[config.task])
109
- parser.add_argument(
110
- '--save_dir', type=str, help='candidate competitors',
111
- default='e_results')
112
- parser.add_argument(
113
- '--check_integrity', type=bool, help='whether to check the file integrity',
114
- default=False)
115
- parser.add_argument(
116
- '--metrics', type=str, help='candidate competitors',
117
- default='+'.join(['S', 'MAE', 'E', 'F', 'WF', 'HCE'][:100 if 'DIS5K' in config.task else -1]))
118
- args = parser.parse_args()
119
-
120
- os.makedirs(args.save_dir, exist_ok=True)
121
- try:
122
- 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]
123
- except:
124
- args.model_lst = [m for m in sorted(os.listdir(args.pred_root))]
125
-
126
- # check the integrity of each candidates
127
- if args.check_integrity:
128
- for _data_name in args.data_lst.split('+'):
129
- for _model_name in args.model_lst:
130
- gt_pth = os.path.join(args.gt_root, _data_name)
131
- pred_pth = os.path.join(args.pred_root, _model_name, _data_name)
132
- if not sorted(os.listdir(gt_pth)) == sorted(os.listdir(pred_pth)):
133
- print(len(sorted(os.listdir(gt_pth))), len(sorted(os.listdir(pred_pth))))
134
- print('The {} Dataset of {} Model is not matching to the ground-truth'.format(_data_name, _model_name))
135
- else:
136
- print('>>> skip check the integrity of each candidates')
137
-
138
- # start engine
139
- do_eval(args)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
BiRefNet_github/evaluation/evaluate.py DELETED
@@ -1,60 +0,0 @@
1
- import os
2
- import prettytable as pt
3
-
4
- from evaluation.metrics import evaluator
5
- from config import Config
6
-
7
-
8
- config = Config()
9
-
10
- def evaluate(pred_dir, method, testset, only_S_MAE=False, epoch=0):
11
- filename = os.path.join('evaluation', 'eval-{}.txt'.format(method))
12
- if os.path.exists(filename):
13
- id_suffix = 1
14
- filename = filename.rstrip('.txt') + '_{}.txt'.format(id_suffix)
15
- while os.path.exists(filename):
16
- id_suffix += 1
17
- filename = filename.replace('_{}.txt'.format(id_suffix-1), '_{}.txt'.format(id_suffix))
18
- gt_paths = sorted([
19
- os.path.join(config.data_root_dir, config.task, testset, 'gt', p)
20
- for p in os.listdir(os.path.join(config.data_root_dir, config.task, testset, 'gt'))
21
- ])
22
- pred_paths = sorted([os.path.join(pred_dir, method, testset, p) for p in os.listdir(os.path.join(pred_dir, method, testset))])
23
- with open(filename, 'a+') as file_to_write:
24
- tb = pt.PrettyTable()
25
- field_names = [
26
- "Dataset", "Method", "maxFm", "wFmeasure", 'MAE', "Smeasure", "meanEm", "maxEm", "meanFm",
27
- "adpEm", "adpFm", 'HCE'
28
- ]
29
- 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'])]
30
- em, sm, fm, mae, wfm, hce = evaluator(
31
- gt_paths=gt_paths[:],
32
- pred_paths=pred_paths[:],
33
- metrics=['S', 'MAE', 'E', 'F', 'HCE'][:10*(not only_S_MAE) + 2], # , 'WF'
34
- verbose=config.verbose_eval,
35
- )
36
- e_max, e_mean, e_adp = em['curve'].max(), em['curve'].mean(), em['adp'].mean()
37
- f_max, f_mean, f_wfm, f_adp = fm['curve'].max(), fm['curve'].mean(), wfm, fm['adp']
38
- tb.add_row(
39
- [
40
- method+str(epoch), testset, f_max.round(3), f_wfm.round(3), mae.round(3), sm.round(3),
41
- e_mean.round(3), e_max.round(3), f_mean.round(3), em['adp'].round(3), f_adp.round(3), hce.round(3)
42
- ] if not only_S_MAE else [method, testset, mae.round(3), sm.round(3)]
43
- )
44
- print(tb)
45
- file_to_write.write(str(tb).replace('+', '|')+'\n')
46
- file_to_write.close()
47
- 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}
48
-
49
-
50
- def main():
51
- only_S_MAE = False
52
- pred_dir = '.'
53
- method = 'tmp_val'
54
- testsets = 'DIS-VD+DIS-TE1'
55
- for testset in testsets.split('+'):
56
- res_dct = evaluate(pred_dir, method, testset, only_S_MAE=only_S_MAE)
57
-
58
-
59
- if __name__ == '__main__':
60
- main()
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
BiRefNet_github/evaluation/metrics.py DELETED
@@ -1,612 +0,0 @@
1
- import os
2
- from tqdm import tqdm
3
- import cv2
4
- import numpy as np
5
- from scipy.ndimage import convolve, distance_transform_edt as bwdist
6
- from skimage.morphology import skeletonize
7
- from skimage.morphology import disk
8
- from skimage.measure import label
9
-
10
-
11
- _EPS = np.spacing(1)
12
- _TYPE = np.float64
13
-
14
-
15
- def evaluator(gt_paths, pred_paths, metrics=['S', 'MAE', 'E', 'F', 'WF', 'HCE'], verbose=False):
16
- # define measures
17
- if 'E' in metrics:
18
- EM = Emeasure()
19
- if 'S' in metrics:
20
- SM = Smeasure()
21
- if 'F' in metrics:
22
- FM = Fmeasure()
23
- if 'MAE' in metrics:
24
- MAE = MAEmeasure()
25
- if 'WF' in metrics:
26
- WFM = WeightedFmeasure()
27
- if 'HCE' in metrics:
28
- HCE = HCEMeasure()
29
-
30
- if isinstance(gt_paths, list) and isinstance(pred_paths, list):
31
- # print(len(gt_paths), len(pred_paths))
32
- assert len(gt_paths) == len(pred_paths)
33
-
34
- for idx_sample in tqdm(range(len(gt_paths)), total=len(gt_paths)) if verbose else range(len(gt_paths)):
35
- gt = gt_paths[idx_sample]
36
- pred = pred_paths[idx_sample]
37
-
38
- pred = pred[:-4] + '.png'
39
- if os.path.exists(pred):
40
- pred_ary = cv2.imread(pred, cv2.IMREAD_GRAYSCALE)
41
- else:
42
- pred_ary = cv2.imread(pred.replace('.png', '.jpg'), cv2.IMREAD_GRAYSCALE)
43
- gt_ary = cv2.imread(gt, cv2.IMREAD_GRAYSCALE)
44
- pred_ary = cv2.resize(pred_ary, (gt_ary.shape[1], gt_ary.shape[0]))
45
-
46
- if 'E' in metrics:
47
- EM.step(pred=pred_ary, gt=gt_ary)
48
- if 'S' in metrics:
49
- SM.step(pred=pred_ary, gt=gt_ary)
50
- if 'F' in metrics:
51
- FM.step(pred=pred_ary, gt=gt_ary)
52
- if 'MAE' in metrics:
53
- MAE.step(pred=pred_ary, gt=gt_ary)
54
- if 'WF' in metrics:
55
- WFM.step(pred=pred_ary, gt=gt_ary)
56
- if 'HCE' in metrics:
57
- ske_path = gt.replace('/gt/', '/ske/')
58
- if os.path.exists(ske_path):
59
- ske_ary = cv2.imread(ske_path, cv2.IMREAD_GRAYSCALE)
60
- ske_ary = ske_ary > 128
61
- else:
62
- ske_ary = skeletonize(gt_ary > 128)
63
- ske_save_dir = os.path.join(*ske_path.split(os.sep)[:-1])
64
- if ske_path[0] == os.sep:
65
- ske_save_dir = os.sep + ske_save_dir
66
- os.makedirs(ske_save_dir, exist_ok=True)
67
- cv2.imwrite(ske_path, ske_ary.astype(np.uint8) * 255)
68
- HCE.step(pred=pred_ary, gt=gt_ary, gt_ske=ske_ary)
69
-
70
- if 'E' in metrics:
71
- em = EM.get_results()['em']
72
- else:
73
- em = {'curve': np.array([np.float64(-1)]), 'adp': np.float64(-1)}
74
- if 'S' in metrics:
75
- sm = SM.get_results()['sm']
76
- else:
77
- sm = np.float64(-1)
78
- if 'F' in metrics:
79
- fm = FM.get_results()['fm']
80
- else:
81
- fm = {'curve': np.array([np.float64(-1)]), 'adp': np.float64(-1)}
82
- if 'MAE' in metrics:
83
- mae = MAE.get_results()['mae']
84
- else:
85
- mae = np.float64(-1)
86
- if 'WF' in metrics:
87
- wfm = WFM.get_results()['wfm']
88
- else:
89
- wfm = np.float64(-1)
90
- if 'HCE' in metrics:
91
- hce = HCE.get_results()['hce']
92
- else:
93
- hce = np.float64(-1)
94
-
95
- return em, sm, fm, mae, wfm, hce
96
-
97
-
98
- def _prepare_data(pred: np.ndarray, gt: np.ndarray) -> tuple:
99
- gt = gt > 128
100
- pred = pred / 255
101
- if pred.max() != pred.min():
102
- pred = (pred - pred.min()) / (pred.max() - pred.min())
103
- return pred, gt
104
-
105
-
106
- def _get_adaptive_threshold(matrix: np.ndarray, max_value: float = 1) -> float:
107
- return min(2 * matrix.mean(), max_value)
108
-
109
-
110
- class Fmeasure(object):
111
- def __init__(self, beta: float = 0.3):
112
- self.beta = beta
113
- self.precisions = []
114
- self.recalls = []
115
- self.adaptive_fms = []
116
- self.changeable_fms = []
117
-
118
- def step(self, pred: np.ndarray, gt: np.ndarray):
119
- pred, gt = _prepare_data(pred, gt)
120
-
121
- adaptive_fm = self.cal_adaptive_fm(pred=pred, gt=gt)
122
- self.adaptive_fms.append(adaptive_fm)
123
-
124
- precisions, recalls, changeable_fms = self.cal_pr(pred=pred, gt=gt)
125
- self.precisions.append(precisions)
126
- self.recalls.append(recalls)
127
- self.changeable_fms.append(changeable_fms)
128
-
129
- def cal_adaptive_fm(self, pred: np.ndarray, gt: np.ndarray) -> float:
130
- adaptive_threshold = _get_adaptive_threshold(pred, max_value=1)
131
- binary_predcition = pred >= adaptive_threshold
132
- area_intersection = binary_predcition[gt].sum()
133
- if area_intersection == 0:
134
- adaptive_fm = 0
135
- else:
136
- pre = area_intersection / np.count_nonzero(binary_predcition)
137
- rec = area_intersection / np.count_nonzero(gt)
138
- adaptive_fm = (1 + self.beta) * pre * rec / (self.beta * pre + rec)
139
- return adaptive_fm
140
-
141
- def cal_pr(self, pred: np.ndarray, gt: np.ndarray) -> tuple:
142
- pred = (pred * 255).astype(np.uint8)
143
- bins = np.linspace(0, 256, 257)
144
- fg_hist, _ = np.histogram(pred[gt], bins=bins)
145
- bg_hist, _ = np.histogram(pred[~gt], bins=bins)
146
- fg_w_thrs = np.cumsum(np.flip(fg_hist), axis=0)
147
- bg_w_thrs = np.cumsum(np.flip(bg_hist), axis=0)
148
- TPs = fg_w_thrs
149
- Ps = fg_w_thrs + bg_w_thrs
150
- Ps[Ps == 0] = 1
151
- T = max(np.count_nonzero(gt), 1)
152
- precisions = TPs / Ps
153
- recalls = TPs / T
154
- numerator = (1 + self.beta) * precisions * recalls
155
- denominator = np.where(numerator == 0, 1, self.beta * precisions + recalls)
156
- changeable_fms = numerator / denominator
157
- return precisions, recalls, changeable_fms
158
-
159
- def get_results(self) -> dict:
160
- adaptive_fm = np.mean(np.array(self.adaptive_fms, _TYPE))
161
- changeable_fm = np.mean(np.array(self.changeable_fms, dtype=_TYPE), axis=0)
162
- precision = np.mean(np.array(self.precisions, dtype=_TYPE), axis=0) # N, 256
163
- recall = np.mean(np.array(self.recalls, dtype=_TYPE), axis=0) # N, 256
164
- return dict(fm=dict(adp=adaptive_fm, curve=changeable_fm),
165
- pr=dict(p=precision, r=recall))
166
-
167
-
168
- class MAEmeasure(object):
169
- def __init__(self):
170
- self.maes = []
171
-
172
- def step(self, pred: np.ndarray, gt: np.ndarray):
173
- pred, gt = _prepare_data(pred, gt)
174
-
175
- mae = self.cal_mae(pred, gt)
176
- self.maes.append(mae)
177
-
178
- def cal_mae(self, pred: np.ndarray, gt: np.ndarray) -> float:
179
- mae = np.mean(np.abs(pred - gt))
180
- return mae
181
-
182
- def get_results(self) -> dict:
183
- mae = np.mean(np.array(self.maes, _TYPE))
184
- return dict(mae=mae)
185
-
186
-
187
- class Smeasure(object):
188
- def __init__(self, alpha: float = 0.5):
189
- self.sms = []
190
- self.alpha = alpha
191
-
192
- def step(self, pred: np.ndarray, gt: np.ndarray):
193
- pred, gt = _prepare_data(pred=pred, gt=gt)
194
-
195
- sm = self.cal_sm(pred, gt)
196
- self.sms.append(sm)
197
-
198
- def cal_sm(self, pred: np.ndarray, gt: np.ndarray) -> float:
199
- y = np.mean(gt)
200
- if y == 0:
201
- sm = 1 - np.mean(pred)
202
- elif y == 1:
203
- sm = np.mean(pred)
204
- else:
205
- sm = self.alpha * self.object(pred, gt) + (1 - self.alpha) * self.region(pred, gt)
206
- sm = max(0, sm)
207
- return sm
208
-
209
- def object(self, pred: np.ndarray, gt: np.ndarray) -> float:
210
- fg = pred * gt
211
- bg = (1 - pred) * (1 - gt)
212
- u = np.mean(gt)
213
- object_score = u * self.s_object(fg, gt) + (1 - u) * self.s_object(bg, 1 - gt)
214
- return object_score
215
-
216
- def s_object(self, pred: np.ndarray, gt: np.ndarray) -> float:
217
- x = np.mean(pred[gt == 1])
218
- sigma_x = np.std(pred[gt == 1], ddof=1)
219
- score = 2 * x / (np.power(x, 2) + 1 + sigma_x + _EPS)
220
- return score
221
-
222
- def region(self, pred: np.ndarray, gt: np.ndarray) -> float:
223
- x, y = self.centroid(gt)
224
- part_info = self.divide_with_xy(pred, gt, x, y)
225
- w1, w2, w3, w4 = part_info['weight']
226
- pred1, pred2, pred3, pred4 = part_info['pred']
227
- gt1, gt2, gt3, gt4 = part_info['gt']
228
- score1 = self.ssim(pred1, gt1)
229
- score2 = self.ssim(pred2, gt2)
230
- score3 = self.ssim(pred3, gt3)
231
- score4 = self.ssim(pred4, gt4)
232
-
233
- return w1 * score1 + w2 * score2 + w3 * score3 + w4 * score4
234
-
235
- def centroid(self, matrix: np.ndarray) -> tuple:
236
- h, w = matrix.shape
237
- area_object = np.count_nonzero(matrix)
238
- if area_object == 0:
239
- x = np.round(w / 2)
240
- y = np.round(h / 2)
241
- else:
242
- # More details can be found at: https://www.yuque.com/lart/blog/gpbigm
243
- y, x = np.argwhere(matrix).mean(axis=0).round()
244
- return int(x) + 1, int(y) + 1
245
-
246
- def divide_with_xy(self, pred: np.ndarray, gt: np.ndarray, x, y) -> dict:
247
- h, w = gt.shape
248
- area = h * w
249
-
250
- gt_LT = gt[0:y, 0:x]
251
- gt_RT = gt[0:y, x:w]
252
- gt_LB = gt[y:h, 0:x]
253
- gt_RB = gt[y:h, x:w]
254
-
255
- pred_LT = pred[0:y, 0:x]
256
- pred_RT = pred[0:y, x:w]
257
- pred_LB = pred[y:h, 0:x]
258
- pred_RB = pred[y:h, x:w]
259
-
260
- w1 = x * y / area
261
- w2 = y * (w - x) / area
262
- w3 = (h - y) * x / area
263
- w4 = 1 - w1 - w2 - w3
264
-
265
- return dict(gt=(gt_LT, gt_RT, gt_LB, gt_RB),
266
- pred=(pred_LT, pred_RT, pred_LB, pred_RB),
267
- weight=(w1, w2, w3, w4))
268
-
269
- def ssim(self, pred: np.ndarray, gt: np.ndarray) -> float:
270
- h, w = pred.shape
271
- N = h * w
272
-
273
- x = np.mean(pred)
274
- y = np.mean(gt)
275
-
276
- sigma_x = np.sum((pred - x) ** 2) / (N - 1)
277
- sigma_y = np.sum((gt - y) ** 2) / (N - 1)
278
- sigma_xy = np.sum((pred - x) * (gt - y)) / (N - 1)
279
-
280
- alpha = 4 * x * y * sigma_xy
281
- beta = (x ** 2 + y ** 2) * (sigma_x + sigma_y)
282
-
283
- if alpha != 0:
284
- score = alpha / (beta + _EPS)
285
- elif alpha == 0 and beta == 0:
286
- score = 1
287
- else:
288
- score = 0
289
- return score
290
-
291
- def get_results(self) -> dict:
292
- sm = np.mean(np.array(self.sms, dtype=_TYPE))
293
- return dict(sm=sm)
294
-
295
-
296
- class Emeasure(object):
297
- def __init__(self):
298
- self.adaptive_ems = []
299
- self.changeable_ems = []
300
-
301
- def step(self, pred: np.ndarray, gt: np.ndarray):
302
- pred, gt = _prepare_data(pred=pred, gt=gt)
303
- self.gt_fg_numel = np.count_nonzero(gt)
304
- self.gt_size = gt.shape[0] * gt.shape[1]
305
-
306
- changeable_ems = self.cal_changeable_em(pred, gt)
307
- self.changeable_ems.append(changeable_ems)
308
- adaptive_em = self.cal_adaptive_em(pred, gt)
309
- self.adaptive_ems.append(adaptive_em)
310
-
311
- def cal_adaptive_em(self, pred: np.ndarray, gt: np.ndarray) -> float:
312
- adaptive_threshold = _get_adaptive_threshold(pred, max_value=1)
313
- adaptive_em = self.cal_em_with_threshold(pred, gt, threshold=adaptive_threshold)
314
- return adaptive_em
315
-
316
- def cal_changeable_em(self, pred: np.ndarray, gt: np.ndarray) -> np.ndarray:
317
- changeable_ems = self.cal_em_with_cumsumhistogram(pred, gt)
318
- return changeable_ems
319
-
320
- def cal_em_with_threshold(self, pred: np.ndarray, gt: np.ndarray, threshold: float) -> float:
321
- binarized_pred = pred >= threshold
322
- fg_fg_numel = np.count_nonzero(binarized_pred & gt)
323
- fg_bg_numel = np.count_nonzero(binarized_pred & ~gt)
324
-
325
- fg___numel = fg_fg_numel + fg_bg_numel
326
- bg___numel = self.gt_size - fg___numel
327
-
328
- if self.gt_fg_numel == 0:
329
- enhanced_matrix_sum = bg___numel
330
- elif self.gt_fg_numel == self.gt_size:
331
- enhanced_matrix_sum = fg___numel
332
- else:
333
- parts_numel, combinations = self.generate_parts_numel_combinations(
334
- fg_fg_numel=fg_fg_numel, fg_bg_numel=fg_bg_numel,
335
- pred_fg_numel=fg___numel, pred_bg_numel=bg___numel,
336
- )
337
-
338
- results_parts = []
339
- for i, (part_numel, combination) in enumerate(zip(parts_numel, combinations)):
340
- align_matrix_value = 2 * (combination[0] * combination[1]) / \
341
- (combination[0] ** 2 + combination[1] ** 2 + _EPS)
342
- enhanced_matrix_value = (align_matrix_value + 1) ** 2 / 4
343
- results_parts.append(enhanced_matrix_value * part_numel)
344
- enhanced_matrix_sum = sum(results_parts)
345
-
346
- em = enhanced_matrix_sum / (self.gt_size - 1 + _EPS)
347
- return em
348
-
349
- def cal_em_with_cumsumhistogram(self, pred: np.ndarray, gt: np.ndarray) -> np.ndarray:
350
- pred = (pred * 255).astype(np.uint8)
351
- bins = np.linspace(0, 256, 257)
352
- fg_fg_hist, _ = np.histogram(pred[gt], bins=bins)
353
- fg_bg_hist, _ = np.histogram(pred[~gt], bins=bins)
354
- fg_fg_numel_w_thrs = np.cumsum(np.flip(fg_fg_hist), axis=0)
355
- fg_bg_numel_w_thrs = np.cumsum(np.flip(fg_bg_hist), axis=0)
356
-
357
- fg___numel_w_thrs = fg_fg_numel_w_thrs + fg_bg_numel_w_thrs
358
- bg___numel_w_thrs = self.gt_size - fg___numel_w_thrs
359
-
360
- if self.gt_fg_numel == 0:
361
- enhanced_matrix_sum = bg___numel_w_thrs
362
- elif self.gt_fg_numel == self.gt_size:
363
- enhanced_matrix_sum = fg___numel_w_thrs
364
- else:
365
- parts_numel_w_thrs, combinations = self.generate_parts_numel_combinations(
366
- fg_fg_numel=fg_fg_numel_w_thrs, fg_bg_numel=fg_bg_numel_w_thrs,
367
- pred_fg_numel=fg___numel_w_thrs, pred_bg_numel=bg___numel_w_thrs,
368
- )
369
-
370
- results_parts = np.empty(shape=(4, 256), dtype=np.float64)
371
- for i, (part_numel, combination) in enumerate(zip(parts_numel_w_thrs, combinations)):
372
- align_matrix_value = 2 * (combination[0] * combination[1]) / \
373
- (combination[0] ** 2 + combination[1] ** 2 + _EPS)
374
- enhanced_matrix_value = (align_matrix_value + 1) ** 2 / 4
375
- results_parts[i] = enhanced_matrix_value * part_numel
376
- enhanced_matrix_sum = results_parts.sum(axis=0)
377
-
378
- em = enhanced_matrix_sum / (self.gt_size - 1 + _EPS)
379
- return em
380
-
381
- def generate_parts_numel_combinations(self, fg_fg_numel, fg_bg_numel, pred_fg_numel, pred_bg_numel):
382
- bg_fg_numel = self.gt_fg_numel - fg_fg_numel
383
- bg_bg_numel = pred_bg_numel - bg_fg_numel
384
-
385
- parts_numel = [fg_fg_numel, fg_bg_numel, bg_fg_numel, bg_bg_numel]
386
-
387
- mean_pred_value = pred_fg_numel / self.gt_size
388
- mean_gt_value = self.gt_fg_numel / self.gt_size
389
-
390
- demeaned_pred_fg_value = 1 - mean_pred_value
391
- demeaned_pred_bg_value = 0 - mean_pred_value
392
- demeaned_gt_fg_value = 1 - mean_gt_value
393
- demeaned_gt_bg_value = 0 - mean_gt_value
394
-
395
- combinations = [
396
- (demeaned_pred_fg_value, demeaned_gt_fg_value),
397
- (demeaned_pred_fg_value, demeaned_gt_bg_value),
398
- (demeaned_pred_bg_value, demeaned_gt_fg_value),
399
- (demeaned_pred_bg_value, demeaned_gt_bg_value)
400
- ]
401
- return parts_numel, combinations
402
-
403
- def get_results(self) -> dict:
404
- adaptive_em = np.mean(np.array(self.adaptive_ems, dtype=_TYPE))
405
- changeable_em = np.mean(np.array(self.changeable_ems, dtype=_TYPE), axis=0)
406
- return dict(em=dict(adp=adaptive_em, curve=changeable_em))
407
-
408
-
409
- class WeightedFmeasure(object):
410
- def __init__(self, beta: float = 1):
411
- self.beta = beta
412
- self.weighted_fms = []
413
-
414
- def step(self, pred: np.ndarray, gt: np.ndarray):
415
- pred, gt = _prepare_data(pred=pred, gt=gt)
416
-
417
- if np.all(~gt):
418
- wfm = 0
419
- else:
420
- wfm = self.cal_wfm(pred, gt)
421
- self.weighted_fms.append(wfm)
422
-
423
- def cal_wfm(self, pred: np.ndarray, gt: np.ndarray) -> float:
424
- # [Dst,IDXT] = bwdist(dGT);
425
- Dst, Idxt = bwdist(gt == 0, return_indices=True)
426
-
427
- # %Pixel dependency
428
- # E = abs(FG-dGT);
429
- E = np.abs(pred - gt)
430
- Et = np.copy(E)
431
- Et[gt == 0] = Et[Idxt[0][gt == 0], Idxt[1][gt == 0]]
432
-
433
- # K = fspecial('gaussian',7,5);
434
- # EA = imfilter(Et,K);
435
- K = self.matlab_style_gauss2D((7, 7), sigma=5)
436
- EA = convolve(Et, weights=K, mode="constant", cval=0)
437
- # MIN_E_EA = E;
438
- # MIN_E_EA(GT & EA<E) = EA(GT & EA<E);
439
- MIN_E_EA = np.where(gt & (EA < E), EA, E)
440
-
441
- # %Pixel importance
442
- B = np.where(gt == 0, 2 - np.exp(np.log(0.5) / 5 * Dst), np.ones_like(gt))
443
- Ew = MIN_E_EA * B
444
-
445
- TPw = np.sum(gt) - np.sum(Ew[gt == 1])
446
- FPw = np.sum(Ew[gt == 0])
447
-
448
-
449
- R = 1 - np.mean(Ew[gt == 1])
450
- P = TPw / (TPw + FPw + _EPS)
451
-
452
- # % Q = (1+Beta^2)*(R*P)./(eps+R+(Beta.*P));
453
- Q = (1 + self.beta) * R * P / (R + self.beta * P + _EPS)
454
-
455
- return Q
456
-
457
- def matlab_style_gauss2D(self, shape: tuple = (7, 7), sigma: int = 5) -> np.ndarray:
458
- """
459
- 2D gaussian mask - should give the same result as MATLAB's
460
- fspecial('gaussian',[shape],[sigma])
461
- """
462
- m, n = [(ss - 1) / 2 for ss in shape]
463
- y, x = np.ogrid[-m: m + 1, -n: n + 1]
464
- h = np.exp(-(x * x + y * y) / (2 * sigma * sigma))
465
- h[h < np.finfo(h.dtype).eps * h.max()] = 0
466
- sumh = h.sum()
467
- if sumh != 0:
468
- h /= sumh
469
- return h
470
-
471
- def get_results(self) -> dict:
472
- weighted_fm = np.mean(np.array(self.weighted_fms, dtype=_TYPE))
473
- return dict(wfm=weighted_fm)
474
-
475
-
476
- class HCEMeasure(object):
477
- def __init__(self):
478
- self.hces = []
479
-
480
- def step(self, pred: np.ndarray, gt: np.ndarray, gt_ske):
481
- # pred, gt = _prepare_data(pred, gt)
482
-
483
- hce = self.cal_hce(pred, gt, gt_ske)
484
- self.hces.append(hce)
485
-
486
- def get_results(self) -> dict:
487
- hce = np.mean(np.array(self.hces, _TYPE))
488
- return dict(hce=hce)
489
-
490
-
491
- def cal_hce(self, pred: np.ndarray, gt: np.ndarray, gt_ske: np.ndarray, relax=5, epsilon=2.0) -> float:
492
- # Binarize gt
493
- if(len(gt.shape)>2):
494
- gt = gt[:, :, 0]
495
-
496
- epsilon_gt = 128#(np.amin(gt)+np.amax(gt))/2.0
497
- gt = (gt>epsilon_gt).astype(np.uint8)
498
-
499
- # Binarize pred
500
- if(len(pred.shape)>2):
501
- pred = pred[:, :, 0]
502
- epsilon_pred = 128#(np.amin(pred)+np.amax(pred))/2.0
503
- pred = (pred>epsilon_pred).astype(np.uint8)
504
-
505
- Union = np.logical_or(gt, pred)
506
- TP = np.logical_and(gt, pred)
507
- FP = pred - TP
508
- FN = gt - TP
509
-
510
- # relax the Union of gt and pred
511
- Union_erode = Union.copy()
512
- Union_erode = cv2.erode(Union_erode.astype(np.uint8), disk(1), iterations=relax)
513
-
514
- # --- get the relaxed False Positive regions for computing the human efforts in correcting them ---
515
- FP_ = np.logical_and(FP, Union_erode) # get the relaxed FP
516
- for i in range(0, relax):
517
- FP_ = cv2.dilate(FP_.astype(np.uint8), disk(1))
518
- FP_ = np.logical_and(FP_, 1-np.logical_or(TP, FN))
519
- FP_ = np.logical_and(FP, FP_)
520
-
521
- # --- get the relaxed False Negative regions for computing the human efforts in correcting them ---
522
- FN_ = np.logical_and(FN, Union_erode) # preserve the structural components of FN
523
- ## recover the FN, where pixels are not close to the TP borders
524
- for i in range(0, relax):
525
- FN_ = cv2.dilate(FN_.astype(np.uint8), disk(1))
526
- FN_ = np.logical_and(FN_, 1-np.logical_or(TP, FP))
527
- FN_ = np.logical_and(FN, FN_)
528
- FN_ = np.logical_or(FN_, np.logical_xor(gt_ske, np.logical_and(TP, gt_ske))) # preserve the structural components of FN
529
-
530
- ## 2. =============Find exact polygon control points and independent regions==============
531
- ## find contours from FP_
532
- ctrs_FP, hier_FP = cv2.findContours(FP_.astype(np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)
533
- ## find control points and independent regions for human correction
534
- bdies_FP, indep_cnt_FP = self.filter_bdy_cond(ctrs_FP, FP_, np.logical_or(TP,FN_))
535
- ## find contours from FN_
536
- ctrs_FN, hier_FN = cv2.findContours(FN_.astype(np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)
537
- ## find control points and independent regions for human correction
538
- bdies_FN, indep_cnt_FN = self.filter_bdy_cond(ctrs_FN, FN_, 1-np.logical_or(np.logical_or(TP, FP_), FN_))
539
-
540
- poly_FP, poly_FP_len, poly_FP_point_cnt = self.approximate_RDP(bdies_FP, epsilon=epsilon)
541
- poly_FN, poly_FN_len, poly_FN_point_cnt = self.approximate_RDP(bdies_FN, epsilon=epsilon)
542
-
543
- # FP_points+FP_indep+FN_points+FN_indep
544
- return poly_FP_point_cnt+indep_cnt_FP+poly_FN_point_cnt+indep_cnt_FN
545
-
546
- def filter_bdy_cond(self, bdy_, mask, cond):
547
-
548
- cond = cv2.dilate(cond.astype(np.uint8), disk(1))
549
- labels = label(mask) # find the connected regions
550
- lbls = np.unique(labels) # the indices of the connected regions
551
- indep = np.ones(lbls.shape[0]) # the label of each connected regions
552
- indep[0] = 0 # 0 indicate the background region
553
-
554
- boundaries = []
555
- h,w = cond.shape[0:2]
556
- ind_map = np.zeros((h, w))
557
- indep_cnt = 0
558
-
559
- for i in range(0, len(bdy_)):
560
- tmp_bdies = []
561
- tmp_bdy = []
562
- for j in range(0, bdy_[i].shape[0]):
563
- r, c = bdy_[i][j,0,1],bdy_[i][j,0,0]
564
-
565
- if(np.sum(cond[r, c])==0 or ind_map[r, c]!=0):
566
- if(len(tmp_bdy)>0):
567
- tmp_bdies.append(tmp_bdy)
568
- tmp_bdy = []
569
- continue
570
- tmp_bdy.append([c, r])
571
- ind_map[r, c] = ind_map[r, c] + 1
572
- indep[labels[r, c]] = 0 # indicates part of the boundary of this region needs human correction
573
- if(len(tmp_bdy)>0):
574
- tmp_bdies.append(tmp_bdy)
575
-
576
- # check if the first and the last boundaries are connected
577
- # if yes, invert the first boundary and attach it after the last boundary
578
- if(len(tmp_bdies)>1):
579
- first_x, first_y = tmp_bdies[0][0]
580
- last_x, last_y = tmp_bdies[-1][-1]
581
- if((abs(first_x-last_x)==1 and first_y==last_y) or
582
- (first_x==last_x and abs(first_y-last_y)==1) or
583
- (abs(first_x-last_x)==1 and abs(first_y-last_y)==1)
584
- ):
585
- tmp_bdies[-1].extend(tmp_bdies[0][::-1])
586
- del tmp_bdies[0]
587
-
588
- for k in range(0, len(tmp_bdies)):
589
- tmp_bdies[k] = np.array(tmp_bdies[k])[:, np.newaxis, :]
590
- if(len(tmp_bdies)>0):
591
- boundaries.extend(tmp_bdies)
592
-
593
- return boundaries, np.sum(indep)
594
-
595
- # this function approximate each boundary by DP algorithm
596
- # https://en.wikipedia.org/wiki/Ramer%E2%80%93Douglas%E2%80%93Peucker_algorithm
597
- def approximate_RDP(self, boundaries, epsilon=1.0):
598
-
599
- boundaries_ = []
600
- boundaries_len_ = []
601
- pixel_cnt_ = 0
602
-
603
- # polygon approximate of each boundary
604
- for i in range(0, len(boundaries)):
605
- boundaries_.append(cv2.approxPolyDP(boundaries[i], epsilon, False))
606
-
607
- # count the control points number of each boundary and the total control points number of all the boundaries
608
- for i in range(0, len(boundaries_)):
609
- boundaries_len_.append(len(boundaries_[i]))
610
- pixel_cnt_ = pixel_cnt_ + len(boundaries_[i])
611
-
612
- return boundaries_, boundaries_len_, pixel_cnt_
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
BiRefNet_github/evaluation/valid.py DELETED
@@ -1,9 +0,0 @@
1
- from inference import inference
2
- from evaluation.evaluate import evaluate
3
-
4
-
5
- def valid(model, data_loader_test, pred_dir, method='tmp_val', testset='DIS-VD', only_S_MAE=True, device=0):
6
- model.eval()
7
- inference(model, data_loader_test, pred_dir, method, testset, device=device)
8
- performance_dict = evaluate(pred_dir, method, testset, only_S_MAE=only_S_MAE, epoch=model.epoch)
9
- return performance_dict
 
 
 
 
 
 
 
 
 
 
BiRefNet_github/gen_best_ep.py DELETED
@@ -1,85 +0,0 @@
1
- import os
2
- from glob import glob
3
- import numpy as np
4
- from config import Config
5
-
6
-
7
- config = Config()
8
-
9
- eval_txts = sorted(glob('e_results/*_eval.txt'))
10
- print('eval_txts:', [_.split(os.sep)[-1] for _ in eval_txts])
11
- score_panel = {}
12
- sep = '&'
13
- metrics = ['sm', 'wfm', 'hce'] # we used HCE for DIS and wFm for others.
14
- if 'DIS5K' not in config.task:
15
- metrics.remove('hce')
16
-
17
- for metric in metrics:
18
- print('Metric:', metric)
19
- current_line_nums = []
20
- for idx_et, eval_txt in enumerate(eval_txts):
21
- with open(eval_txt, 'r') as f:
22
- lines = [l for l in f.readlines()[3:] if '.' in l]
23
- current_line_nums.append(len(lines))
24
- for idx_et, eval_txt in enumerate(eval_txts):
25
- with open(eval_txt, 'r') as f:
26
- lines = [l for l in f.readlines()[3:] if '.' in l]
27
- for idx_line, line in enumerate(lines[:min(current_line_nums)]): # Consist line numbers by the minimal result file.
28
- properties = line.strip().strip(sep).split(sep)
29
- dataset = properties[0].strip()
30
- ckpt = properties[1].strip()
31
- if int(ckpt.split('--epoch_')[-1].strip()) < 0:
32
- continue
33
- targe_idx = {
34
- 'sm': [5, 2, 2, 5, 2],
35
- 'wfm': [3, 3, 8, 3, 8],
36
- 'hce': [7, -1, -1, 7, -1]
37
- }[metric][['DIS5K', 'COD', 'HRSOD', 'DIS5K+HRSOD+HRS10K', 'P3M-10k'].index(config.task)]
38
- if metric != 'hce':
39
- score_sm = float(properties[targe_idx].strip())
40
- else:
41
- score_sm = int(properties[targe_idx].strip().strip('.'))
42
- if idx_et == 0:
43
- score_panel[ckpt] = []
44
- score_panel[ckpt].append(score_sm)
45
-
46
- metrics_min = ['hce', 'mae']
47
- max_or_min = min if metric in metrics_min else max
48
- score_max = max_or_min(score_panel.values(), key=lambda x: np.sum(x))
49
-
50
- good_models = []
51
- for k, v in score_panel.items():
52
- if (np.sum(v) <= np.sum(score_max)) if metric in metrics_min else (np.sum(v) >= np.sum(score_max)):
53
- print(k, v)
54
- good_models.append(k)
55
-
56
- # Write
57
- with open(eval_txt, 'r') as f:
58
- lines = f.readlines()
59
- info4good_models = lines[:3]
60
- metric_names = [m.strip() for m in lines[1].strip().strip('&').split('&')[2:]]
61
- testset_mean_values = {metric_name: [] for metric_name in metric_names}
62
- for good_model in good_models:
63
- for idx_et, eval_txt in enumerate(eval_txts):
64
- with open(eval_txt, 'r') as f:
65
- lines = f.readlines()
66
- for line in lines:
67
- if set([good_model]) & set([_.strip() for _ in line.split(sep)]):
68
- info4good_models.append(line)
69
- metric_scores = [float(m.strip()) for m in line.strip().strip('&').split('&')[2:]]
70
- for idx_score, metric_score in enumerate(metric_scores):
71
- testset_mean_values[metric_names[idx_score]].append(metric_score)
72
-
73
- if 'DIS5K' in config.task:
74
- 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
75
- sample_line_for_placing_mean_values = info4good_models[-2]
76
- numbers_placed_well = sample_line_for_placing_mean_values.replace(sample_line_for_placing_mean_values.split('&')[1].strip(), 'DIS-TEs').strip().split('&')[3:]
77
- for idx_number, (number_placed_well, testset_mean_value) in enumerate(zip(numbers_placed_well, testset_mean_values_lst)):
78
- numbers_placed_well[idx_number] = number_placed_well.replace(number_placed_well.strip(), testset_mean_value)
79
- 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'
80
- info4good_models.append(testset_mean_line)
81
- info4good_models.append(lines[-1])
82
- info = ''.join(info4good_models)
83
- print(info)
84
- with open(os.path.join('e_results', 'eval-{}_best_on_{}.txt'.format(config.task, metric)), 'w') as f:
85
- f.write(info + '\n')
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
BiRefNet_github/inference.py DELETED
@@ -1,105 +0,0 @@
1
- import os
2
- import argparse
3
- from glob import glob
4
- from tqdm import tqdm
5
- import cv2
6
- import torch
7
-
8
- from dataset import MyData
9
- from models.birefnet import BiRefNet
10
- from utils import save_tensor_img, check_state_dict
11
- from config import Config
12
-
13
-
14
- config = Config()
15
-
16
-
17
- def inference(model, data_loader_test, pred_root, method, testset, device=0):
18
- model_training = model.training
19
- if model_training:
20
- model.eval()
21
- for batch in tqdm(data_loader_test, total=len(data_loader_test)) if 1 or config.verbose_eval else data_loader_test:
22
- inputs = batch[0].to(device)
23
- # gts = batch[1].to(device)
24
- label_paths = batch[-1]
25
- with torch.no_grad():
26
- scaled_preds = model(inputs)[-1].sigmoid()
27
-
28
- os.makedirs(os.path.join(pred_root, method, testset), exist_ok=True)
29
-
30
- for idx_sample in range(scaled_preds.shape[0]):
31
- res = torch.nn.functional.interpolate(
32
- scaled_preds[idx_sample].unsqueeze(0),
33
- size=cv2.imread(label_paths[idx_sample], cv2.IMREAD_GRAYSCALE).shape[:2],
34
- mode='bilinear',
35
- align_corners=True
36
- )
37
- 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
38
- if model_training:
39
- model.train()
40
- return None
41
-
42
-
43
- def main(args):
44
- # Init model
45
-
46
- device = config.device
47
- if args.ckpt_folder:
48
- print('Testing with models in {}'.format(args.ckpt_folder))
49
- else:
50
- print('Testing with model {}'.format(args.ckpt))
51
-
52
- if config.model == 'BiRefNet':
53
- model = BiRefNet(bb_pretrained=False)
54
- weights_lst = sorted(
55
- glob(os.path.join(args.ckpt_folder, '*.pth')) if args.ckpt_folder else [args.ckpt],
56
- key=lambda x: int(x.split('epoch_')[-1].split('.pth')[0]),
57
- reverse=True
58
- )
59
- for testset in args.testsets.split('+'):
60
- print('>>>> Testset: {}...'.format(testset))
61
- data_loader_test = torch.utils.data.DataLoader(
62
- dataset=MyData(testset, image_size=config.size, is_train=False),
63
- batch_size=config.batch_size_valid, shuffle=False, num_workers=config.num_workers, pin_memory=True
64
- )
65
- for weights in weights_lst:
66
- if int(weights.strip('.pth').split('epoch_')[-1]) % 1 != 0:
67
- continue
68
- print('\tInferencing {}...'.format(weights))
69
- # model.load_state_dict(torch.load(weights, map_location='cpu'))
70
- state_dict = torch.load(weights, map_location='cpu')
71
- state_dict = check_state_dict(state_dict)
72
- model.load_state_dict(state_dict)
73
- model = model.to(device)
74
- inference(
75
- model, data_loader_test=data_loader_test, pred_root=args.pred_root,
76
- method='--'.join([w.rstrip('.pth') for w in weights.split(os.sep)[-2:]]),
77
- testset=testset, device=config.device
78
- )
79
-
80
-
81
- if __name__ == '__main__':
82
- # Parameter from command line
83
- parser = argparse.ArgumentParser(description='')
84
- parser.add_argument('--ckpt', type=str, help='model folder')
85
- parser.add_argument('--ckpt_folder', default=sorted(glob(os.path.join('ckpt', '*')))[-1], type=str, help='model folder')
86
- parser.add_argument('--pred_root', default='e_preds', type=str, help='Output folder')
87
- parser.add_argument('--testsets',
88
- default={
89
- 'DIS5K': 'DIS-VD+DIS-TE1+DIS-TE2+DIS-TE3+DIS-TE4',
90
- 'COD': 'TE-COD10K+NC4K+TE-CAMO+CHAMELEON',
91
- 'HRSOD': 'DAVIS-S+TE-HRSOD+TE-UHRSD+TE-DUTS+DUT-OMRON',
92
- 'DIS5K+HRSOD+HRS10K': 'DIS-VD',
93
- 'P3M-10k': 'TE-P3M-500-P+TE-P3M-500-NP',
94
- 'DIS5K-': 'DIS-VD',
95
- 'COD-': 'TE-COD10K',
96
- 'SOD-': 'DAVIS-S+TE-HRSOD+TE-UHRSD',
97
- }[config.task + ''],
98
- type=str,
99
- help="Test all sets: , 'DIS-VD+DIS-TE1+DIS-TE2+DIS-TE3+DIS-TE4'")
100
-
101
- args = parser.parse_args()
102
-
103
- if config.precisionHigh:
104
- torch.set_float32_matmul_precision('high')
105
- main(args)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
BiRefNet_github/loss.py DELETED
@@ -1,274 +0,0 @@
1
- import torch
2
- from torch import nn
3
- import torch.nn.functional as F
4
- from torch.autograd import Variable
5
- from math import exp
6
- from config import Config
7
-
8
-
9
- class Discriminator(nn.Module):
10
- def __init__(self, channels=1, img_size=256):
11
- super(Discriminator, self).__init__()
12
-
13
- def discriminator_block(in_filters, out_filters, bn=Config().batch_size > 1):
14
- block = [nn.Conv2d(in_filters, out_filters, 3, 2, 1), nn.LeakyReLU(0.2, inplace=True), nn.Dropout2d(0.25)]
15
- if bn:
16
- block.append(nn.BatchNorm2d(out_filters, 0.8))
17
- return block
18
-
19
- self.model = nn.Sequential(
20
- *discriminator_block(channels, 16, bn=False),
21
- *discriminator_block(16, 32),
22
- *discriminator_block(32, 64),
23
- *discriminator_block(64, 128),
24
- )
25
-
26
- # The height and width of downsampled image
27
- ds_size = img_size // 2 ** 4
28
- self.adv_layer = nn.Sequential(nn.Linear(128 * ds_size ** 2, 1), nn.Sigmoid())
29
-
30
- def forward(self, img):
31
- out = self.model(img)
32
- out = out.view(out.shape[0], -1)
33
- validity = self.adv_layer(out)
34
-
35
- return validity
36
-
37
-
38
- class ContourLoss(torch.nn.Module):
39
- def __init__(self):
40
- super(ContourLoss, self).__init__()
41
-
42
- def forward(self, pred, target, weight=10):
43
- '''
44
- target, pred: tensor of shape (B, C, H, W), where target[:,:,region_in_contour] == 1,
45
- target[:,:,region_out_contour] == 0.
46
- weight: scalar, length term weight.
47
- '''
48
- # length term
49
- delta_r = pred[:,:,1:,:] - pred[:,:,:-1,:] # horizontal gradient (B, C, H-1, W)
50
- delta_c = pred[:,:,:,1:] - pred[:,:,:,:-1] # vertical gradient (B, C, H, W-1)
51
-
52
- delta_r = delta_r[:,:,1:,:-2]**2 # (B, C, H-2, W-2)
53
- delta_c = delta_c[:,:,:-2,1:]**2 # (B, C, H-2, W-2)
54
- delta_pred = torch.abs(delta_r + delta_c)
55
-
56
- epsilon = 1e-8 # where is a parameter to avoid square root is zero in practice.
57
- length = torch.mean(torch.sqrt(delta_pred + epsilon)) # eq.(11) in the paper, mean is used instead of sum.
58
-
59
- c_in = torch.ones_like(pred)
60
- c_out = torch.zeros_like(pred)
61
-
62
- region_in = torch.mean( pred * (target - c_in )**2 ) # equ.(12) in the paper, mean is used instead of sum.
63
- region_out = torch.mean( (1-pred) * (target - c_out)**2 )
64
- region = region_in + region_out
65
-
66
- loss = weight * length + region
67
-
68
- return loss
69
-
70
-
71
- class IoULoss(torch.nn.Module):
72
- def __init__(self):
73
- super(IoULoss, self).__init__()
74
-
75
- def forward(self, pred, target):
76
- b = pred.shape[0]
77
- IoU = 0.0
78
- for i in range(0, b):
79
- # compute the IoU of the foreground
80
- Iand1 = torch.sum(target[i, :, :, :] * pred[i, :, :, :])
81
- Ior1 = torch.sum(target[i, :, :, :]) + torch.sum(pred[i, :, :, :]) - Iand1
82
- IoU1 = Iand1 / Ior1
83
- # IoU loss is (1-IoU1)
84
- IoU = IoU + (1-IoU1)
85
- # return IoU/b
86
- return IoU
87
-
88
-
89
- class StructureLoss(torch.nn.Module):
90
- def __init__(self):
91
- super(StructureLoss, self).__init__()
92
-
93
- def forward(self, pred, target):
94
- weit = 1+5*torch.abs(F.avg_pool2d(target, kernel_size=31, stride=1, padding=15)-target)
95
- wbce = F.binary_cross_entropy_with_logits(pred, target, reduction='none')
96
- wbce = (weit*wbce).sum(dim=(2,3))/weit.sum(dim=(2,3))
97
-
98
- pred = torch.sigmoid(pred)
99
- inter = ((pred * target) * weit).sum(dim=(2, 3))
100
- union = ((pred + target) * weit).sum(dim=(2, 3))
101
- wiou = 1-(inter+1)/(union-inter+1)
102
-
103
- return (wbce+wiou).mean()
104
-
105
-
106
- class PatchIoULoss(torch.nn.Module):
107
- def __init__(self):
108
- super(PatchIoULoss, self).__init__()
109
- self.iou_loss = IoULoss()
110
-
111
- def forward(self, pred, target):
112
- win_y, win_x = 64, 64
113
- iou_loss = 0.
114
- for anchor_y in range(0, target.shape[0], win_y):
115
- for anchor_x in range(0, target.shape[1], win_y):
116
- patch_pred = pred[:, :, anchor_y:anchor_y+win_y, anchor_x:anchor_x+win_x]
117
- patch_target = target[:, :, anchor_y:anchor_y+win_y, anchor_x:anchor_x+win_x]
118
- patch_iou_loss = self.iou_loss(patch_pred, patch_target)
119
- iou_loss += patch_iou_loss
120
- return iou_loss
121
-
122
-
123
- class ThrReg_loss(torch.nn.Module):
124
- def __init__(self):
125
- super(ThrReg_loss, self).__init__()
126
-
127
- def forward(self, pred, gt=None):
128
- return torch.mean(1 - ((pred - 0) ** 2 + (pred - 1) ** 2))
129
-
130
-
131
- class ClsLoss(nn.Module):
132
- """
133
- Auxiliary classification loss for each refined class output.
134
- """
135
- def __init__(self):
136
- super(ClsLoss, self).__init__()
137
- self.config = Config()
138
- self.lambdas_cls = self.config.lambdas_cls
139
-
140
- self.criterions_last = {
141
- 'ce': nn.CrossEntropyLoss()
142
- }
143
-
144
- def forward(self, preds, gt):
145
- loss = 0.
146
- for _, pred_lvl in enumerate(preds):
147
- if pred_lvl is None:
148
- continue
149
- for criterion_name, criterion in self.criterions_last.items():
150
- loss += criterion(pred_lvl, gt) * self.lambdas_cls[criterion_name]
151
- return loss
152
-
153
-
154
- class PixLoss(nn.Module):
155
- """
156
- Pixel loss for each refined map output.
157
- """
158
- def __init__(self):
159
- super(PixLoss, self).__init__()
160
- self.config = Config()
161
- self.lambdas_pix_last = self.config.lambdas_pix_last
162
-
163
- self.criterions_last = {}
164
- if 'bce' in self.lambdas_pix_last and self.lambdas_pix_last['bce']:
165
- self.criterions_last['bce'] = nn.BCELoss() if not self.config.use_fp16 else nn.BCEWithLogitsLoss()
166
- if 'iou' in self.lambdas_pix_last and self.lambdas_pix_last['iou']:
167
- self.criterions_last['iou'] = IoULoss()
168
- if 'iou_patch' in self.lambdas_pix_last and self.lambdas_pix_last['iou_patch']:
169
- self.criterions_last['iou_patch'] = PatchIoULoss()
170
- if 'ssim' in self.lambdas_pix_last and self.lambdas_pix_last['ssim']:
171
- self.criterions_last['ssim'] = SSIMLoss()
172
- if 'mse' in self.lambdas_pix_last and self.lambdas_pix_last['mse']:
173
- self.criterions_last['mse'] = nn.MSELoss()
174
- if 'reg' in self.lambdas_pix_last and self.lambdas_pix_last['reg']:
175
- self.criterions_last['reg'] = ThrReg_loss()
176
- if 'cnt' in self.lambdas_pix_last and self.lambdas_pix_last['cnt']:
177
- self.criterions_last['cnt'] = ContourLoss()
178
- if 'structure' in self.lambdas_pix_last and self.lambdas_pix_last['structure']:
179
- self.criterions_last['structure'] = StructureLoss()
180
-
181
- def forward(self, scaled_preds, gt):
182
- loss = 0.
183
- criterions_embedded_with_sigmoid = ['structure', ] + ['bce'] if self.config.use_fp16 else []
184
- for _, pred_lvl in enumerate(scaled_preds):
185
- if pred_lvl.shape != gt.shape:
186
- pred_lvl = nn.functional.interpolate(pred_lvl, size=gt.shape[2:], mode='bilinear', align_corners=True)
187
- for criterion_name, criterion in self.criterions_last.items():
188
- _loss = criterion(pred_lvl.sigmoid() if criterion_name not in criterions_embedded_with_sigmoid else pred_lvl, gt) * self.lambdas_pix_last[criterion_name]
189
- loss += _loss
190
- # print(criterion_name, _loss.item())
191
- return loss
192
-
193
-
194
- class SSIMLoss(torch.nn.Module):
195
- def __init__(self, window_size=11, size_average=True):
196
- super(SSIMLoss, self).__init__()
197
- self.window_size = window_size
198
- self.size_average = size_average
199
- self.channel = 1
200
- self.window = create_window(window_size, self.channel)
201
-
202
- def forward(self, img1, img2):
203
- (_, channel, _, _) = img1.size()
204
- if channel == self.channel and self.window.data.type() == img1.data.type():
205
- window = self.window
206
- else:
207
- window = create_window(self.window_size, channel)
208
- if img1.is_cuda:
209
- window = window.cuda(img1.get_device())
210
- window = window.type_as(img1)
211
- self.window = window
212
- self.channel = channel
213
- return 1 - _ssim(img1, img2, window, self.window_size, channel, self.size_average)
214
-
215
-
216
- def gaussian(window_size, sigma):
217
- gauss = torch.Tensor([exp(-(x - window_size//2)**2/float(2*sigma**2)) for x in range(window_size)])
218
- return gauss/gauss.sum()
219
-
220
-
221
- def create_window(window_size, channel):
222
- _1D_window = gaussian(window_size, 1.5).unsqueeze(1)
223
- _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0)
224
- window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous())
225
- return window
226
-
227
-
228
- def _ssim(img1, img2, window, window_size, channel, size_average=True):
229
- mu1 = F.conv2d(img1, window, padding = window_size//2, groups=channel)
230
- mu2 = F.conv2d(img2, window, padding = window_size//2, groups=channel)
231
-
232
- mu1_sq = mu1.pow(2)
233
- mu2_sq = mu2.pow(2)
234
- mu1_mu2 = mu1*mu2
235
-
236
- sigma1_sq = F.conv2d(img1*img1, window, padding=window_size//2, groups=channel) - mu1_sq
237
- sigma2_sq = F.conv2d(img2*img2, window, padding=window_size//2, groups=channel) - mu2_sq
238
- sigma12 = F.conv2d(img1*img2, window, padding=window_size//2, groups=channel) - mu1_mu2
239
-
240
- C1 = 0.01**2
241
- C2 = 0.03**2
242
-
243
- ssim_map = ((2*mu1_mu2 + C1)*(2*sigma12 + C2))/((mu1_sq + mu2_sq + C1)*(sigma1_sq + sigma2_sq + C2))
244
-
245
- if size_average:
246
- return ssim_map.mean()
247
- else:
248
- return ssim_map.mean(1).mean(1).mean(1)
249
-
250
-
251
- def SSIM(x, y):
252
- C1 = 0.01 ** 2
253
- C2 = 0.03 ** 2
254
-
255
- mu_x = nn.AvgPool2d(3, 1, 1)(x)
256
- mu_y = nn.AvgPool2d(3, 1, 1)(y)
257
- mu_x_mu_y = mu_x * mu_y
258
- mu_x_sq = mu_x.pow(2)
259
- mu_y_sq = mu_y.pow(2)
260
-
261
- sigma_x = nn.AvgPool2d(3, 1, 1)(x * x) - mu_x_sq
262
- sigma_y = nn.AvgPool2d(3, 1, 1)(y * y) - mu_y_sq
263
- sigma_xy = nn.AvgPool2d(3, 1, 1)(x * y) - mu_x_mu_y
264
-
265
- SSIM_n = (2 * mu_x_mu_y + C1) * (2 * sigma_xy + C2)
266
- SSIM_d = (mu_x_sq + mu_y_sq + C1) * (sigma_x + sigma_y + C2)
267
- SSIM = SSIM_n / SSIM_d
268
-
269
- return torch.clamp((1 - SSIM) / 2, 0, 1)
270
-
271
-
272
- def saliency_structure_consistency(x, y):
273
- ssim = torch.mean(SSIM(x,y))
274
- return ssim
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
BiRefNet_github/make_a_copy.sh DELETED
@@ -1,18 +0,0 @@
1
- #!/bin/bash
2
- # Set dst repo here.
3
- repo=$1
4
- mkdir ../${repo}
5
- mkdir ../${repo}/evaluation
6
- mkdir ../${repo}/models
7
- mkdir ../${repo}/models/backbones
8
- mkdir ../${repo}/models/modules
9
- mkdir ../${repo}/models/refinement
10
-
11
- cp ./*.sh ../${repo}
12
- cp ./*.py ../${repo}
13
- cp ./evaluation/*.py ../${repo}/evaluation
14
- cp ./models/*.py ../${repo}/models
15
- cp ./models/backbones/*.py ../${repo}/models/backbones
16
- cp ./models/modules/*.py ../${repo}/models/modules
17
- cp ./models/refinement/*.py ../${repo}/models/refinement
18
- cp -r ./.git* ../${repo}
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
BiRefNet_github/requirements.txt DELETED
@@ -1,15 +0,0 @@
1
- --extra-index-url https://download.pytorch.org/whl/cu118
2
- torch==2.0.1
3
- --extra-index-url https://download.pytorch.org/whl/cu118
4
- torchvision==0.15.2
5
- numpy<2
6
- opencv-python
7
- timm
8
- scipy
9
- scikit-image
10
- kornia
11
-
12
- tqdm
13
- prettytable
14
-
15
- huggingface_hub
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
BiRefNet_github/rm_cache.sh DELETED
@@ -1,20 +0,0 @@
1
- #!/bin/bash
2
- rm -rf __pycache__ */__pycache__
3
-
4
- # Val
5
- rm -r tmp*
6
-
7
- # Train
8
- rm slurm*
9
- rm -r ckpt
10
- rm nohup.out*
11
-
12
- # Eval
13
- rm -r evaluation/eval-*
14
- rm -r tmp*
15
- rm -r e_logs/
16
-
17
- # System
18
- rm core-*-python-*
19
-
20
- clear
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
BiRefNet_github/sub.sh DELETED
@@ -1,19 +0,0 @@
1
- #!/bin/sh
2
- # Example: ./sub.sh tmp_proj 0,1,2,3 3 --> Use 0,1,2,3 for training, release GPUs, use GPU:3 for inference.
3
-
4
- module load compilers/cuda/11.8
5
-
6
- export PYTHONUNBUFFERED=1
7
- export LD_PRELOAD=/home/bingxing2/apps/compilers/gcc/12.2.0/lib64/libstdc++.so.6
8
- export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:${HOME}/miniconda3/lib:/home/bingxing2/apps/cudnn/8.4.0.27_cuda11.x/lib
9
-
10
- method=${1:-"BSL"}
11
- devices=${2:-0}
12
-
13
- sbatch --nodes=1 -p vip_gpu_ailab -A ai4bio \
14
- --ntasks-per-node=1 \
15
- --gres=gpu:$(($(echo ${devices%%,} | grep -o "," | wc -l)+1)) \
16
- --cpus-per-task=32 \
17
- ./train_test.sh ${method} ${devices}
18
-
19
- hostname
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
BiRefNet_github/test.sh DELETED
@@ -1,28 +0,0 @@
1
- devices=${1:-0}
2
- pred_root=${2:-e_preds}
3
-
4
- # Inference
5
-
6
- CUDA_VISIBLE_DEVICES=${devices} python inference.py --pred_root ${pred_root}
7
-
8
- echo Inference finished at $(date)
9
-
10
- # Evaluation
11
- log_dir=e_logs && mkdir ${log_dir}
12
-
13
- task=$(python3 config.py)
14
- case "${task}" in
15
- "DIS5K") testsets='DIS-VD,DIS-TE1,DIS-TE2,DIS-TE3,DIS-TE4' ;;
16
- "COD") testsets='CHAMELEON,NC4K,TE-CAMO,TE-COD10K' ;;
17
- "HRSOD") testsets='DAVIS-S,TE-HRSOD,TE-UHRSD,DUT-OMRON,TE-DUTS' ;;
18
- "DIS5K+HRSOD+HRS10K") testsets='DIS-VD' ;;
19
- "P3M-10k") testsets='TE-P3M-500-P,TE-P3M-500-NP' ;;
20
- esac
21
- testsets=(`echo ${testsets} | tr ',' ' '`) && testsets=${testsets[@]}
22
-
23
- for testset in ${testsets}; do
24
- nohup python eval_existingOnes.py --pred_root ${pred_root} --data_lst ${testset} > ${log_dir}/eval_${testset}.out 2>&1 &
25
- done
26
-
27
-
28
- echo Evaluation started at $(date)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
BiRefNet_github/train.py DELETED
@@ -1,377 +0,0 @@
1
- import os
2
- import datetime
3
- import argparse
4
- import torch
5
- import torch.nn as nn
6
- import torch.optim as optim
7
- from torch.autograd import Variable
8
-
9
- from config import Config
10
- from loss import PixLoss, ClsLoss
11
- from dataset import MyData
12
- from models.birefnet import BiRefNet
13
- from utils import Logger, AverageMeter, set_seed, check_state_dict
14
- from evaluation.valid import valid
15
-
16
- from torch.utils.data.distributed import DistributedSampler
17
- from torch.nn.parallel import DistributedDataParallel as DDP
18
- from torch.distributed import init_process_group, destroy_process_group, get_rank
19
- from torch.cuda import amp
20
-
21
-
22
- parser = argparse.ArgumentParser(description='')
23
- parser.add_argument('--resume', default=None, type=str, help='path to latest checkpoint')
24
- parser.add_argument('--epochs', default=120, type=int)
25
- parser.add_argument('--trainset', default='DIS5K', type=str, help="Options: 'DIS5K'")
26
- parser.add_argument('--ckpt_dir', default=None, help='Temporary folder')
27
- parser.add_argument('--testsets', default='DIS-VD+DIS-TE1+DIS-TE2+DIS-TE3+DIS-TE4', type=str)
28
- parser.add_argument('--dist', default=False, type=lambda x: x == 'True')
29
- args = parser.parse_args()
30
-
31
-
32
- config = Config()
33
- if config.rand_seed:
34
- set_seed(config.rand_seed)
35
-
36
- if config.use_fp16:
37
- # Half Precision
38
- scaler = amp.GradScaler(enabled=config.use_fp16)
39
-
40
- # DDP
41
- to_be_distributed = args.dist
42
- if to_be_distributed:
43
- init_process_group(backend="nccl", timeout=datetime.timedelta(seconds=3600*10))
44
- device = int(os.environ["LOCAL_RANK"])
45
- else:
46
- device = config.device
47
-
48
- epoch_st = 1
49
- # make dir for ckpt
50
- os.makedirs(args.ckpt_dir, exist_ok=True)
51
-
52
- # Init log file
53
- logger = Logger(os.path.join(args.ckpt_dir, "log.txt"))
54
- logger_loss_idx = 1
55
-
56
- # log model and optimizer params
57
- # logger.info("Model details:"); logger.info(model)
58
- logger.info("datasets: load_all={}, compile={}.".format(config.load_all, config.compile))
59
- logger.info("Other hyperparameters:"); logger.info(args)
60
- print('batch size:', config.batch_size)
61
-
62
-
63
- if os.path.exists(os.path.join(config.data_root_dir, config.task, args.testsets.strip('+').split('+')[0])):
64
- args.testsets = args.testsets.strip('+').split('+')
65
- else:
66
- args.testsets = []
67
-
68
- # Init model
69
- def prepare_dataloader(dataset: torch.utils.data.Dataset, batch_size: int, to_be_distributed=False, is_train=True):
70
- if to_be_distributed:
71
- return torch.utils.data.DataLoader(
72
- dataset=dataset, batch_size=batch_size, num_workers=min(config.num_workers, batch_size), pin_memory=True,
73
- shuffle=False, sampler=DistributedSampler(dataset), drop_last=True
74
- )
75
- else:
76
- return torch.utils.data.DataLoader(
77
- dataset=dataset, batch_size=batch_size, num_workers=min(config.num_workers, batch_size, 0), pin_memory=True,
78
- shuffle=is_train, drop_last=True
79
- )
80
-
81
-
82
- def init_data_loaders(to_be_distributed):
83
- # Prepare dataset
84
- train_loader = prepare_dataloader(
85
- MyData(datasets=config.training_set, image_size=config.size, is_train=True),
86
- config.batch_size, to_be_distributed=to_be_distributed, is_train=True
87
- )
88
- print(len(train_loader), "batches of train dataloader {} have been created.".format(config.training_set))
89
- test_loaders = {}
90
- for testset in args.testsets:
91
- _data_loader_test = prepare_dataloader(
92
- MyData(datasets=testset, image_size=config.size, is_train=False),
93
- config.batch_size_valid, is_train=False
94
- )
95
- print(len(_data_loader_test), "batches of valid dataloader {} have been created.".format(testset))
96
- test_loaders[testset] = _data_loader_test
97
- return train_loader, test_loaders
98
-
99
-
100
- def init_models_optimizers(epochs, to_be_distributed):
101
- model = BiRefNet(bb_pretrained=True)
102
- if args.resume:
103
- if os.path.isfile(args.resume):
104
- logger.info("=> loading checkpoint '{}'".format(args.resume))
105
- state_dict = torch.load(args.resume, map_location='cpu')
106
- state_dict = check_state_dict(state_dict)
107
- model.load_state_dict(state_dict)
108
- global epoch_st
109
- epoch_st = int(args.resume.rstrip('.pth').split('epoch_')[-1]) + 1
110
- else:
111
- logger.info("=> no checkpoint found at '{}'".format(args.resume))
112
- if to_be_distributed:
113
- model = model.to(device)
114
- model = DDP(model, device_ids=[device])
115
- else:
116
- model = model.to(device)
117
- if config.compile:
118
- model = torch.compile(model, mode=['default', 'reduce-overhead', 'max-autotune'][0])
119
- if config.precisionHigh:
120
- torch.set_float32_matmul_precision('high')
121
-
122
-
123
- # Setting optimizer
124
- if config.optimizer == 'AdamW':
125
- optimizer = optim.AdamW(params=model.parameters(), lr=config.lr, weight_decay=1e-2)
126
- elif config.optimizer == 'Adam':
127
- optimizer = optim.Adam(params=model.parameters(), lr=config.lr, weight_decay=0)
128
- lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
129
- optimizer,
130
- milestones=[lde if lde > 0 else epochs + lde + 1 for lde in config.lr_decay_epochs],
131
- gamma=config.lr_decay_rate
132
- )
133
- logger.info("Optimizer details:"); logger.info(optimizer)
134
- logger.info("Scheduler details:"); logger.info(lr_scheduler)
135
-
136
- return model, optimizer, lr_scheduler
137
-
138
-
139
- class Trainer:
140
- def __init__(
141
- self, data_loaders, model_opt_lrsch,
142
- ):
143
- self.model, self.optimizer, self.lr_scheduler = model_opt_lrsch
144
- self.train_loader, self.test_loaders = data_loaders
145
- if config.out_ref:
146
- self.criterion_gdt = nn.BCELoss() if not config.use_fp16 else nn.BCEWithLogitsLoss()
147
-
148
- # Setting Losses
149
- self.pix_loss = PixLoss()
150
- self.cls_loss = ClsLoss()
151
-
152
- # Others
153
- self.loss_log = AverageMeter()
154
- if config.lambda_adv_g:
155
- self.optimizer_d, self.lr_scheduler_d, self.disc, self.adv_criterion = self._load_adv_components()
156
- self.disc_update_for_odd = 0
157
-
158
- def _load_adv_components(self):
159
- # AIL
160
- from loss import Discriminator
161
- disc = Discriminator(channels=3, img_size=config.size)
162
- if to_be_distributed:
163
- disc = disc.to(device)
164
- disc = DDP(disc, device_ids=[device], broadcast_buffers=False)
165
- else:
166
- disc = disc.to(device)
167
- if config.compile:
168
- disc = torch.compile(disc, mode=['default', 'reduce-overhead', 'max-autotune'][0])
169
- adv_criterion = nn.BCELoss() if not config.use_fp16 else nn.BCEWithLogitsLoss()
170
- if config.optimizer == 'AdamW':
171
- optimizer_d = optim.AdamW(params=disc.parameters(), lr=config.lr, weight_decay=1e-2)
172
- elif config.optimizer == 'Adam':
173
- optimizer_d = optim.Adam(params=disc.parameters(), lr=config.lr, weight_decay=0)
174
- lr_scheduler_d = torch.optim.lr_scheduler.MultiStepLR(
175
- optimizer_d,
176
- milestones=[lde if lde > 0 else args.epochs + lde + 1 for lde in config.lr_decay_epochs],
177
- gamma=config.lr_decay_rate
178
- )
179
- return optimizer_d, lr_scheduler_d, disc, adv_criterion
180
-
181
- def _train_batch(self, batch):
182
- inputs = batch[0].to(device)
183
- gts = batch[1].to(device)
184
- class_labels = batch[2].to(device)
185
- if config.use_fp16:
186
- with amp.autocast(enabled=config.use_fp16):
187
- scaled_preds, class_preds_lst = self.model(inputs)
188
- if config.out_ref:
189
- (outs_gdt_pred, outs_gdt_label), scaled_preds = scaled_preds
190
- for _idx, (_gdt_pred, _gdt_label) in enumerate(zip(outs_gdt_pred, outs_gdt_label)):
191
- _gdt_pred = nn.functional.interpolate(_gdt_pred, size=_gdt_label.shape[2:], mode='bilinear', align_corners=True)#.sigmoid()
192
- # _gdt_label = _gdt_label.sigmoid()
193
- loss_gdt = self.criterion_gdt(_gdt_pred, _gdt_label) if _idx == 0 else self.criterion_gdt(_gdt_pred, _gdt_label) + loss_gdt
194
- # self.loss_dict['loss_gdt'] = loss_gdt.item()
195
- if None in class_preds_lst:
196
- loss_cls = 0.
197
- else:
198
- loss_cls = self.cls_loss(class_preds_lst, class_labels) * 1.0
199
- self.loss_dict['loss_cls'] = loss_cls.item()
200
-
201
- # Loss
202
- loss_pix = self.pix_loss(scaled_preds, torch.clamp(gts, 0, 1)) * 1.0
203
- self.loss_dict['loss_pix'] = loss_pix.item()
204
- # since there may be several losses for sal, the lambdas for them (lambdas_pix) are inside the loss.py
205
- loss = loss_pix + loss_cls
206
- if config.out_ref:
207
- loss = loss + loss_gdt * 1.0
208
-
209
- if config.lambda_adv_g:
210
- # gen
211
- valid = Variable(torch.cuda.FloatTensor(scaled_preds[-1].shape[0], 1).fill_(1.0), requires_grad=False).to(device)
212
- adv_loss_g = self.adv_criterion(self.disc(scaled_preds[-1] * inputs), valid) * config.lambda_adv_g
213
- loss += adv_loss_g
214
- self.loss_dict['loss_adv'] = adv_loss_g.item()
215
- self.disc_update_for_odd += 1
216
- # self.loss_log.update(loss.item(), inputs.size(0))
217
- # self.optimizer.zero_grad()
218
- # loss.backward()
219
- # self.optimizer.step()
220
- self.optimizer.zero_grad()
221
- scaler.scale(loss).backward()
222
- scaler.step(self.optimizer)
223
- scaler.update()
224
-
225
- if config.lambda_adv_g and self.disc_update_for_odd % 2 == 0:
226
- # disc
227
- fake = Variable(torch.cuda.FloatTensor(scaled_preds[-1].shape[0], 1).fill_(0.0), requires_grad=False).to(device)
228
- adv_loss_real = self.adv_criterion(self.disc(gts * inputs), valid)
229
- adv_loss_fake = self.adv_criterion(self.disc(scaled_preds[-1].detach() * inputs.detach()), fake)
230
- adv_loss_d = (adv_loss_real + adv_loss_fake) / 2 * config.lambda_adv_d
231
- self.loss_dict['loss_adv_d'] = adv_loss_d.item()
232
- # self.optimizer_d.zero_grad()
233
- # adv_loss_d.backward()
234
- # self.optimizer_d.step()
235
- self.optimizer_d.zero_grad()
236
- scaler.scale(adv_loss_d).backward()
237
- scaler.step(self.optimizer_d)
238
- scaler.update()
239
- else:
240
- scaled_preds, class_preds_lst = self.model(inputs)
241
- if config.out_ref:
242
- (outs_gdt_pred, outs_gdt_label), scaled_preds = scaled_preds
243
- for _idx, (_gdt_pred, _gdt_label) in enumerate(zip(outs_gdt_pred, outs_gdt_label)):
244
- _gdt_pred = nn.functional.interpolate(_gdt_pred, size=_gdt_label.shape[2:], mode='bilinear', align_corners=True).sigmoid()
245
- _gdt_label = _gdt_label.sigmoid()
246
- loss_gdt = self.criterion_gdt(_gdt_pred, _gdt_label) if _idx == 0 else self.criterion_gdt(_gdt_pred, _gdt_label) + loss_gdt
247
- # self.loss_dict['loss_gdt'] = loss_gdt.item()
248
- if None in class_preds_lst:
249
- loss_cls = 0.
250
- else:
251
- loss_cls = self.cls_loss(class_preds_lst, class_labels) * 1.0
252
- self.loss_dict['loss_cls'] = loss_cls.item()
253
-
254
- # Loss
255
- loss_pix = self.pix_loss(scaled_preds, torch.clamp(gts, 0, 1)) * 1.0
256
- self.loss_dict['loss_pix'] = loss_pix.item()
257
- # since there may be several losses for sal, the lambdas for them (lambdas_pix) are inside the loss.py
258
- loss = loss_pix + loss_cls
259
- if config.out_ref:
260
- loss = loss + loss_gdt * 1.0
261
-
262
- if config.lambda_adv_g:
263
- # gen
264
- valid = Variable(torch.cuda.FloatTensor(scaled_preds[-1].shape[0], 1).fill_(1.0), requires_grad=False).to(device)
265
- adv_loss_g = self.adv_criterion(self.disc(scaled_preds[-1] * inputs), valid) * config.lambda_adv_g
266
- loss += adv_loss_g
267
- self.loss_dict['loss_adv'] = adv_loss_g.item()
268
- self.disc_update_for_odd += 1
269
- self.loss_log.update(loss.item(), inputs.size(0))
270
- self.optimizer.zero_grad()
271
- loss.backward()
272
- self.optimizer.step()
273
-
274
- if config.lambda_adv_g and self.disc_update_for_odd % 2 == 0:
275
- # disc
276
- fake = Variable(torch.cuda.FloatTensor(scaled_preds[-1].shape[0], 1).fill_(0.0), requires_grad=False).to(device)
277
- adv_loss_real = self.adv_criterion(self.disc(gts * inputs), valid)
278
- adv_loss_fake = self.adv_criterion(self.disc(scaled_preds[-1].detach() * inputs.detach()), fake)
279
- adv_loss_d = (adv_loss_real + adv_loss_fake) / 2 * config.lambda_adv_d
280
- self.loss_dict['loss_adv_d'] = adv_loss_d.item()
281
- self.optimizer_d.zero_grad()
282
- adv_loss_d.backward()
283
- self.optimizer_d.step()
284
-
285
- def train_epoch(self, epoch):
286
- global logger_loss_idx
287
- self.model.train()
288
- self.loss_dict = {}
289
- if epoch > args.epochs + config.IoU_finetune_last_epochs:
290
- self.pix_loss.lambdas_pix_last['bce'] *= 0
291
- self.pix_loss.lambdas_pix_last['ssim'] *= 1
292
- self.pix_loss.lambdas_pix_last['iou'] *= 0.5
293
-
294
- for batch_idx, batch in enumerate(self.train_loader):
295
- self._train_batch(batch)
296
- # Logger
297
- if batch_idx % 20 == 0:
298
- info_progress = 'Epoch[{0}/{1}] Iter[{2}/{3}].'.format(epoch, args.epochs, batch_idx, len(self.train_loader))
299
- info_loss = 'Training Losses'
300
- for loss_name, loss_value in self.loss_dict.items():
301
- info_loss += ', {}: {:.3f}'.format(loss_name, loss_value)
302
- logger.info(' '.join((info_progress, info_loss)))
303
- info_loss = '@==Final== Epoch[{0}/{1}] Training Loss: {loss.avg:.3f} '.format(epoch, args.epochs, loss=self.loss_log)
304
- logger.info(info_loss)
305
-
306
- self.lr_scheduler.step()
307
- if config.lambda_adv_g:
308
- self.lr_scheduler_d.step()
309
- return self.loss_log.avg
310
-
311
- def validate_model(self, epoch):
312
- num_image_testset_all = {'DIS-VD': 470, 'DIS-TE1': 500, 'DIS-TE2': 500, 'DIS-TE3': 500, 'DIS-TE4': 500}
313
- num_image_testset = {}
314
- for testset in args.testsets:
315
- if 'DIS-TE' in testset:
316
- num_image_testset[testset] = num_image_testset_all[testset]
317
- weighted_scores = {'f_max': 0, 'f_mean': 0, 'f_wfm': 0, 'sm': 0, 'e_max': 0, 'e_mean': 0, 'mae': 0}
318
- len_all_data_loaders = 0
319
- self.model.epoch = epoch
320
- for testset, data_loader_test in self.test_loaders.items():
321
- print('Validating {}...'.format(testset))
322
- performance_dict = valid(
323
- self.model,
324
- data_loader_test,
325
- pred_dir='.',
326
- method=args.ckpt_dir.split('/')[-1] if args.ckpt_dir.split('/')[-1].strip('.').strip('/') else 'tmp_val',
327
- testset=testset,
328
- only_S_MAE=config.only_S_MAE,
329
- device=device
330
- )
331
- print('Test set: {}:'.format(testset))
332
- if config.only_S_MAE:
333
- print('Smeasure: {:.4f}, MAE: {:.4f}'.format(
334
- performance_dict['sm'], performance_dict['mae']
335
- ))
336
- else:
337
- print('Fmax: {:.4f}, Fwfm: {:.4f}, Smeasure: {:.4f}, Emean: {:.4f}, MAE: {:.4f}'.format(
338
- performance_dict['f_max'], performance_dict['f_wfm'], performance_dict['sm'], performance_dict['e_mean'], performance_dict['mae']
339
- ))
340
- if '-TE' in testset:
341
- for metric in ['sm', 'mae'] if config.only_S_MAE else ['f_max', 'f_mean', 'f_wfm', 'sm', 'e_max', 'e_mean', 'mae']:
342
- weighted_scores[metric] += performance_dict[metric] * len(data_loader_test)
343
- len_all_data_loaders += len(data_loader_test)
344
- print('Weighted Scores:')
345
- for metric, score in weighted_scores.items():
346
- if score:
347
- print('\t{}: {:.4f}.'.format(metric, score / len_all_data_loaders))
348
-
349
-
350
- def main():
351
-
352
- trainer = Trainer(
353
- data_loaders=init_data_loaders(to_be_distributed),
354
- model_opt_lrsch=init_models_optimizers(args.epochs, to_be_distributed)
355
- )
356
-
357
- for epoch in range(epoch_st, args.epochs+1):
358
- train_loss = trainer.train_epoch(epoch)
359
- # Save checkpoint
360
- # DDP
361
- if epoch >= args.epochs - config.save_last and epoch % config.save_step == 0:
362
- torch.save(
363
- trainer.model.module.state_dict() if to_be_distributed else trainer.model.state_dict(),
364
- os.path.join(args.ckpt_dir, 'epoch_{}.pth'.format(epoch))
365
- )
366
- if config.val_step and epoch >= args.epochs - config.save_last and (args.epochs - epoch) % config.val_step == 0:
367
- if to_be_distributed:
368
- if get_rank() == 0:
369
- print('Validating at rank-{}...'.format(get_rank()))
370
- trainer.validate_model(epoch)
371
- else:
372
- trainer.validate_model(epoch)
373
- if to_be_distributed:
374
- destroy_process_group()
375
-
376
- if __name__ == '__main__':
377
- main()
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
BiRefNet_github/train_test.sh DELETED
@@ -1,11 +0,0 @@
1
- #!/bin/sh
2
-
3
- method=${1:-"BSL"}
4
- devices=${2:-"0,1,2,3,4,5,6,7"}
5
-
6
- bash train.sh ${method} ${devices}
7
-
8
- devices_test=${3:-0}
9
- bash test.sh ${devices_test}
10
-
11
- hostname
 
 
 
 
 
 
 
 
 
 
 
 
BiRefNet_github/waiting4eval.py DELETED
@@ -1,141 +0,0 @@
1
- # --------------------------------------------------------
2
- # Make evaluation along with training. Swith time with space/computation.
3
- # Licensed under The MIT License [see LICENSE for details]
4
- # Written by Peng Zheng
5
- # --------------------------------------------------------
6
- import os
7
- from glob import glob
8
- from time import sleep
9
- import argparse
10
- import torch
11
-
12
- from config import Config
13
- from models.birefnet import BiRefNet
14
- from dataset import MyData
15
- from evaluation.valid import valid
16
-
17
-
18
- parser = argparse.ArgumentParser(description='')
19
- parser.add_argument('--cuda_idx', default=-1, type=int)
20
- parser.add_argument('--val_step', default=5*1, type=int)
21
- parser.add_argument('--program_id', default=0, type=int)
22
- # id-th one of this program will evaluate val_step * N + program_id -th epoch model.
23
- # 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.
24
- # Test fastest, only sequentially searched val_step*N -th models -- set all program_id as the same.
25
- parser.add_argument('--testsets', default='DIS-VD+DIS-TE1+DIS-TE2+DIS-TE3+DIS-TE4', type=str)
26
- args_eval = parser.parse_args()
27
-
28
- args_eval.program_id = (args_eval.val_step - args_eval.program_id) % args_eval.val_step
29
-
30
- config = Config()
31
- config.only_S_MAE = True
32
- device = 'cpu' if args_eval.cuda_idx < 0 else 'cuda:{}'.format(args_eval.cuda_idx)
33
- ckpt_dir, testsets = glob(os.path.join('ckpt', '*'))[0], args_eval.testsets
34
-
35
-
36
- def validate_model(model, test_loaders, epoch):
37
- num_image_testset_all = {'DIS-VD': 470, 'DIS-TE1': 500, 'DIS-TE2': 500, 'DIS-TE3': 500, 'DIS-TE4': 500}
38
- num_image_testset = {}
39
- for testset in testsets.split('+'):
40
- if 'DIS-TE' in testset:
41
- num_image_testset[testset] = num_image_testset_all[testset]
42
- weighted_scores = {'f_max': 0, 'sm': 0, 'e_max': 0, 'mae': 0}
43
- len_all_data_loaders = 0
44
- model.epoch = epoch
45
- for testset, data_loader_test in test_loaders.items():
46
- print('Validating {}...'.format(testset))
47
- performance_dict = valid(
48
- model,
49
- data_loader_test,
50
- pred_dir='.',
51
- method=ckpt_dir.split('/')[-1] if ckpt_dir.split('/')[-1].strip('.').strip('/') else 'tmp_val',
52
- testset=testset,
53
- only_S_MAE=config.only_S_MAE,
54
- device=device
55
- )
56
- print('Test set: {}:'.format(testset))
57
- if config.only_S_MAE:
58
- print('Smeasure: {:.4f}, MAE: {:.4f}'.format(
59
- performance_dict['sm'], performance_dict['mae']
60
- ))
61
- else:
62
- print('Fmax: {:.4f}, Fwfm: {:.4f}, Smeasure: {:.4f}, Emean: {:.4f}, MAE: {:.4f}'.format(
63
- performance_dict['f_max'], performance_dict['f_wfm'], performance_dict['sm'], performance_dict['e_mean'], performance_dict['mae']
64
- ))
65
- if '-TE' in testset:
66
- for metric in ['sm', 'mae'] if config.only_S_MAE else ['f_max', 'f_wfm', 'sm', 'e_mean', 'mae']:
67
- weighted_scores[metric] += performance_dict[metric] * len(data_loader_test)
68
- len_all_data_loaders += len(data_loader_test)
69
- print('Weighted Scores:')
70
- for metric, score in weighted_scores.items():
71
- if score:
72
- print('\t{}: {:.4f}.'.format(metric, score / len_all_data_loaders))
73
-
74
- @torch.no_grad()
75
- def main():
76
- config = Config()
77
- # Dataloader
78
- test_loaders = {}
79
- for testset in testsets.split('+'):
80
- dataset = MyData(
81
- datasets=testset,
82
- image_size=config.size, is_train=False
83
- )
84
- _data_loader_test = torch.utils.data.DataLoader(
85
- dataset=dataset, batch_size=config.batch_size_valid, num_workers=min(config.num_workers, config.batch_size_valid),
86
- pin_memory=device != 'cpu', shuffle=False
87
- )
88
- print(len(_data_loader_test), "batches of valid dataloader {} have been created.".format(testset))
89
- test_loaders[testset] = _data_loader_test
90
-
91
- # Model, 3070MiB GPU memory for inference
92
- model = BiRefNet(bb_pretrained=False).to(device)
93
- models_evaluated = []
94
- continous_sleep_time = 0
95
- while True:
96
- if (
97
- (models_evaluated and continous_sleep_time > 60*60*2) or
98
- (not models_evaluated and continous_sleep_time > 60*60*24)
99
- ):
100
- # If no ckpt has been saved, we wait for 24h;
101
- # elif some ckpts have been saved, we wait for 2h for new ones;
102
- # else: exit this waiting.
103
- print('Exiting the waiting for evaluation.')
104
- break
105
- models_evaluated_record = 'tmp_models_evaluated.txt'
106
- if os.path.exists(models_evaluated_record):
107
- with open(models_evaluated_record, 'r') as f:
108
- models_evaluated_global = f.read().splitlines()
109
- else:
110
- models_evaluated_global = []
111
- models_detected = [
112
- m for idx_m, m in enumerate(sorted(
113
- glob(os.path.join(ckpt_dir, '*.pth')),
114
- key=lambda x: int(x.rstrip('.pth').split('epoch_')[-1]), reverse=True
115
- )) if idx_m % args_eval.val_step == args_eval.program_id and m not in models_evaluated + models_evaluated_global
116
- ]
117
- if models_detected:
118
- from time import time
119
- time_st = time()
120
- # register the evaluated models
121
- model_not_evaluated_latest = models_detected[0]
122
- with open('tmp_models_evaluated.txt', 'a') as f:
123
- f.write(model_not_evaluated_latest + '\n')
124
- models_evaluated.append(model_not_evaluated_latest)
125
- print('Loading {} for validation...'.format(model_not_evaluated_latest))
126
-
127
- # evaluate the current model
128
- state_dict = torch.load(model_not_evaluated_latest, map_location=device)
129
- model.load_state_dict(state_dict, strict=False)
130
- validate_model(model, test_loaders, int(model_not_evaluated_latest.rstrip('.pth').split('epoch_')[-1]))
131
- continous_sleep_time = 0
132
- print('Duration of this evaluation:', time() - time_st)
133
- else:
134
- sleep_interval = 60 * 2
135
- sleep(sleep_interval)
136
- continous_sleep_time += sleep_interval
137
- continue
138
-
139
-
140
- if __name__ == '__main__':
141
- main()
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
BiRefNet_github/models/birefnet.py β†’ birefnet.py RENAMED
File without changes
config.json CHANGED
@@ -5,7 +5,7 @@
5
  ],
6
  "auto_map": {
7
  "AutoConfig": "BiRefNet_config.BiRefNetConfig",
8
- "AutoModelForImageSegmentation": "BiRefNet_github.models.birefnet.BiRefNet"
9
  },
10
  "custom_pipelines": {
11
  "image-segmentation": {
 
5
  ],
6
  "auto_map": {
7
  "AutoConfig": "BiRefNet_config.BiRefNetConfig",
8
+ "AutoModelForImageSegmentation": "birefnet.BiRefNet"
9
  },
10
  "custom_pipelines": {
11
  "image-segmentation": {
BiRefNet_github/config.py β†’ config.py RENAMED
File without changes
BiRefNet_github/dataset.py β†’ dataset.py RENAMED
File without changes
{BiRefNet_github/models β†’ models}/backbones/build_backbone.py RENAMED
File without changes
{BiRefNet_github/models β†’ models}/backbones/pvt_v2.py RENAMED
File without changes
{BiRefNet_github/models β†’ models}/backbones/swin_v1.py RENAMED
File without changes
{BiRefNet_github/models β†’ models}/modules/aspp.py RENAMED
File without changes
{BiRefNet_github/models β†’ models}/modules/attentions.py RENAMED
File without changes
{BiRefNet_github/models β†’ models}/modules/decoder_blocks.py RENAMED
File without changes
{BiRefNet_github/models β†’ models}/modules/deform_conv.py RENAMED
File without changes
{BiRefNet_github/models β†’ models}/modules/ing.py RENAMED
File without changes
{BiRefNet_github/models β†’ models}/modules/lateral_blocks.py RENAMED
File without changes
{BiRefNet_github/models β†’ models}/modules/mlp.py RENAMED
File without changes
{BiRefNet_github/models β†’ models}/modules/prompt_encoder.py RENAMED
File without changes
{BiRefNet_github/models β†’ models}/modules/utils.py RENAMED
File without changes
{BiRefNet_github/models β†’ models}/refinement/refiner.py RENAMED
File without changes
{BiRefNet_github/models β†’ models}/refinement/stem_layer.py RENAMED
File without changes
BiRefNet_github/preproc.py β†’ preproc.py RENAMED
File without changes
BiRefNet_github/train.sh β†’ train.sh RENAMED
File without changes
BiRefNet_github/utils.py β†’ utils.py RENAMED
File without changes