j-morano/rrwnet-rite

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", by José Morano, Guilherme Aresta, and Hrvoje Bogunović, published in Expert Systems with Applications (2024).

[arXiv] ESWA] [GitHub] [BibTeX]

RRWNet models

Model	Dataset	Resolution	Weights
RRWNet	RITE	720x576 (original)	Download
RRWNet	HRF	1024 width (resized)	Download

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).

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-rite")

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.

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