Update README.md

README.md

@@ -2,10 +2,73 @@
 license: cc-by-nc-4.0
 pipeline_tag: image-segmentation
 tags:
+- artery-vein
+- retinal-imaging
+- segmentation
 - model_hub_mixin
 - pytorch_model_hub_mixin
 ---
 
-This model has been pushed to the Hub using the [PytorchModelHubMixin](https://huggingface.co/docs/huggingface_hub/package_reference/mixins#huggingface_hub.PyTorchModelHubMixin) integration:
-- Library: https://huggingface.co/j-morano/rrwnet-hrf
-- Docs: [More Information Needed]
+# rrwnet-rite
+
+This repo contains the the official weights of the RRWNet model trained on the RITE dataset, from the paper ["RRWNet: Recursive Refinement Network for Effective Retinal Artery/Vein Segmentation and Classification"](https://doi.org/10.1016/j.eswa.2024.124970), by José Morano, Guilherme Aresta, and Hrvoje Bogunović, published in _Expert Systems with Applications_ (2024).
+
+[[`arXiv`](https://doi.org/10.48550/arXiv.2402.03166)] [`ESWA`](https://doi.org/10.1016/j.eswa.2024.124970)]  [[`GitHub`](https://github.com/j-morano/rrwnet)] [[`BibTeX`](#citation)]
+
+
+![Overview](overview.png)
+
+
+
+## RRWNet models
+
+Model | Dataset | Train resolution | Weights
+--- | --- | --- | ---
+RRWNet | RITE | 720x576 (original) | [Download](https://huggingface.co/j-morano/rrwnet-rite)
+RRWNet | HRF | 1024 width (resized) | [Download](https://huggingface.co/j-morano/rrwnet-hrf)
+
+Please note that the size of the images used for training is important when using the weights for predictions.
+
+
+## Usage
+
+The model can be loaded using the `PyTorchModelHubMixin` from the `huggingface_hub` package and the code from the `model.py` file in our repo (<https://github.com/j-morano/rrwnet>).
+
+```python
+from huggingface_hub import PyTorchModelHubMixin
+from model import RRWNet as RRWNetModel
+
+
+class RRWNet(RRWNetModel, PyTorchModelHubMixin):
+    def __init__(self, input_ch=3, output_ch=3, base_ch=64, iterations=5):
+        super().__init__(input_ch, output_ch, base_ch, iterations)
+
+
+model = RRWNet.from_pretrained("j-morano/rrwnet-hrf")
+```
+
+
+## Preprocessing
+
+Models are trained using enhanced images and masks.
+You can preprocess the images offline using the `preprocessing.py` script in the repo.
+The script will enhance the images and masks and save them in the specified directory.
+
+```bash
+python3 preprocessing.py --images-path data/images/ --masks-path data/masks/ --save-path data/enhanced
+```
+
+
+## Citation
+
+If you use this code, the weights, the preprocessed data, or the predictions in your research, we would greatly appreciate it if you give a star to the repo and cite our work:
+
+```
+@article{morano2024rrwnet,
+    title={RRWNet: Recursive Refinement Network for Effective Retinal Artery/Vein Segmentation and Classification},
+    author={Morano, Jos{\'e} and Aresta, Guilherme and Bogunovi{\'c}, Hrvoje},
+    journal={Expert Systems with Applications},
+    year={2024},
+    doi={10.1016/j.eswa.2024.124970}
+}
+```