Migrate repo

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,12 @@
+*.pyc
+__pycache__
+*.egg-info
+*.zip
+.DS_Store
+.cache/
+.mplconfig/
+model_releases/
+output_*/
+output_*.zip
+/samples/fundus/*
+!/samples/fundus/original

README.md

@@ -1,3 +1,193 @@
----
-license: mit
----
+---
+license: agpl-3.0
+pipeline_tag: image-segmentation
+tags:
+- medical
+- biology
+---
+
+# 👁️ VascX Fork
+
+This repository contains the instructions for using the VascX models from the paper [VascX Models: Model Ensembles for Retinal Vascular Analysis from Color Fundus Images](https://arxiv.org/abs/2409.16016). This fork is published as `zyf0717/vascx-fork` on the Hugging Face Hub.
+
+The model weights are in [huggingface](https://huggingface.co/zyf0717/vascx-fork).
+
+<img src="imgs/samples_vascx_hrf.png">
+
+## 🛠️ Installation
+
+To install the entire fundus analysis pipeline including fundus preprocessing, model inference code and vascular biomarker extraction:
+
+1. Create a conda or virtualenv virtual environment, or otherwise ensure a clean environment.
+
+2. Install `torch` and `torchvision` for your platform.
+
+3. Install the pipeline runtime packages:
+
+```bash
+pip install retinalysis-fundusprep retinalysis-inference
+pip install -e .
+```
+
+The `environment.yml` in this repository includes the same runtime dependencies.
+
+
+## 🚀 `vascx run` Command
+
+The repository name is `vascx-fork`, but the installed CLI entry point remains `vascx` for compatibility.
+
+The `run` command provides a comprehensive pipeline for processing fundus images, performing various analyses, and creating visualizations.
+
+### Usage
+
+```bash
+vascx run DATA_PATH OUTPUT_PATH [OPTIONS]
+```
+
+If `config.yaml` exists in the current working directory or at the repository root, `vascx run` loads it automatically. You can also point to a specific file with `--config /path/to/config.yaml`.
+
+### Arguments
+
+- `DATA_PATH`: Path to input data. Can be either:
+  - A directory containing fundus images
+  - A CSV file with a 'path' column containing paths to images
+
+- `OUTPUT_PATH`: Directory where processed results will be stored
+
+### Options
+
+| Option | Default | Description |
+|--------|---------|-------------|
+| `--preprocess/--no-preprocess` | `--preprocess` | Run preprocessing to standardize images for model input |
+| `--vessels/--no-vessels` | `--vessels` | Run vessel segmentation and artery-vein classification |
+| `--disc/--no-disc` | `--disc` | Run optic disc segmentation |
+| `--quality/--no-quality` | `--quality` | Run image quality assessment |
+| `--fovea/--no-fovea` | `--fovea` | Run fovea detection |
+| `--overlay/--no-overlay` | `config.yaml` or `--overlay` | Create visualization overlays combining all results |
+| `--config PATH` | auto-detect `config.yaml` | Load pipeline configuration from YAML |
+| `--n_jobs` | `4` | Number of preprocessing workers for parallel processing |
+
+### 📁 Output Structure
+
+When run with default options, the command creates the following structure in `OUTPUT_PATH`:
+
+```
+OUTPUT_PATH/
+├── preprocessed_rgb/     # Standardized fundus images
+├── vessels/              # Vessel segmentation results
+├── artery_vein/          # Artery-vein classification
+├── disc/                 # Optic disc segmentation
+├── disc_ring_2r/         # Binary masks for the 2r optic-disc ring
+├── disc_ring_3r/         # Binary masks for the 3r optic-disc ring
+├── overlays/             # Visualization images
+├── bounds.csv            # Image boundary information
+├── disc_geometry.csv     # Disc center and radius estimates in pixels
+├── quality.csv           # Image quality scores
+└── fovea.csv             # Fovea coordinates
+```
+
+### 🔄 Processing Stages
+
+1. **Preprocessing**: 
+   - Standardizes input images for consistent analysis
+   - Outputs preprocessed images and boundary information
+
+2. **Quality Assessment**:
+   - Evaluates image quality with three quality metrics (q1, q2, q3)
+   - Higher scores indicate better image quality
+
+3. **Vessel Segmentation and Artery-Vein Classification**:
+   - Identifies blood vessels in the retina
+   - Classifies vessels as arteries (1) or veins (2) with intersections (3)
+
+4. **Optic Disc Segmentation**:
+   - Identifies the optic disc location and boundaries
+   - Estimates disc center and radius from the disc mask
+   - Generates 2r and 3r ring masks around the disc
+
+5. **Fovea Detection**:
+   - Determines the coordinates of the fovea (center of vision)
+
+6. **Visualization Overlays**:
+   - Creates color-coded images showing:
+     - Arteries in red
+     - Veins in blue
+     - Optic disc in white
+     - 2r ring in green
+     - 3r ring in magenta
+     - Fovea marked with yellow X
+   - Overlay layers and colors can be controlled from `config.yaml`
+
+### ⚙️ `config.yaml`
+
+The repository root now includes a `config.yaml` file for overlay settings. The default file looks like this:
+
+```yaml
+overlay:
+  enabled: true
+  layers:
+    arteries: true
+    veins: true
+    disc: true
+    ring_2r: true
+    ring_3r: true
+    fovea: true
+  colours:
+    artery: "#FF0000"
+    vein: "#0000FF"
+    disc: "#FFFFFF"
+    ring_2r: "#00FF00"
+    ring_3r: "#FF00FF"
+    fovea: "#FFFF00"
+```
+
+Notes:
+
+- `overlay.enabled` controls whether overlays are produced when `--overlay/--no-overlay` is not set explicitly.
+- `overlay.layers` lets you choose which predictions are drawn.
+- `overlay.colors` and `overlay.colours` are both accepted.
+- Colors can be written as `#RRGGBB` strings or 3-value RGB arrays such as `[255, 0, 0]`.
+
+### 💻 Examples
+
+**Process a directory of images with all analyses:**
+```bash
+vascx run /path/to/images /path/to/output
+```
+
+**Process specific images listed in a CSV:**
+```bash
+vascx run /path/to/image_list.csv /path/to/output
+```
+
+**Only run preprocessing and vessel segmentation:**
+```bash
+vascx run /path/to/images /path/to/output --no-disc --no-quality --no-fovea --no-overlay
+```
+
+**Skip preprocessing on already preprocessed images:**
+```bash
+vascx run /path/to/preprocessed/images /path/to/output --no-preprocess
+```
+
+**Increase parallel processing workers:**
+```bash
+vascx run /path/to/images /path/to/output --n_jobs 8
+```
+
+### 📝 Notes
+
+- The CSV input must contain a 'path' column with image file paths
+- If the CSV includes an 'id' column, these IDs will be used instead of filenames
+- When `--no-preprocess` is used, input images must already be in the proper format
+- The overlay visualization requires at least one analysis component to be enabled
+
+## 📓 Notebooks
+
+For more advanced usage, we have Jupyter notebooks showing how preprocessing and inference are run.
+
+To speed up re-execution of vascx we recommend to run the preprocessing and segmentation steps separately:
+
+1. Preprocessing. See [this notebook](./notebooks/0_preprocess.ipynb). This step is CPU-heavy and benefits from parallelization (see notebook).
+
+2. Inference. See [this notebook](./notebooks/1_segment_preprocessed.ipynb). All models can be ran in a single GPU with >10GB VRAM.

artery_vein/av_july24.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b11d4e26ada8e1f0f279747afa5f4ef348d9d7350c4bf80e7ae6b5ac8d0b95b5
+size 352774102

artery_vein/av_july24_AVRDB.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:21fd6a693be9a5ffbf1b56e624612a493470d6e63025ce11f2e3886bd6f18b4c
+size 352791110

artery_vein/av_july24_IOSTAR.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:21fd6a693be9a5ffbf1b56e624612a493470d6e63025ce11f2e3886bd6f18b4c
+size 352791110

artery_vein/av_july24_LEUVEN.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:446580e6cda2acec8dc2ab30d9526735fc670f296048055cbb5ebb9ccac28d0b
+size 352830466

artery_vein/av_july24_RS.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c3de549b74ebd9c9a4f49c17043b831665cc7a1981773ff8c17db2416b9dfe48
+size 352805874

config.yaml

@@ -0,0 +1,16 @@
+overlay:
+  enabled: true
+  layers:
+    arteries: true
+    veins: true
+    disc: true
+    ring_2r: true
+    ring_3r: true
+    fovea: true
+  colours:
+    artery: "#FF0000"
+    vein: "#0000FF"
+    disc: "#FFFFFF"
+    ring_2r: "#00FF00"
+    ring_3r: "#FF00FF"
+    fovea: "#FFFF00"

disc/disc_july24.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6892b1bdb3bb68b666ea9a7891b0fb2f6fbb5fd4f05038c013c1c69ec6c7910c
+size 352801898

disc/disc_july24_ADAM.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9f423e047c88d9f6c03ada2c706bc84265d61ec473eaab28e8ca12a0f1738401
+size 352819138

disc/disc_july24_IDRID.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dbf239cfec2ee9b550aa09e3623af01cd79de688cac1c79902d0feb7d24bb3f7
+size 352835178

disc/disc_july24_ORIGA.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:18cc9ffac522fc9b91e46fe3aed8d6bb8fbf00e44bfb768b622f5b881713add6
+size 352826358

disc/disc_july24_PAPILA.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:080de64a1c005a921cb09eb04e42ce9c087dca2c1090209f77a3635af9eb1d19
+size 352832490

discedge/discedge_july24.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:891d8ef9bbc0676b019a81b1eceb349c1cdf4b5665a834196dd252915af64392
+size 352723146

environment.yml

@@ -0,0 +1,19 @@
+name: vascx-fork
+channels:
+  - conda-forge
+dependencies:
+  - python=3.11
+  - pip
+  - numpy=2.*
+  - pandas=2.*
+  - tqdm=4.*
+  - pillow=11.*
+  - click=8.*
+  - pyyaml=6.*
+  - pip:
+      - torch==2.11.0
+      - torchvision==0.26.0
+      - torchaudio==2.11.0
+      - retinalysis-fundusprep
+      - retinalysis-inference
+      - -e .

fovea/fovea_july24.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1af042f7e2a398f512be8a8d54cc480300312c3f3692c13bd98be65439a33222
+size 352714676

notebooks/0_preprocess.ipynb

@@ -0,0 +1,138 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from pathlib import Path\n",
+    "\n",
+    "import pandas as pd\n",
+    "\n",
+    "from rtnls_fundusprep.preprocessor import parallel_preprocess"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Preprocessing\n",
+    "\n",
+    "This code will preprocess the images and write .png files with the square fundus image and the contrast enhanced version\n",
+    "\n",
+    "This step is not strictly necessary, but it is useful if you want to run the preprocessing step separately before model inference\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Create a list of files to be preprocessed:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ds_path = Path(\"../samples/fundus\")\n",
+    "files = list((ds_path / \"original\").glob(\"*\"))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Images with .dcm extension will be read as dicom and the pixel_array will be read as RGB. All other images will be read using PIL's Image.open"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "0it [00:00, ?it/s][Parallel(n_jobs=4)]: Using backend LokyBackend with 4 concurrent workers.\n",
+      "6it [00:00, 154.80it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Error with image ../samples/fundus/original/HRF_07_dr.jpg\n",
+      "Error with image ../samples/fundus/original/HRF_04_g.jpg\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "[Parallel(n_jobs=4)]: Done   2 out of   6 | elapsed:    0.9s remaining:    1.8s\n",
+      "[Parallel(n_jobs=4)]: Done   3 out of   6 | elapsed:    1.5s remaining:    1.5s\n",
+      "[Parallel(n_jobs=4)]: Done   4 out of   6 | elapsed:    1.5s remaining:    0.8s\n",
+      "[Parallel(n_jobs=4)]: Done   6 out of   6 | elapsed:    1.6s finished\n"
+     ]
+    }
+   ],
+   "source": [
+    "bounds = parallel_preprocess(\n",
+    "    files,  # List of image files\n",
+    "    rgb_path=ds_path / \"rgb\",  # Output path for RGB images\n",
+    "    ce_path=ds_path / \"ce\",  # Output path for Contrast Enhanced images\n",
+    "    n_jobs=4,  # number of preprocessing workers\n",
+    ")\n",
+    "df_bounds = pd.DataFrame(bounds).set_index(\"id\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "The preprocessor will produce RGB and contrast-enhanced preprocessed images cropped to a square and return a dataframe with the image bounds that can be used to reconstruct the original image. Output files will be named the same as input images, but with .png extension. Be careful with providing multiple inputs with the same filename without extension as this will result in over-written images. Any exceptions during pre-processing will not stop execution but will print error. Images that failed pre-processing for any reason will be marked with `success=False` in the df_bounds dataframe."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "df_bounds.to_csv(ds_path / \"meta.csv\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "retinalysis",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.13"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

notebooks/1_segment_preprocessed.ipynb

@@ -0,0 +1,217 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from pathlib import Path\n",
+    "\n",
+    "import torch\n",
+    "\n",
+    "from rtnls_inference import (\n",
+    "    HeatmapRegressionEnsemble,\n",
+    "    SegmentationEnsemble,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Segmentation of preprocessed images\n",
+    "\n",
+    "Here we segment images preprocessed using 0_preprocess.ipynb\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ds_path = Path(\"../samples/fundus\")\n",
+    "\n",
+    "# input folders. these are the folders where we stored the preprocessed images\n",
+    "rgb_path = ds_path / \"rgb\"\n",
+    "ce_path = ds_path / \"ce\"\n",
+    "\n",
+    "# these are the output folders for:\n",
+    "av_path = ds_path / \"av\"  # artery-vein segmentations\n",
+    "discs_path = ds_path / \"discs\"  # optic disc segmentations\n",
+    "overlays_path = ds_path / \"overlays\"  # optional overlay visualizations\n",
+    "\n",
+    "device = torch.device(\"cuda:0\")  # device to use for inference"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "rgb_paths = sorted(list(rgb_path.glob(\"*.png\")))\n",
+    "ce_paths = sorted(list(ce_path.glob(\"*.png\")))\n",
+    "paired_paths = list(zip(rgb_paths, ce_paths))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "paired_paths[0]  # important to make sure that the paths are paired correctly"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Artery-vein segmentation\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "av_ensemble = SegmentationEnsemble.from_huggingface('zyf0717/vascx-fork:artery_vein/av_july24.pt').to(device)\n",
+    "\n",
+    "av_ensemble.predict_preprocessed(paired_paths, dest_path=av_path, num_workers=2)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Disc segmentation\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "disc_ensemble = SegmentationEnsemble.from_huggingface('zyf0717/vascx-fork:disc/disc_july24.pt').to(device)\n",
+    "disc_ensemble.predict_preprocessed(paired_paths, dest_path=discs_path, num_workers=2)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Fovea detection\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fovea_ensemble = HeatmapRegressionEnsemble.from_huggingface('zyf0717/vascx-fork:fovea/fovea_july24.pt').to(device)\n",
+    "# note: this model does not use contrast enhanced images\n",
+    "df = fovea_ensemble.predict_preprocessed(paired_paths, num_workers=2)\n",
+    "df.columns = [\"mean_x\", \"mean_y\"]\n",
+    "df.to_csv(ds_path / \"fovea.csv\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "df"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Plotting the retinas (optional)\n",
+    "\n",
+    "This will only work if you ran all the models and stored the outputs using the same folder/file names as above\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from vascx.fundus.loader import RetinaLoader\n",
+    "\n",
+    "from rtnls_enface.utils.plotting import plot_gridfns\n",
+    "\n",
+    "loader = RetinaLoader.from_folder(ds_path)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "plot_gridfns([ret.plot for ret in loader[:6]])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Storing visualizations (optional)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "if not overlays_path.exists():\n",
+    "    overlays_path.mkdir()\n",
+    "for ret in loader:\n",
+    "    fig, _ = ret.plot()\n",
+    "    fig.savefig(overlays_path / f\"{ret.id}.png\", bbox_inches=\"tight\", pad_inches=0)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "retinalysis",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.13"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

odfd/odfd_march25.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:aa2be119eb915bc9da6ba42234f703b5cc270d53b62c1e7d7e1bdff52c1e0edd
+size 855538988

quality/quality.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:80034ccb21a57522ba0cb86be0d46d2b659e193b86db909ad6ec2e85e61f87aa
+size 855578258

run.sh

@@ -0,0 +1,60 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+CONDA_ENV="${CONDA_ENV:-vascx-fork}"
+SAMPLE_INPUT_PATH="$REPO_ROOT/samples/fundus/original"
+DEFAULT_INPUT_PATH="$SAMPLE_INPUT_PATH"
+INPUT_PATH="${INPUT_PATH:-$DEFAULT_INPUT_PATH}"
+TIMESTAMP="$(date +"%Y%m%d_%H%M%S")"
+DEFAULT_OUTPUT_PATH="$REPO_ROOT/output_$TIMESTAMP"
+OUTPUT_PATH="${OUTPUT_PATH:-$DEFAULT_OUTPUT_PATH}"
+N_JOBS="${N_JOBS:-1}"
+MODEL_RELEASES_DIR="$REPO_ROOT/model_releases"
+
+while [[ $# -gt 0 ]]; do
+  case "$1" in
+    --sample-run)
+      INPUT_PATH="$SAMPLE_INPUT_PATH"
+      shift
+      ;;
+    *)
+      echo "Unknown argument: $1" >&2
+      echo "Usage: $0 [--sample-run]" >&2
+      exit 1
+      ;;
+  esac
+done
+
+if [[ ! -d "$INPUT_PATH" ]]; then
+  echo "Input directory does not exist: $INPUT_PATH" >&2
+  exit 1
+fi
+
+mkdir -p "$REPO_ROOT/.mplconfig" "$REPO_ROOT/.cache" "$MODEL_RELEASES_DIR" "$OUTPUT_PATH"
+
+for model_path in "$REPO_ROOT"/*/*.pt; do
+  [[ -e "$model_path" ]] || continue
+  if [[ "$model_path" == "$MODEL_RELEASES_DIR"/* ]]; then
+    continue
+  fi
+  ln -sf "$model_path" "$MODEL_RELEASES_DIR/$(basename "$model_path")"
+done
+
+export MPLCONFIGDIR="$REPO_ROOT/.mplconfig"
+export XDG_CACHE_HOME="$REPO_ROOT/.cache"
+export RTNLS_MODEL_RELEASES="$MODEL_RELEASES_DIR"
+
+echo "Running VascX Fork"
+echo "  conda env:   $CONDA_ENV"
+echo "  input path:  $INPUT_PATH"
+echo "  output path: $OUTPUT_PATH"
+echo "  n_jobs:      $N_JOBS"
+echo "  models dir:  $RTNLS_MODEL_RELEASES"
+
+CONDA_BASE="$(conda info --base)"
+# shellcheck disable=SC1091
+source "$CONDA_BASE/etc/profile.d/conda.sh"
+conda activate "$CONDA_ENV"
+
+exec python -c "from vascx_models.cli import cli; cli()" run "$INPUT_PATH" "$OUTPUT_PATH" --n_jobs "$N_JOBS"

setup.py

@@ -0,0 +1,36 @@
+from setuptools import find_packages, setup
+
+with open("README.md", "r") as fh:
+    long_description = fh.read()
+
+setup(
+    name="vascx_models",
+    # using versioneer for versioning using git tags
+    # https://github.com/python-versioneer/python-versioneer/blob/master/INSTALL.md
+    # version=versioneer.get_version(),
+    # cmdclass=versioneer.get_cmdclass(),
+    author="Jose Vargas",
+    author_email="j.vargasquiros@erasmusmc.nl",
+    description="Retinal analysis toolbox for Python",
+    long_description=long_description,
+    long_description_content_type="text/markdown",
+    packages=find_packages(),
+    include_package_data=True,
+    zip_safe=False,
+    entry_points={
+        "console_scripts": [
+            "vascx = vascx_models.cli:cli",
+        ]
+    },
+    install_requires=[
+        "numpy == 2.*",
+        "pandas == 2.*",
+        "tqdm == 4.*",
+        "Pillow == 11.*",
+        "click==8.*",
+        "PyYAML == 6.*",
+        "retinalysis-fundusprep",
+        "retinalysis-inference",
+    ],
+    python_requires=">=3.10, <3.13",
+)

vascx_models/cli.py

@@ -0,0 +1,259 @@
+import logging
+import warnings
+from pathlib import Path
+
+import click
+import pandas as pd
+
+from rtnls_fundusprep.cli import _run_preprocessing
+
+from .config import load_app_config
+from .disc_rings import generate_disc_rings
+from .inference import (
+    preferred_device,
+    run_fovea_detection,
+    run_quality_estimation,
+    run_segmentation_disc,
+    run_segmentation_vessels_and_av,
+)
+from .utils import batch_create_overlays
+
+logger = logging.getLogger(__name__)
+
+
+def configure_logging() -> None:
+    logging.basicConfig(level=logging.INFO, format="[%(levelname)s] %(message)s")
+    warnings.filterwarnings(
+        "ignore",
+        message=(
+            "Using a non-tuple sequence for multidimensional indexing is deprecated "
+            "and will be changed in pytorch 2.9; use x\\[tuple\\(seq\\)\\] instead of x\\[seq\\].*"
+        ),
+        category=UserWarning,
+        module=r"monai\.inferers\.utils",
+    )
+
+
+@click.group(name="vascx")
+def cli():
+    configure_logging()
+
+
+@cli.command()
+@click.argument("data_path", type=click.Path(exists=True))
+@click.argument("output_path", type=click.Path())
+@click.option(
+    "--config",
+    "config_path",
+    type=click.Path(exists=True, dir_okay=False, path_type=Path),
+    default=None,
+    help="Path to a YAML config file. Defaults to ./config.yaml or the repo-root config.yaml when present.",
+)
+@click.option(
+    "--preprocess/--no-preprocess",
+    default=True,
+    help="Run preprocessing or use preprocessed images",
+)
+@click.option(
+    "--vessels/--no-vessels", default=True, help="Run vessels and AV segmentation"
+)
+@click.option("--disc/--no-disc", default=True, help="Run optic disc segmentation")
+@click.option(
+    "--quality/--no-quality", default=True, help="Run image quality estimation"
+)
+@click.option("--fovea/--no-fovea", default=True, help="Run fovea detection")
+@click.option(
+    "--overlay/--no-overlay",
+    default=None,
+    help="Create visualization overlays. Defaults to the config value when set.",
+)
+@click.option("--n_jobs", type=int, default=4, help="Number of preprocessing workers")
+def run(
+    data_path,
+    output_path,
+    config_path,
+    preprocess,
+    vessels,
+    disc,
+    quality,
+    fovea,
+    overlay,
+    n_jobs,
+):
+    """Run the complete inference pipeline on fundus images.
+
+    DATA_PATH is either a directory containing images or a CSV file with 'path' column.
+    OUTPUT_PATH is the directory where results will be stored.
+    """
+
+    output_path = Path(output_path)
+    output_path.mkdir(exist_ok=True, parents=True)
+    try:
+        app_config = load_app_config(config_path)
+    except (FileNotFoundError, ValueError) as exc:
+        raise click.ClickException(str(exc)) from exc
+    overlay_enabled = app_config.overlay.enabled if overlay is None else overlay
+    if app_config.source_path is not None:
+        logger.info("Loaded config from %s", app_config.source_path)
+
+    # Setup output directories
+    preprocess_rgb_path = output_path / "preprocessed_rgb"
+    vessels_path = output_path / "vessels"
+    av_path = output_path / "artery_vein"
+    disc_path = output_path / "disc"
+    disc_ring_2r_path = output_path / "disc_ring_2r"
+    disc_ring_3r_path = output_path / "disc_ring_3r"
+    overlay_path = output_path / "overlays"
+
+    # Create required directories
+    if preprocess:
+        preprocess_rgb_path.mkdir(exist_ok=True, parents=True)
+    if vessels:
+        av_path.mkdir(exist_ok=True, parents=True)
+        vessels_path.mkdir(exist_ok=True, parents=True)
+    if disc:
+        disc_path.mkdir(exist_ok=True, parents=True)
+        disc_ring_2r_path.mkdir(exist_ok=True, parents=True)
+        disc_ring_3r_path.mkdir(exist_ok=True, parents=True)
+    if overlay_enabled:
+        overlay_path.mkdir(exist_ok=True, parents=True)
+
+    bounds_path = output_path / "bounds.csv" if preprocess else None
+    quality_path = output_path / "quality.csv" if quality else None
+    fovea_path = output_path / "fovea.csv" if fovea else None
+    disc_geometry_path = output_path / "disc_geometry.csv" if disc else None
+
+    # Determine if input is a folder or CSV file
+    data_path = Path(data_path)
+    is_csv = data_path.suffix.lower() == ".csv"
+
+    # Get files to process
+    files = []
+    ids = None
+
+    if is_csv:
+        logger.info("Reading file paths from CSV: %s", data_path)
+        try:
+            df = pd.read_csv(data_path)
+            if "path" not in df.columns:
+                logger.error("CSV must contain a 'path' column")
+                return
+
+            # Get file paths and convert to Path objects
+            files = [Path(p) for p in df["path"]]
+
+            if "id" in df.columns:
+                ids = df["id"].tolist()
+                logger.info("Using IDs from CSV 'id' column")
+
+        except Exception as e:
+            logger.exception("Error reading CSV file: %s", e)
+            return
+    else:
+        logger.info("Finding files in directory: %s", data_path)
+        files = list(data_path.glob("*"))
+        ids = [f.stem for f in files]
+
+    if not files:
+        logger.warning("No files found to process")
+        return
+
+    logger.info("Found %d files to process", len(files))
+
+    # Step 1: Preprocess images if requested
+    if preprocess:
+        logger.info("Running preprocessing")
+        _run_preprocessing(
+            files=files,
+            ids=ids,
+            rgb_path=preprocess_rgb_path,
+            bounds_path=bounds_path,
+            n_jobs=n_jobs,
+        )
+        # Use the preprocessed images for subsequent steps
+        preprocessed_files = list(preprocess_rgb_path.glob("*.png"))
+    else:
+        # Use the input files directly
+        preprocessed_files = files
+    ids = [f.stem for f in preprocessed_files]
+    logger.info("Prepared %d images for inference", len(preprocessed_files))
+
+    # Prefer hardware acceleration when the active torch build supports it.
+    device = preferred_device()
+    logger.info("Using device: %s", device)
+
+    # Step 2: Run quality estimation if requested
+    if quality:
+        logger.info("Running quality estimation")
+        df_quality = run_quality_estimation(
+            fpaths=preprocessed_files, ids=ids, device=device
+        )
+        df_quality.to_csv(quality_path)
+        logger.info("Quality results saved to %s", quality_path)
+
+    # Step 3: Run vessels and AV segmentation if requested
+    if vessels:
+        logger.info("Running vessels and AV segmentation")
+        run_segmentation_vessels_and_av(
+            rgb_paths=preprocessed_files,
+            ids=ids,
+            av_path=av_path,
+            vessels_path=vessels_path,
+            device=device,
+        )
+        logger.info("Vessel segmentation saved to %s", vessels_path)
+        logger.info("AV segmentation saved to %s", av_path)
+
+    # Step 4: Run optic disc segmentation if requested
+    if disc:
+        logger.info("Running optic disc segmentation")
+        run_segmentation_disc(
+            rgb_paths=preprocessed_files, ids=ids, output_path=disc_path, device=device
+        )
+        logger.info("Disc segmentation saved to %s", disc_path)
+        generate_disc_rings(
+            disc_dir=disc_path,
+            ring_2r_dir=disc_ring_2r_path,
+            ring_3r_dir=disc_ring_3r_path,
+            measurements_path=disc_geometry_path,
+        )
+        logger.info("2r disc rings saved to %s", disc_ring_2r_path)
+        logger.info("3r disc rings saved to %s", disc_ring_3r_path)
+
+    # Step 5: Run fovea detection if requested
+    df_fovea = None
+    if fovea:
+        logger.info("Running fovea detection")
+        df_fovea = run_fovea_detection(
+            rgb_paths=preprocessed_files, ids=ids, device=device
+        )
+        df_fovea.to_csv(fovea_path)
+        logger.info("Fovea detection results saved to %s", fovea_path)
+
+    # Step 6: Create overlays if requested
+    if overlay_enabled:
+        logger.info("Creating visualization overlays")
+
+        # Prepare fovea data if available
+        fovea_data = None
+        if df_fovea is not None:
+            fovea_data = {
+                idx: (row["x_fovea"], row["y_fovea"])
+                for idx, row in df_fovea.iterrows()
+            }
+
+        # Create visualization overlays
+        batch_create_overlays(
+            rgb_dir=preprocess_rgb_path if preprocess else data_path,
+            output_dir=overlay_path,
+            av_dir=av_path,
+            disc_dir=disc_path,
+            ring_2r_dir=disc_ring_2r_path,
+            ring_3r_dir=disc_ring_3r_path,
+            fovea_data=fovea_data,
+            overlay_config=app_config.overlay,
+        )
+
+        logger.info("Visualization overlays saved to %s", overlay_path)
+
+    logger.info("All requested processing complete. Results saved to %s", output_path)

vascx_models/config.py

@@ -0,0 +1,196 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Iterable, Mapping
+
+import yaml
+
+
+DEFAULT_CONFIG_NAME = "config.yaml"
+
+
+def _repo_root() -> Path:
+    return Path(__file__).resolve().parent.parent
+
+
+@dataclass(frozen=True)
+class OverlayLayers:
+    arteries: bool = True
+    veins: bool = True
+    disc: bool = True
+    ring_2r: bool = True
+    ring_3r: bool = True
+    fovea: bool = True
+
+
+@dataclass(frozen=True)
+class OverlayColors:
+    artery: tuple[int, int, int] = (255, 0, 0)
+    vein: tuple[int, int, int] = (0, 0, 255)
+    disc: tuple[int, int, int] = (255, 255, 255)
+    ring_2r: tuple[int, int, int] = (0, 255, 0)
+    ring_3r: tuple[int, int, int] = (255, 0, 255)
+    fovea: tuple[int, int, int] = (255, 255, 0)
+
+
+@dataclass(frozen=True)
+class OverlayConfig:
+    enabled: bool = True
+    layers: OverlayLayers = field(default_factory=OverlayLayers)
+    colors: OverlayColors = field(default_factory=OverlayColors)
+
+
+@dataclass(frozen=True)
+class AppConfig:
+    overlay: OverlayConfig = field(default_factory=OverlayConfig)
+    source_path: Path | None = None
+
+
+def default_config_candidates() -> list[Path]:
+    candidates = [Path.cwd() / DEFAULT_CONFIG_NAME, _repo_root() / DEFAULT_CONFIG_NAME]
+    unique_candidates: list[Path] = []
+    seen: set[Path] = set()
+    for candidate in candidates:
+        resolved = candidate.resolve()
+        if resolved not in seen:
+            unique_candidates.append(candidate)
+            seen.add(resolved)
+    return unique_candidates
+
+
+def resolve_config_path(config_path: str | Path | None) -> Path | None:
+    if config_path is not None:
+        candidate = Path(config_path).expanduser()
+        if not candidate.exists():
+            raise FileNotFoundError(f"Config file not found: {candidate}")
+        return candidate
+
+    for candidate in default_config_candidates():
+        if candidate.exists():
+            return candidate
+    return None
+
+
+def load_app_config(config_path: str | Path | None = None) -> AppConfig:
+    resolved_path = resolve_config_path(config_path)
+    if resolved_path is None:
+        return AppConfig()
+
+    with resolved_path.open("r", encoding="utf-8") as handle:
+        raw_config = yaml.safe_load(handle) or {}
+
+    if not isinstance(raw_config, dict):
+        raise ValueError("Config root must be a mapping")
+
+    overlay_raw = raw_config.get("overlay", {})
+    if overlay_raw is None:
+        overlay_raw = {}
+    if not isinstance(overlay_raw, dict):
+        raise ValueError("'overlay' must be a mapping")
+
+    layer_overrides = overlay_raw.get("layers", {})
+    if layer_overrides is None:
+        layer_overrides = {}
+    if not isinstance(layer_overrides, dict):
+        raise ValueError("'overlay.layers' must be a mapping")
+
+    color_overrides = overlay_raw.get("colors", overlay_raw.get("colours", {}))
+    if color_overrides is None:
+        color_overrides = {}
+    if not isinstance(color_overrides, dict):
+        raise ValueError("'overlay.colors' must be a mapping")
+
+    return AppConfig(
+        overlay=OverlayConfig(
+            enabled=_coerce_bool(overlay_raw.get("enabled", True), "overlay.enabled"),
+            layers=_build_overlay_layers(layer_overrides),
+            colors=_build_overlay_colors(color_overrides),
+        ),
+        source_path=resolved_path,
+    )
+
+
+def _build_overlay_layers(raw_layers: Mapping[str, object]) -> OverlayLayers:
+    defaults = OverlayLayers()
+    alias_map = {
+        "artery": "arteries",
+        "arteries": "arteries",
+        "vein": "veins",
+        "veins": "veins",
+        "disc": "disc",
+        "ring_2r": "ring_2r",
+        "disc_ring_2r": "ring_2r",
+        "ring_3r": "ring_3r",
+        "disc_ring_3r": "ring_3r",
+        "fovea": "fovea",
+    }
+    values = defaults.__dict__.copy()
+    for raw_key, raw_value in raw_layers.items():
+        if raw_key not in alias_map:
+            raise ValueError(f"Unsupported overlay layer '{raw_key}'")
+        normalized_key = alias_map[raw_key]
+        values[normalized_key] = _coerce_bool(
+            raw_value, f"overlay.layers.{raw_key}"
+        )
+    return OverlayLayers(**values)
+
+
+def _build_overlay_colors(raw_colors: Mapping[str, object]) -> OverlayColors:
+    defaults = OverlayColors()
+    alias_map = {
+        "artery": "artery",
+        "arteries": "artery",
+        "vein": "vein",
+        "veins": "vein",
+        "disc": "disc",
+        "ring_2r": "ring_2r",
+        "disc_ring_2r": "ring_2r",
+        "ring_3r": "ring_3r",
+        "disc_ring_3r": "ring_3r",
+        "fovea": "fovea",
+    }
+    values = defaults.__dict__.copy()
+    for raw_key, raw_value in raw_colors.items():
+        if raw_key not in alias_map:
+            raise ValueError(f"Unsupported overlay color '{raw_key}'")
+        normalized_key = alias_map[raw_key]
+        values[normalized_key] = _parse_rgb(raw_value, f"overlay.colors.{raw_key}")
+    return OverlayColors(**values)
+
+
+def _coerce_bool(value: object, field_name: str) -> bool:
+    if isinstance(value, bool):
+        return value
+    raise ValueError(f"'{field_name}' must be a boolean")
+
+
+def _parse_rgb(value: object, field_name: str) -> tuple[int, int, int]:
+    if isinstance(value, str):
+        return _parse_hex_color(value, field_name)
+    if isinstance(value, Iterable) and not isinstance(value, (str, bytes, dict)):
+        channels = tuple(value)
+        if len(channels) != 3:
+            raise ValueError(f"'{field_name}' must contain exactly 3 channels")
+        return tuple(_coerce_channel(channel, field_name) for channel in channels)
+    raise ValueError(
+        f"'{field_name}' must be a '#RRGGBB' string or a 3-item RGB sequence"
+    )
+
+
+def _parse_hex_color(value: str, field_name: str) -> tuple[int, int, int]:
+    normalized = value.strip()
+    if normalized.startswith("#"):
+        normalized = normalized[1:]
+    if len(normalized) != 6:
+        raise ValueError(f"'{field_name}' must be a 6-digit hex color")
+    try:
+        return tuple(int(normalized[index : index + 2], 16) for index in (0, 2, 4))
+    except ValueError as exc:
+        raise ValueError(f"'{field_name}' must be a valid hex color") from exc
+
+
+def _coerce_channel(value: object, field_name: str) -> int:
+    if isinstance(value, int) and 0 <= value <= 255:
+        return value
+    raise ValueError(f"'{field_name}' channels must be integers between 0 and 255")

vascx_models/disc_rings.py

@@ -0,0 +1,118 @@
+import logging
+from pathlib import Path
+from typing import Dict, Optional, Tuple
+
+import numpy as np
+import pandas as pd
+from PIL import Image, ImageDraw
+
+logger = logging.getLogger(__name__)
+
+
+def estimate_disc_geometry(
+    disc_mask: np.ndarray,
+) -> Optional[Tuple[float, float, float]]:
+    """Estimate optic disc center and radius from a binary mask."""
+    mask = disc_mask > 0
+    if not np.any(mask):
+        return None
+
+    ys, xs = np.nonzero(mask)
+    center_x = float(xs.mean())
+    center_y = float(ys.mean())
+
+    # Use the equivalent-circle radius so the estimate is stable for irregular masks.
+    radius = float(np.sqrt(mask.sum() / np.pi))
+    return center_x, center_y, radius
+
+
+def create_ring_mask(
+    image_shape: Tuple[int, int],
+    center: Tuple[float, float],
+    radius: float,
+    thickness: int,
+) -> np.ndarray:
+    """Create a binary ring mask for a circle outline."""
+    height, width = image_shape
+    ring = Image.new("L", (width, height), 0)
+    draw = ImageDraw.Draw(ring)
+
+    center_x, center_y = center
+    bbox = (
+        center_x - radius,
+        center_y - radius,
+        center_x + radius,
+        center_y + radius,
+    )
+    draw.ellipse(bbox, outline=255, width=thickness)
+    return np.array(ring, dtype=np.uint8)
+
+
+def generate_disc_rings(
+    disc_dir: Path,
+    ring_2r_dir: Path,
+    ring_3r_dir: Path,
+    measurements_path: Optional[Path] = None,
+) -> pd.DataFrame:
+    """Generate 2r and 3r optic-disc ring masks from saved disc segmentations."""
+    ring_2r_dir.mkdir(exist_ok=True, parents=True)
+    ring_3r_dir.mkdir(exist_ok=True, parents=True)
+
+    disc_files = list(disc_dir.glob("*.png"))
+    if not disc_files:
+        logger.warning("No disc masks found for ring generation in %s", disc_dir)
+        columns = [
+            "x_disc_center",
+            "y_disc_center",
+            "disc_radius_px",
+            "ring_2r_px",
+            "ring_3r_px",
+        ]
+        return pd.DataFrame(columns=columns)
+
+    records: Dict[str, Dict[str, float]] = {}
+    logger.info("Generating 2r and 3r rings for %d disc masks", len(disc_files))
+
+    for disc_file in disc_files:
+        image_id = disc_file.stem
+        disc_mask = np.array(Image.open(disc_file)) > 0
+        geometry = estimate_disc_geometry(disc_mask)
+
+        if geometry is None:
+            logger.warning("Disc mask is empty for %s; writing blank ring masks", image_id)
+            blank = np.zeros(disc_mask.shape, dtype=np.uint8)
+            Image.fromarray(blank).save(ring_2r_dir / f"{image_id}.png")
+            Image.fromarray(blank).save(ring_3r_dir / f"{image_id}.png")
+            records[image_id] = {
+                "x_disc_center": np.nan,
+                "y_disc_center": np.nan,
+                "disc_radius_px": np.nan,
+                "ring_2r_px": np.nan,
+                "ring_3r_px": np.nan,
+            }
+            continue
+
+        center_x, center_y, disc_radius = geometry
+        line_width = max(1, int(round(disc_radius * 0.08)))
+        ring_2r = create_ring_mask(
+            disc_mask.shape, (center_x, center_y), radius=disc_radius * 2.0, thickness=line_width
+        )
+        ring_3r = create_ring_mask(
+            disc_mask.shape, (center_x, center_y), radius=disc_radius * 3.0, thickness=line_width
+        )
+
+        Image.fromarray(ring_2r).save(ring_2r_dir / f"{image_id}.png")
+        Image.fromarray(ring_3r).save(ring_3r_dir / f"{image_id}.png")
+        records[image_id] = {
+            "x_disc_center": center_x,
+            "y_disc_center": center_y,
+            "disc_radius_px": disc_radius,
+            "ring_2r_px": disc_radius * 2.0,
+            "ring_3r_px": disc_radius * 3.0,
+        }
+
+    df_measurements = pd.DataFrame.from_dict(records, orient="index")
+    if measurements_path is not None:
+        df_measurements.to_csv(measurements_path)
+        logger.info("Disc ring measurements saved to %s", measurements_path)
+    return df_measurements

vascx_models/inference.py

@@ -0,0 +1,292 @@
+import logging
+import os
+from contextlib import nullcontext
+from pathlib import Path
+from typing import List, Optional
+
+import numpy as np
+import pandas as pd
+import torch
+from PIL import Image
+from tqdm import tqdm
+
+from rtnls_inference.ensembles.ensemble_classification import ClassificationEnsemble
+from rtnls_inference.ensembles.ensemble_heatmap_regression import (
+    HeatmapRegressionEnsemble,
+)
+from rtnls_inference.ensembles.ensemble_segmentation import SegmentationEnsemble
+from rtnls_inference.utils import decollate_batch, extract_keypoints_from_heatmaps
+
+logger = logging.getLogger(__name__)
+
+
+def preferred_device() -> torch.device:
+    if torch.cuda.is_available():
+        return torch.device("cuda:0")
+    if torch.backends.mps.is_available():
+        return torch.device("mps")
+    return torch.device("cpu")
+
+
+def _inference_num_workers(device: torch.device) -> int:
+    # Torch shared-memory workers can fail in restricted CPU environments.
+    return 8 if device.type in {"cuda", "mps"} else 0
+
+
+def _autocast_context(device: torch.device):
+    return torch.autocast(device_type=device.type) if device.type == "cuda" else nullcontext()
+
+
+def run_quality_estimation(fpaths, ids, device: torch.device):
+    logger.info("Loading quality model on %s", device)
+    ensemble_quality = ClassificationEnsemble.from_release("quality.pt").to(device)
+    dataloader = ensemble_quality._make_inference_dataloader(
+        fpaths,
+        ids=ids,
+        num_workers=_inference_num_workers(device),
+        preprocess=False,
+        batch_size=16,
+    )
+    logger.info("Quality dataloader ready with %d images", len(fpaths))
+
+    output_ids, outputs = [], []
+    with torch.no_grad():
+        for batch in tqdm(dataloader):
+            if len(batch) == 0:
+                continue
+
+            im = batch["image"].to(device)
+
+            # QUALITY
+            quality = ensemble_quality.predict_step(im)
+            quality = torch.mean(quality, dim=0)
+
+            items = {"id": batch["id"], "quality": quality}
+            items = decollate_batch(items)
+
+            for item in items:
+                output_ids.append(item["id"])
+                outputs.append(item["quality"].tolist())
+
+    return pd.DataFrame(
+        outputs,
+        index=output_ids,
+        columns=["q1", "q2", "q3"],
+    )
+
+
+def run_segmentation_vessels_and_av(
+    rgb_paths: List[Path],
+    ce_paths: Optional[List[Path]] = None,
+    ids: Optional[List[str]] = None,
+    av_path: Optional[Path] = None,
+    vessels_path: Optional[Path] = None,
+    device: torch.device = preferred_device(),
+) -> None:
+    """
+    Run AV and vessel segmentation on the provided images.
+
+    Args:
+        rgb_paths: List of paths to RGB fundus images
+        ce_paths: Optional list of paths to contrast enhanced images
+        ids: Optional list of ids to pass to _make_inference_dataloader
+        av_path: Folder where to store output AV segmentations
+        vessels_path: Folder where to store output vessel segmentations
+        device: Device to run inference on
+    """
+    # Create output directories if they don't exist
+    if av_path is not None:
+        av_path.mkdir(exist_ok=True, parents=True)
+    if vessels_path is not None:
+        vessels_path.mkdir(exist_ok=True, parents=True)
+
+    # Load models
+    logger.info("Loading AV and vessel models on %s", device)
+    ensemble_av = SegmentationEnsemble.from_release("av_july24.pt").to(device).eval()
+    ensemble_vessels = (
+        SegmentationEnsemble.from_release("vessels_july24.pt").to(device).eval()
+    )
+
+    # Prepare input paths
+    if ce_paths is None:
+        # If CE paths are not provided, use RGB paths for both inputs
+        fpaths = rgb_paths
+    else:
+        # If CE paths are provided, pair them with RGB paths
+        if len(rgb_paths) != len(ce_paths):
+            raise ValueError("rgb_paths and ce_paths must have the same length")
+        fpaths = list(zip(rgb_paths, ce_paths))
+
+    # Create dataloader
+    dataloader = ensemble_av._make_inference_dataloader(
+        fpaths,
+        ids=ids,
+        num_workers=_inference_num_workers(device),
+        preprocess=False,
+        batch_size=8,
+    )
+    logger.info("AV and vessel dataloader ready with %d images", len(fpaths))
+
+    # Run inference
+    with torch.no_grad():
+        for batch in tqdm(dataloader):
+            # AV segmentation
+            if av_path is not None:
+                with _autocast_context(device):
+                    proba = ensemble_av.forward(batch["image"].to(device))
+                proba = torch.mean(proba, dim=1)  # average over models
+                proba = torch.permute(proba, (0, 2, 3, 1))  # NCHW -> NHWC
+                proba = torch.nn.functional.softmax(proba, dim=-1)
+
+                items = {
+                    "id": batch["id"],
+                    "image": proba,
+                }
+
+                items = decollate_batch(items)
+                for i, item in enumerate(items):
+                    fpath = os.path.join(av_path, f"{item['id']}.png")
+                    mask = np.argmax(item["image"], -1)
+                    Image.fromarray(mask.squeeze().astype(np.uint8)).save(fpath)
+
+            # Vessel segmentation
+            if vessels_path is not None:
+                with _autocast_context(device):
+                    proba = ensemble_vessels.forward(batch["image"].to(device))
+                proba = torch.mean(proba, dim=1)  # average over models
+                proba = torch.permute(proba, (0, 2, 3, 1))  # NCHW -> NHWC
+                proba = torch.nn.functional.softmax(proba, dim=-1)
+
+                items = {
+                    "id": batch["id"],
+                    "image": proba,
+                }
+
+                items = decollate_batch(items)
+                for i, item in enumerate(items):
+                    fpath = os.path.join(vessels_path, f"{item['id']}.png")
+                    mask = np.argmax(item["image"], -1)
+                    Image.fromarray(mask.squeeze().astype(np.uint8)).save(fpath)
+
+
+def run_segmentation_disc(
+    rgb_paths: List[Path],
+    ce_paths: Optional[List[Path]] = None,
+    ids: Optional[List[str]] = None,
+    output_path: Optional[Path] = None,
+    device: torch.device = preferred_device(),
+) -> None:
+    logger.info("Loading disc model on %s", device)
+    ensemble_disc = (
+        SegmentationEnsemble.from_release("disc_july24.pt").to(device).eval()
+    )
+
+    # Prepare input paths
+    if ce_paths is None:
+        # If CE paths are not provided, use RGB paths for both inputs
+        fpaths = rgb_paths
+    else:
+        # If CE paths are provided, pair them with RGB paths
+        if len(rgb_paths) != len(ce_paths):
+            raise ValueError("rgb_paths and ce_paths must have the same length")
+        fpaths = list(zip(rgb_paths, ce_paths))
+
+    dataloader = ensemble_disc._make_inference_dataloader(
+        fpaths,
+        ids=ids,
+        num_workers=_inference_num_workers(device),
+        preprocess=False,
+        batch_size=8,
+    )
+    logger.info("Disc dataloader ready with %d images", len(fpaths))
+
+    with torch.no_grad():
+        for batch in tqdm(dataloader):
+            # AV
+            with _autocast_context(device):
+                proba = ensemble_disc.forward(batch["image"].to(device))
+            proba = torch.mean(proba, dim=1)  # average over models
+            proba = torch.permute(proba, (0, 2, 3, 1))  # NCHW -> NHWC
+            proba = torch.nn.functional.softmax(proba, dim=-1)
+
+            items = {
+                "id": batch["id"],
+                "image": proba,
+            }
+
+            items = decollate_batch(items)
+            items = [dataloader.dataset.transform.undo_item(item) for item in items]
+            for i, item in enumerate(items):
+                fpath = os.path.join(output_path, f"{item['id']}.png")
+
+                mask = np.argmax(item["image"], -1)
+                Image.fromarray(mask.squeeze().astype(np.uint8)).save(fpath)
+
+
+def run_fovea_detection(
+    rgb_paths: List[Path],
+    ce_paths: Optional[List[Path]] = None,
+    ids: Optional[List[str]] = None,
+    device: torch.device = preferred_device(),
+) -> None:
+    # def run_fovea_detection(fpaths, ids, device: torch.device):
+    logger.info("Loading fovea model on %s", device)
+    ensemble_fovea = HeatmapRegressionEnsemble.from_release("fovea_july24.pt").to(
+        device
+    )
+
+    # Prepare input paths
+    if ce_paths is None:
+        # If CE paths are not provided, use RGB paths for both inputs
+        fpaths = rgb_paths
+    else:
+        # If CE paths are provided, pair them with RGB paths
+        if len(rgb_paths) != len(ce_paths):
+            raise ValueError("rgb_paths and ce_paths must have the same length")
+        fpaths = list(zip(rgb_paths, ce_paths))
+
+    dataloader = ensemble_fovea._make_inference_dataloader(
+        fpaths,
+        ids=ids,
+        num_workers=_inference_num_workers(device),
+        preprocess=False,
+        batch_size=8,
+    )
+    logger.info("Fovea dataloader ready with %d images", len(fpaths))
+
+    output_ids, outputs = [], []
+    with torch.no_grad():
+        for batch in tqdm(dataloader):
+            if len(batch) == 0:
+                continue
+
+            im = batch["image"].to(device)
+
+            # FOVEA DETECTION
+            with _autocast_context(device):
+                heatmap = ensemble_fovea.forward(im)
+            keypoints = extract_keypoints_from_heatmaps(heatmap)
+
+            kp_fovea = torch.mean(keypoints, dim=1)  # average over models
+
+            items = {
+                "id": batch["id"],
+                "keypoints": kp_fovea,
+                "metadata": batch["metadata"],
+            }
+            items = decollate_batch(items)
+
+            items = [dataloader.dataset.transform.undo_item(item) for item in items]
+
+            for item in items:
+                output_ids.append(item["id"])
+                outputs.append(
+                    [
+                        *item["keypoints"][0].tolist(),
+                    ]
+                )
+    return pd.DataFrame(
+        outputs,
+        index=output_ids,
+        columns=["x_fovea", "y_fovea"],
+    )

vascx_models/utils.py

@@ -0,0 +1,196 @@
+import logging
+from pathlib import Path
+from typing import Dict, Optional, Tuple
+
+import numpy as np
+from PIL import Image, ImageDraw
+
+from .config import OverlayConfig
+
+logger = logging.getLogger(__name__)
+
+
+def create_fundus_overlay(
+    rgb_path: str,
+    av_path: Optional[str] = None,
+    disc_path: Optional[str] = None,
+    ring_2r_path: Optional[str] = None,
+    ring_3r_path: Optional[str] = None,
+    fovea_location: Optional[Tuple[int, int]] = None,
+    output_path: Optional[str] = None,
+    overlay_config: Optional[OverlayConfig] = None,
+) -> np.ndarray:
+    """
+    Create a visualization of a fundus image with overlaid segmentations and markers.
+
+    Args:
+        rgb_path: Path to the RGB fundus image
+        av_path: Optional path to artery-vein segmentation (1=artery, 2=vein, 3=intersection)
+        disc_path: Optional path to binary disc segmentation
+        ring_2r_path: Optional path to a binary 2r ring mask
+        ring_3r_path: Optional path to a binary 3r ring mask
+        fovea_location: Optional (x,y) tuple indicating the location of the fovea
+        output_path: Optional path to save the visualization image
+        overlay_config: Overlay display configuration including enabled layers and colors
+
+    Returns:
+        Numpy array containing the visualization image
+    """
+    overlay_config = overlay_config or OverlayConfig()
+
+    # Load RGB image
+    rgb_img = np.array(Image.open(rgb_path))
+
+    # Create output image starting with the RGB image
+    output_img = rgb_img.copy()
+
+    # Load and overlay AV segmentation if provided
+    if av_path and (overlay_config.layers.arteries or overlay_config.layers.veins):
+        av_mask = np.array(Image.open(av_path))
+
+        # Create masks for arteries (1), veins (2) and intersections (3)
+        artery_mask = av_mask == 1
+        vein_mask = av_mask == 2
+        intersection_mask = av_mask == 3
+
+        if overlay_config.layers.arteries:
+            artery_combined = np.logical_or(artery_mask, intersection_mask)
+            output_img[artery_combined, :] = overlay_config.colors.artery
+
+        if overlay_config.layers.veins:
+            vein_combined = np.logical_or(vein_mask, intersection_mask)
+            output_img[vein_combined, :] = overlay_config.colors.vein
+
+    # Load and overlay optic disc segmentation if provided
+    if disc_path and overlay_config.layers.disc:
+        disc_mask = np.array(Image.open(disc_path)) > 0
+        output_img[disc_mask, :] = overlay_config.colors.disc
+
+    if ring_2r_path and overlay_config.layers.ring_2r:
+        ring_2r_mask = np.array(Image.open(ring_2r_path)) > 0
+        output_img[ring_2r_mask, :] = overlay_config.colors.ring_2r
+
+    if ring_3r_path and overlay_config.layers.ring_3r:
+        ring_3r_mask = np.array(Image.open(ring_3r_path)) > 0
+        output_img[ring_3r_mask, :] = overlay_config.colors.ring_3r
+
+    # Convert to PIL image for drawing the fovea marker
+    pil_img = Image.fromarray(output_img)
+
+    # Add fovea marker if provided
+    if fovea_location and overlay_config.layers.fovea:
+        draw = ImageDraw.Draw(pil_img)
+        x, y = fovea_location
+        marker_size = (
+            min(pil_img.width, pil_img.height) // 50
+        )  # Scale marker with image
+
+        # Draw yellow X at fovea location
+        draw.line(
+            [(x - marker_size, y - marker_size), (x + marker_size, y + marker_size)],
+            fill=overlay_config.colors.fovea,
+            width=2,
+        )
+        draw.line(
+            [(x - marker_size, y + marker_size), (x + marker_size, y - marker_size)],
+            fill=overlay_config.colors.fovea,
+            width=2,
+        )
+
+    # Convert back to numpy array
+    output_img = np.array(pil_img)
+
+    # Save output if path provided
+    if output_path:
+        Image.fromarray(output_img).save(output_path)
+
+    return output_img
+
+
+def batch_create_overlays(
+    rgb_dir: Path,
+    output_dir: Path,
+    av_dir: Optional[Path] = None,
+    disc_dir: Optional[Path] = None,
+    ring_2r_dir: Optional[Path] = None,
+    ring_3r_dir: Optional[Path] = None,
+    fovea_data: Optional[Dict[str, Tuple[int, int]]] = None,
+    overlay_config: Optional[OverlayConfig] = None,
+) -> None:
+    """
+    Create visualization overlays for a batch of images.
+
+    Args:
+        rgb_dir: Directory containing RGB fundus images
+        output_dir: Directory to save visualization images
+        av_dir: Optional directory containing AV segmentations
+        disc_dir: Optional directory containing disc segmentations
+        ring_2r_dir: Optional directory containing 2r ring masks
+        ring_3r_dir: Optional directory containing 3r ring masks
+        fovea_data: Optional dictionary mapping image IDs to fovea coordinates
+        overlay_config: Overlay display configuration including enabled layers and colors
+
+    Returns:
+        List of paths to created visualization images
+    """
+    # Create output directory if it doesn't exist
+    output_dir.mkdir(exist_ok=True, parents=True)
+    overlay_config = overlay_config or OverlayConfig()
+
+    # Get all RGB images
+    rgb_files = list(rgb_dir.glob("*.png"))
+    if not rgb_files:
+        logger.warning("No RGB images found for overlays in %s", rgb_dir)
+        return []
+    logger.info("Creating overlays for %d images", len(rgb_files))
+
+    # Process each image
+    for rgb_file in rgb_files:
+        image_id = rgb_file.stem
+
+        # Check for corresponding AV segmentation
+        av_file = None
+        if av_dir:
+            av_file_path = av_dir / f"{image_id}.png"
+            if av_file_path.exists():
+                av_file = str(av_file_path)
+
+        # Check for corresponding disc segmentation
+        disc_file = None
+        if disc_dir:
+            disc_file_path = disc_dir / f"{image_id}.png"
+            if disc_file_path.exists():
+                disc_file = str(disc_file_path)
+
+        ring_2r_file = None
+        if ring_2r_dir:
+            ring_2r_file_path = ring_2r_dir / f"{image_id}.png"
+            if ring_2r_file_path.exists():
+                ring_2r_file = str(ring_2r_file_path)
+
+        ring_3r_file = None
+        if ring_3r_dir:
+            ring_3r_file_path = ring_3r_dir / f"{image_id}.png"
+            if ring_3r_file_path.exists():
+                ring_3r_file = str(ring_3r_file_path)
+
+        # Get fovea location if available
+        fovea_location = None
+        if fovea_data and image_id in fovea_data:
+            fovea_location = fovea_data[image_id]
+
+        # Create output path
+        output_file = output_dir / f"{image_id}.png"
+
+        # Create and save overlay
+        create_fundus_overlay(
+            rgb_path=str(rgb_file),
+            av_path=av_file,
+            disc_path=disc_file,
+            ring_2r_path=ring_2r_file,
+            ring_3r_path=ring_3r_file,
+            fovea_location=fovea_location,
+            output_path=str(output_file),
+            overlay_config=overlay_config,
+        )
+    logger.info("Finished overlay generation in %s", output_dir)

vessels/vessels_july24.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bdabae77502648acd1c176bff6d5e3c8295da60f18f2f84fe1bcd6181d2b2ca4
+size 352821632

vessels/vessels_july24_DRHAGIS.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:406cf89ed2c35713296096cdfdbc2d6e67c164e25dbd6542612f31e2bfa0c85e
+size 352848262