How j-morano/rrwnet works
1. Repository purpose
- This repo packages the full RRWNet workflow for retinal artery/vein segmentation and classification: core model definition, preprocessing, training, inference, and evaluation.
- The architecture is centered on a single PyTorch model that can be run standalone, trained from scratch, or used as a refinement module on existing predictions.
2. Core model boundary
- The root-level model.py is the main reusable model artifact.
- It is the piece exposed for loading pretrained weights from Hugging Face and is the dependency used by the top-level get_predictions.py script.
- This makes model.py the “production inference” boundary, separate from training and evaluation code.
3. Inference path
- get_predictions.py at the repository root is the entrypoint for applying pretrained RRWNet weights to images.
- It supports two usage modes:
- direct prediction from fundus images
- refinement of existing predictions from another model
- This script is the main bridge from dataset inputs to saved RGB prediction masks.
4. Preprocessing subsystem
- preprocessing.py performs offline enhancement of images and masks before training or inference.
- The README makes clear this preprocessing is optional for custom training, but required to reproduce the paper’s trained-model behavior.
- Architecturally, this is a separate data preparation stage that sits upstream of both training and inference.
5. Data contract
- Inputs are expected as RGB fundus images with pixel values in [0, 255].
- Masks/predictions are also RGB, but channel semantics matter:
- red = arteries
- green = veins
- blue = vessels union
- This channelized label format is a key boundary in the system because both training targets and evaluation inputs depend on it.
6. Training subsystem
- The train/ package contains the full training pipeline.
- train/train.py is the orchestrator; train/config.py centralizes hyperparameters and command-line arguments.
- Supporting modules in this package cover data handling, model construction, losses, transformations, and utility functions, making this directory a self-contained training stack.
7. Training data and model assembly
- train/data_utils.py, train/data_vessels.py, and train/transformations.py indicate a standard PyTorch dataset/augmentation pipeline feeding the model.
- train/factories.py, train/models.py, and train/losses.py suggest a factory-based assembly of model variants and objectives.
- train/r2av.py likely encapsulates the recursive refinement behavior used by RRWNet, distinguishing the network from a plain segmentation model.
8. Prediction generation during training
- train/get_predictions.py is separate from the top-level inference script and appears intended for generating test-set outputs after training.
- This creates a clean split between “model training” and “evaluation-ready predictions,” which is important for reproducing benchmark results.
9. Evaluation subsystem
- The eval/ package is a distinct postprocessing/benchmarking stage that consumes saved predictions and ground truth.
- eval/evaluate.py is the main evaluator, with supporting modules for topology metrics, efficiency analysis, significance tests, and result comparison.
- This suggests the repo evaluates not just pixel overlap, but also topology-aware quality and statistical robustness.
10. Evaluation data boundary
- eval/_Evaluation_Data/ is reserved for evaluation assets, separate from the raw training data and model outputs.
- The presence of compare_results.py, topo_metric.py, and significance_test.py shows that evaluation is designed as an offline analysis pipeline rather than an online runtime component.
11. Data organization
- The data/ directory holds example assets and is split into images/, masks/, and U-Net_predictions/.
- This mirrors the repo’s core use cases: raw inputs, supervision masks, and third-party predictions to be refined.
- The structure reinforces the refinement-centric architecture: RRWNet can be used as a downstream corrective stage on existing segmentation maps.
12. Distribution and runtime dependencies
- The repo relies on PyTorch for model execution and on Hugging Face’s model hub for pretrained weight distribution.
- safetensors is part of the loading path for model checkpoints.
- The architecture therefore supports both local experimentation and remote pretrained model retrieval without changing the model code.
13. Overall system flow
- The practical flow is: prepare/enhance data → train RRWNet or load pretrained weights → generate predictions or refine existing masks → run evaluation against ground truth.
- The key architectural separation is between:
- core model code at the root
- training-specific logic in train/
- evaluation-specific logic in eval/
- data preparation in preprocessing.py and data_processing/
14. What to draw in an architecture diagram
- A clear diagram should show one central RRWNet model component with four surrounding subsystems:
- preprocessing/data preparation
- training pipeline
- inference/refinement pipeline
- evaluation/analysis pipeline
- The most important external boundary is the pretrained-weight distribution via Hugging Face/releases, and the most important internal boundary is the RGB image/mask format contract shared across all stages.
Main entry points: train.py Train runner · get_predictions.py Predict · evaluate.py Evaluate · data_utils.py Data pipeline
Components 17
Core model
Reusable RRWNet model boundary and its pretrained-weight distribution path.
- RRWNet model · Core reusable network definition used for loading pretrained weights and running prediction.
model.py - Pretrained weights · External checkpoint source for pretrained RRWNet parameters.
- RGB contract · Shared image/mask convention: RGB fundus inputs and RGB masks with artery, vein, and vessel channels.
Preprocessing
Offline enhancement and data preparation before training or inference.
- Preprocess · Offline enhancement of images and masks before training or paper-reproduction inference.
preprocessing.py
Training
End-to-end PyTorch training stack and training-time prediction generation.
- Train runner · Main training entrypoint coordinating configuration, data, model, and optimization.
train/train.py - Train config · Central hyperparameters and command-line argument definitions for training.
train/config.py - Data pipeline · Training data loading and vessel-specific dataset assembly.
train/data_utils.py - Augmentations · Training-time image and mask transformations.
train/transformations.py - Model assembly · Factory-based construction of model variants, recursive refinement, and losses wiring.
train/factories.py - Losses · Training objectives used to optimize RRWNet variants.
train/losses.py - Train preds · Training-package script for producing benchmark predictions after training.
train/get_predictions.py
Inference
Standalone prediction and refinement entrypoint.
- Predict · Top-level entrypoint for direct prediction or refinement using pretrained RRWNet.
get_predictions.py
Evaluation
Offline benchmarking, topology analysis, and statistical comparison.
Data assets
Repository data directories and evaluation/training asset stores.
- Repo data · Example images, masks, and refinement inputs shipped with the repository.
data - Ground truth · Utilities for data distribution and ground-truth handling.
data_processing/groundtruth.py
Connections 21
- Pretrained weights → RRWNet model: load · Pretrained checkpoints initialize the core RRWNet model.
- RGB contract → Preprocess: formats · Preprocessing preserves the RGB image and channelized mask contract.
- RGB contract → Data pipeline: supervises · Training datasets consume the same RGB image/mask convention.
- RGB contract → Predict: inputs · Inference expects RGB fundus images and RGB prediction masks.
- RGB contract → Evaluate: compares · Evaluation consumes predictions and ground truth in the shared mask format.
- Preprocess → Data pipeline: feeds · Prepared images and masks feed the training dataset pipeline.
- Preprocess → Predict: preps · Optional offline preprocessing can be applied before inference.
- Train config → Train runner: configures · Training runner reads hyperparameters and CLI arguments.
- Train runner → Data pipeline: loads · Training orchestration builds the dataset pipeline.
- Train runner → Augmentations: augments · Training runner applies transformation pipeline to samples.
- Train runner → Model assembly: builds · Training runner constructs model and refinement components.
- Train runner → Losses: optimizes · Training runner wires in the selected objective functions.
- Model assembly → RRWNet model: instantiates · Factory code assembles the RRWNet architecture around the reusable model boundary.
- Model assembly → Losses: pairs · Factory code connects model variants with training objectives.
- Train preds → Evaluate: generates · Training-package prediction job produces outputs for downstream benchmarking.
- Predict → RRWNet model: runs · Top-level inference script applies the core model to images or prior predictions.
- Predict → Repo data: reads · Inference consumes repository images or refinement inputs.
- Repo data → Ground truth: contains · Repository data directories hold image and mask assets managed by data utilities.
- Evaluate → Metrics: analyzes · Evaluation driver invokes metric and statistical analysis modules.
- Evaluate → Eval data: consumes · Evaluation reads dedicated offline assets and benchmark outputs.
- Ground truth → Data pipeline: supplies · Ground-truth utilities support dataset preparation and label handling.