README.md

---
license: cc-by-nc-4.0
pipeline_tag: image-segmentation
tags:
- arxiv:2402.03166
- artery-vein
- retinal-imaging
- segmentation
- model_hub_mixin
- pytorch_model_hub_mixin
---

# RRWNet HRF

This repo contains the the official weights of the RRWNet model trained on the HRF 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://arxiv.org/abs/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 | Resolution | Weights
--- | --- | --- | ---
RRWNet | RITE | 584x565 (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>), which can be downloaded from [here](https://raw.githubusercontent.com/j-morano/rrwnet/refs/heads/main/model.py).

```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}
}
```