Improve the README.

README.md

@@ -8,7 +8,7 @@ tags:
 - Salient Object Detection
 - pytorch_model_hub_mixin
 - model_hub_mixin
-repo_url: https://github.com/ZhengPeng7/BiRefNet
+repo_url: https://github.com/ZhengPeng7/BiRefNet-portrait
 pipeline_tag: image-segmentation
 ---
 <h1 align="center">Bilateral Reference for High-Resolution Dichotomous Image Segmentation</h1>
@@ -44,7 +44,7 @@ pipeline_tag: image-segmentation
 | :------------------------------: | :-------------------------------: |
 | <img src="https://drive.google.com/thumbnail?id=1ItXaA26iYnE8XQ_GgNLy71MOWePoS2-g&sz=w400" /> |  <img src="https://drive.google.com/thumbnail?id=1Z-esCujQF_uEa_YJjkibc3NUrW4aR_d4&sz=w400" /> |
 
-This repo is the official implementation of "[**Bilateral Reference for High-Resolution Dichotomous Image Segmentation**](https://arxiv.org/pdf/2401.03407.pdf)" (___arXiv 2024___). 
+This repo is the official implementation of "[**Bilateral Reference for High-Resolution Dichotomous Image Segmentation**](https://arxiv.org/pdf/2401.03407.pdf)" (___CAAI AIR 2024___). 
 
 Visit our GitHub repo: [https://github.com/ZhengPeng7/BiRefNet](https://github.com/ZhengPeng7/BiRefNet) for more details -- **codes**, **docs**, and **model zoo**!
 
@@ -52,6 +52,15 @@ Visit our GitHub repo: [https://github.com/ZhengPeng7/BiRefNet](https://github.c
 
 ### 1. Load BiRefNet:
 
+#### Use codes + weights from HuggingFace
+> Only use the weights on HuggingFace -- Pro: No need to download BiRefNet codes manually; Con: Codes on HuggingFace might not be latest version (I'll try to keep them always latest).
+
+```python
+# Load BiRefNet with weights
+from transformers import AutoModelForImageSegmentation
+birefnet = AutoModelForImageSegmentation.from_pretrained('zhengpeng7/birefnet-portrait', trust_remote_code=True)
+```
+
 #### Use codes from GitHub + weights from HuggingFace
 > Only use the weights on HuggingFace -- Pro: codes are always latest; Con: Need to clone the BiRefNet repo from my GitHub.
 
@@ -62,18 +71,25 @@ cd BiRefNet
 ```
 
 ```python
-# Load weights
+# Use codes locally
 from models.birefnet import BiRefNet
+
+# Load weights from Hugging Face Models
 birefnet = BiRefNet.from_pretrained('zhengpeng7/birefnet-portrait')
 ```
 
-#### Use codes + weights from HuggingFace
-> Only use the weights on HuggingFace -- Pro: No need to download BiRefNet codes manually; Con: Codes on HuggingFace might not be latest version (I'll try to keep them always latest).
+#### Use codes from GitHub + weights from HuggingFace
+> Only use the weights and codes both locally.
 
 ```python
-# Load BiRefNet with weights
-from transformers import AutoModelForImageSegmentation
-birefnet = AutoModelForImageSegmentation.from_pretrained('zhengpeng7/birefnet-portrait', trust_remote_code=True)
+# Use codes and weights locally
+import torch
+from utils import check_state_dict
+
+birefnet = BiRefNet(bb_pretrained=False)
+state_dict = torch.load(PATH_TO_WEIGHT, map_location='cpu')
+state_dict = check_state_dict(state_dict)
+birefnet.load_state_dict(state_dict)
 ```
 
 #### Use the loaded BiRefNet for inference
@@ -121,7 +137,7 @@ plt.show()
 
 > This BiRefNet for standard dichotomous image segmentation (DIS) is trained on **DIS-TR** and validated on **DIS-TEs and DIS-VD**.
 
-## This repo holds the official model weights of "[<ins>Bilateral Reference for High-Resolution Dichotomous Image Segmentation</ins>](https://arxiv.org/pdf/2401.03407)" (_arXiv 2024_).
+## This repo holds the official model weights of "[<ins>Bilateral Reference for High-Resolution Dichotomous Image Segmentation</ins>](https://arxiv.org/pdf/2401.03407)" (_CAAI AIR 2024_).
 
 This repo contains the weights of BiRefNet proposed in our paper, which has achieved the SOTA performance on three tasks (DIS, HRSOD, and COD).
 
@@ -144,10 +160,10 @@ Go to my GitHub page for BiRefNet codes and the latest updates: https://github.c
 ## Citation
 
 ```
-@article{zheng2024birefnet,
+@article{BiRefNet,
   title={Bilateral Reference for High-Resolution Dichotomous Image Segmentation},
   author={Zheng, Peng and Gao, Dehong and Fan, Deng-Ping and Liu, Li and Laaksonen, Jorma and Ouyang, Wanli and Sebe, Nicu},
-  journal={arXiv},
+  journal={CAAI Artificial Intelligence Research},
   year={2024}
 }
 ```