# j-morano/rrwnet architecture

> Architecture diagram of the GitHub repository j-morano/rrwnet, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-06-15). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/j-morano/rrwnet
- Repository: https://github.com/j-morano/rrwnet

## Overview

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/train.py`](https://github.com/j-morano/rrwnet/blob/main/train/train.py) (Train runner), [`get_predictions.py`](https://github.com/j-morano/rrwnet/blob/main/get_predictions.py) (Predict), [`eval/evaluate.py`](https://github.com/j-morano/rrwnet/blob/main/eval/evaluate.py) (Evaluate), [`train/data_utils.py`](https://github.com/j-morano/rrwnet/blob/main/train/data_utils.py) (Data pipeline)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_core["Core model"]
  node_model["RRWNet model<br/>PyTorch model<br/>[model.py]"]
  node_weights[("Pretrained weights<br/>Hub weights")]
  node_contract["RGB contract<br/>Data format"]
end

subgraph group_prep["Preprocessing"]
  node_preprocessing["Preprocess<br/>Data prep<br/>[preprocessing.py]"]
end

subgraph group_train["Training"]
  node_train_main["Train runner<br/>Orchestrator<br/>[train.py]"]
  node_train_config["Train config<br/>Settings<br/>[config.py]"]
  node_train_data["Data pipeline<br/>Dataset stack<br/>[data_utils.py]"]
  node_train_aug["Augmentations<br/>Transforms<br/>[transformations.py]"]
  node_train_factory["Model assembly<br/>Factory<br/>[factories.py]"]
  node_train_loss["Losses<br/>Objective<br/>[losses.py]"]
  node_train_preds["Train preds<br/>Prediction job<br/>[get_predictions.py]"]
end

subgraph group_infer["Inference"]
  node_infer_entry["Predict<br/>Inference script<br/>[get_predictions.py]"]
end

subgraph group_eval["Evaluation"]
  node_eval_main["Evaluate<br/>Evaluator<br/>[evaluate.py]"]
  node_eval_metrics["Metrics<br/>Analysis suite<br/>[topo_metric.py]"]
  node_eval_utils[("Eval data<br/>Asset store")]
end

subgraph group_data["Data assets"]
  node_data_assets[("Repo data<br/>Dataset")]
  node_distribution["Ground truth<br/>Data utilities<br/>[groundtruth.py]"]
end

node_weights -->|"load"| node_model
node_contract -->|"formats"| node_preprocessing
node_contract -->|"supervises"| node_train_data
node_contract -->|"inputs"| node_infer_entry
node_contract -->|"compares"| node_eval_main
node_preprocessing -->|"feeds"| node_train_data
node_preprocessing -.->|"preps"| node_infer_entry
node_train_config -->|"configures"| node_train_main
node_train_main -->|"loads"| node_train_data
node_train_main -->|"augments"| node_train_aug
node_train_main -->|"builds"| node_train_factory
node_train_main -->|"optimizes"| node_train_loss
node_train_factory -->|"instantiates"| node_model
node_train_factory -->|"pairs"| node_train_loss
node_train_preds -->|"generates"| node_eval_main
node_infer_entry -->|"runs"| node_model
node_infer_entry -->|"reads"| node_data_assets
node_data_assets -->|"contains"| node_distribution
node_eval_main -->|"analyzes"| node_eval_metrics
node_eval_main -->|"consumes"| node_eval_utils
node_distribution -->|"supplies"| node_train_data

click node_model "https://github.com/j-morano/rrwnet/blob/main/model.py"
click node_preprocessing "https://github.com/j-morano/rrwnet/blob/main/preprocessing.py"
click node_train_main "https://github.com/j-morano/rrwnet/blob/main/train/train.py"
click node_train_config "https://github.com/j-morano/rrwnet/blob/main/train/config.py"
click node_train_data "https://github.com/j-morano/rrwnet/blob/main/train/data_utils.py"
click node_train_aug "https://github.com/j-morano/rrwnet/blob/main/train/transformations.py"
click node_train_factory "https://github.com/j-morano/rrwnet/blob/main/train/factories.py"
click node_train_loss "https://github.com/j-morano/rrwnet/blob/main/train/losses.py"
click node_train_preds "https://github.com/j-morano/rrwnet/blob/main/train/get_predictions.py"
click node_infer_entry "https://github.com/j-morano/rrwnet/blob/main/get_predictions.py"
click node_eval_main "https://github.com/j-morano/rrwnet/blob/main/eval/evaluate.py"
click node_eval_metrics "https://github.com/j-morano/rrwnet/blob/main/eval/topo_metric.py"
click node_eval_utils "https://github.com/j-morano/rrwnet/tree/main/eval/_Evaluation_Data"
click node_data_assets "https://github.com/j-morano/rrwnet/tree/main/data"
click node_distribution "https://github.com/j-morano/rrwnet/blob/main/data_processing/groundtruth.py"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_model,node_weights,node_contract toneBlue
class node_preprocessing toneAmber
class node_train_main,node_train_config,node_train_data,node_train_aug,node_train_factory,node_train_loss,node_train_preds toneMint
class node_infer_entry toneRose
class node_eval_main,node_eval_metrics,node_eval_utils toneIndigo
class node_data_assets,node_distribution toneTeal
```

## Components

### Core model

Reusable RRWNet model boundary and its pretrained-weight distribution path.

- [RRWNet model](https://github.com/j-morano/rrwnet/blob/main/model.py): PyTorch model · Core reusable network definition used for loading pretrained weights and running prediction. · `model.py`
- **Pretrained weights**: Hub weights · External checkpoint source for pretrained RRWNet parameters.
- **RGB contract**: Data format · 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](https://github.com/j-morano/rrwnet/blob/main/preprocessing.py): Data prep · 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](https://github.com/j-morano/rrwnet/blob/main/train/train.py): Orchestrator · Main training entrypoint coordinating configuration, data, model, and optimization. · `train/train.py`
- [Train config](https://github.com/j-morano/rrwnet/blob/main/train/config.py): Settings · Central hyperparameters and command-line argument definitions for training. · `train/config.py`
- [Data pipeline](https://github.com/j-morano/rrwnet/blob/main/train/data_utils.py): Dataset stack · Training data loading and vessel-specific dataset assembly. · `train/data_utils.py`
- [Augmentations](https://github.com/j-morano/rrwnet/blob/main/train/transformations.py): Transforms · Training-time image and mask transformations. · `train/transformations.py`
- [Model assembly](https://github.com/j-morano/rrwnet/blob/main/train/factories.py): Factory · Factory-based construction of model variants, recursive refinement, and losses wiring. · `train/factories.py`
- [Losses](https://github.com/j-morano/rrwnet/blob/main/train/losses.py): Objective · Training objectives used to optimize RRWNet variants. · `train/losses.py`
- [Train preds](https://github.com/j-morano/rrwnet/blob/main/train/get_predictions.py): Prediction job · Training-package script for producing benchmark predictions after training. · `train/get_predictions.py`

### Inference

Standalone prediction and refinement entrypoint.

- [Predict](https://github.com/j-morano/rrwnet/blob/main/get_predictions.py): Inference script · Top-level entrypoint for direct prediction or refinement using pretrained RRWNet. · `get_predictions.py`

### Evaluation

Offline benchmarking, topology analysis, and statistical comparison.

- [Evaluate](https://github.com/j-morano/rrwnet/blob/main/eval/evaluate.py): Evaluator · Main evaluation driver consuming predictions and ground truth. · `eval/evaluate.py`
- [Metrics](https://github.com/j-morano/rrwnet/blob/main/eval/topo_metric.py): Analysis suite · Topology metrics, significance tests, efficiency analysis, and result comparison. · `eval/topo_metric.py`
- [Eval data](https://github.com/j-morano/rrwnet/tree/main/eval/_Evaluation_Data): Asset store · Reserved evaluation assets and offline benchmark data. · `eval/_Evaluation_Data`

### Data assets

Repository data directories and evaluation/training asset stores.

- [Repo data](https://github.com/j-morano/rrwnet/tree/main/data): Dataset · Example images, masks, and refinement inputs shipped with the repository. · `data`
- [Ground truth](https://github.com/j-morano/rrwnet/blob/main/data_processing/groundtruth.py): Data utilities · Utilities for data distribution and ground-truth handling. · `data_processing/groundtruth.py`

## Connections

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