first commit github branch

.gitattributes

@@ -0,0 +1,37 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz 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,6 @@
+*.pyc
+__pycache__
+*.egg-info
+*.zip
+/samples/fundus/*
+!/samples/fundus/original

README.md

@@ -0,0 +1,143 @@
+---
+license: agpl-3.0
+pipeline_tag: image-segmentation
+tags:
+- medical
+- biology
+---
+
+# VascX models
+
+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).
+
+The model weights are in [huggingface](https://huggingface.co/Eyened/vascx).
+
+<img src="imgs/CHASEDB1_12R_rgb.png" width="240" height="240" style="display:inline"><img src="imgs/CHASEDB1_12R.png" width="240" height="240" style="display:inline">
+
+<img src="imgs/DRIVE_22_rgb.png" width="240" height="240" style="display:inline"><img src="imgs/DRIVE_22.png" width="240" height="240" style="display:inline">
+
+<img src="imgs/HRF_04_g_rgb.png" width="240" height="240" style="display:inline"><img src="imgs/HRF_04_g.png" width="240" height="240" style="display:inline">
+
+## 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 the [rtnls_inference package](https://github.com/Eyened/retinalysis-inference).
+
+
+## `vascx run` Command
+
+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]
+```
+
+### 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` | `--overlay` | Create visualization overlays combining all results |
+| `--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
+├── overlays/             # Visualization images
+├── bounds.csv            # Image boundary information
+├── 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
+
+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
+     - Fovea marked with yellow X
+
+### 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
+
+### 
+
+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.

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('Eyened/vascx: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('Eyened/vascx: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('Eyened/vascx: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
+}

setup.py

@@ -0,0 +1,33 @@
+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 == 1.*",
+        "pandas == 2.*",
+        "tqdm == 4.*",
+        "Pillow == 9.*",
+        "click==8.*",
+    ],
+    python_requires=">=3.10, <3.11",
+)

vascx_models/cli.py

@@ -0,0 +1,198 @@
+from pathlib import Path
+
+import click
+import pandas as pd
+import torch
+
+from rtnls_fundusprep.cli import _run_preprocessing
+
+from .inference import (
+    run_fovea_detection,
+    run_quality_estimation,
+    run_segmentation_disc,
+    run_segmentation_vessels_and_av,
+)
+from .utils import batch_create_overlays
+
+
+@click.group(name="vascx")
+def cli():
+    pass
+
+
+@cli.command()
+@click.argument("data_path", type=click.Path(exists=True))
+@click.argument("output_path", type=click.Path())
+@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=True, help="Create visualization overlays"
+)
+@click.option("--n_jobs", type=int, default=4, help="Number of preprocessing workers")
+def run(
+    data_path, output_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)
+
+    # 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"
+    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)
+    if overlay:
+        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
+
+    # 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:
+        click.echo(f"Reading file paths from CSV: {data_path}")
+        try:
+            df = pd.read_csv(data_path)
+            if "path" not in df.columns:
+                click.echo("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()
+                click.echo("Using IDs from CSV 'id' column")
+
+        except Exception as e:
+            click.echo(f"Error reading CSV file: {e}")
+            return
+    else:
+        click.echo(f"Finding files in directory: {data_path}")
+        files = list(data_path.glob("*"))
+        ids = [f.stem for f in files]
+
+    if not files:
+        click.echo("No files found to process")
+        return
+
+    click.echo(f"Found {len(files)} files to process")
+
+    # Step 1: Preprocess images if requested
+    if preprocess:
+        click.echo("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]
+
+    # Set up GPU device
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    click.echo(f"Using device: {device}")
+
+    # Step 2: Run quality estimation if requested
+    if quality:
+        click.echo("Running quality estimation...")
+        df_quality = run_quality_estimation(
+            fpaths=preprocessed_files, ids=ids, device=device
+        )
+        df_quality.to_csv(quality_path)
+        click.echo(f"Quality results saved to {quality_path}")
+
+    # Step 3: Run vessels and AV segmentation if requested
+    if vessels:
+        click.echo("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,
+        )
+        click.echo(f"Vessel segmentation saved to {vessels_path}")
+        click.echo(f"AV segmentation saved to {av_path}")
+
+    # Step 4: Run optic disc segmentation if requested
+    if disc:
+        click.echo("Running optic disc segmentation...")
+        run_segmentation_disc(
+            rgb_paths=preprocessed_files, ids=ids, output_path=disc_path, device=device
+        )
+        click.echo(f"Disc segmentation saved to {disc_path}")
+
+    # Step 5: Run fovea detection if requested
+    df_fovea = None
+    if fovea:
+        click.echo("Running fovea detection...")
+        df_fovea = run_fovea_detection(
+            rgb_paths=preprocessed_files, ids=ids, device=device
+        )
+        df_fovea.to_csv(fovea_path)
+        click.echo(f"Fovea detection results saved to {fovea_path}")
+
+    # Step 6: Create overlays if requested
+    if overlay:
+        click.echo("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,
+            fovea_data=fovea_data,
+        )
+
+        click.echo(f"Visualization overlays saved to {overlay_path}")
+
+    click.echo(f"All requested processing complete. Results saved to {output_path}")

vascx_models/inference.py

@@ -0,0 +1,269 @@
+import os
+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
+
+
+def run_quality_estimation(fpaths, ids, device: torch.device):
+    ensemble_quality = ClassificationEnsemble.from_release("quality.pt").to(device)
+    dataloader = ensemble_quality._make_inference_dataloader(
+        fpaths,
+        ids=ids,
+        num_workers=8,
+        preprocess=False,
+        batch_size=16,
+    )
+
+    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 = torch.device(
+        "cuda:0" if torch.cuda.is_available() else "cpu"
+    ),
+) -> 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
+    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=8,
+        preprocess=False,
+        batch_size=8,
+    )
+
+    # Run inference
+    with torch.no_grad():
+        for batch in tqdm(dataloader):
+            # AV segmentation
+            if av_path is not None:
+                with torch.autocast(device_type=device.type):
+                    proba = ensemble_av.forward(batch["image"].to(device))
+                proba = torch.mean(proba, dim=0)  # 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 torch.autocast(device_type=device.type):
+                    proba = ensemble_vessels.forward(batch["image"].to(device))
+                proba = torch.mean(proba, dim=0)  # 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 = torch.device(
+        "cuda:0" if torch.cuda.is_available() else "cpu"
+    ),
+) -> None:
+    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=8,
+        preprocess=False,
+        batch_size=8,
+    )
+
+    with torch.no_grad():
+        for batch in tqdm(dataloader):
+            # AV
+            with torch.autocast(device_type=device.type):
+                proba = ensemble_disc.forward(batch["image"].to(device))
+            proba = torch.mean(proba, dim=0)  # 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 = torch.device(
+        "cuda:0" if torch.cuda.is_available() else "cpu"
+    ),
+) -> None:
+    # def run_fovea_detection(fpaths, ids, device: torch.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=8,
+        preprocess=False,
+        batch_size=8,
+    )
+
+    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 torch.autocast(device_type=device.type):
+                heatmap = ensemble_fovea.forward(im)
+            keypoints = extract_keypoints_from_heatmaps(heatmap)
+
+            kp_fovea = torch.mean(keypoints, dim=0)  # 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,160 @@
+from pathlib import Path
+from typing import Dict, Optional, Tuple
+
+import numpy as np
+from PIL import Image, ImageDraw
+
+
+def create_fundus_overlay(
+    rgb_path: str,
+    av_path: Optional[str] = None,
+    disc_path: Optional[str] = None,
+    fovea_location: Optional[Tuple[int, int]] = None,
+    output_path: Optional[str] = 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
+        fovea_location: Optional (x,y) tuple indicating the location of the fovea
+        output_path: Optional path to save the visualization image
+
+    Returns:
+        Numpy array containing the visualization image
+    """
+    print(rgb_path, av_path, disc_path, fovea_location, output_path)
+    # 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:
+        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
+
+        # Combine artery and intersection for visualization
+        artery_combined = np.logical_or(artery_mask, intersection_mask)
+        vein_combined = np.logical_or(vein_mask, intersection_mask)
+
+        # Apply colors: red for arteries, blue for veins
+        # Red channel - increase for arteries
+        output_img[artery_combined, 0] = 255
+        output_img[artery_combined, 1] = 0
+        output_img[artery_combined, 2] = 0
+
+        # Blue channel - increase for veins
+        output_img[vein_combined, 0] = 0
+        output_img[vein_combined, 1] = 0
+        output_img[vein_combined, 2] = 255
+
+    # Load and overlay optic disc segmentation if provided
+    if disc_path:
+        disc_mask = np.array(Image.open(disc_path)) > 0
+
+        # Apply white color for disc
+        output_img[disc_mask, :] = [255, 255, 255]  # White
+
+    # Convert to PIL image for drawing the fovea marker
+    pil_img = Image.fromarray(output_img)
+
+    # Add fovea marker if provided
+    if fovea_location:
+        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=(255, 255, 0),
+            width=2,
+        )
+        draw.line(
+            [(x - marker_size, y + marker_size), (x + marker_size, y - marker_size)],
+            fill=(255, 255, 0),
+            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,
+    fovea_data: Optional[Dict[str, Tuple[int, int]]] = 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
+        fovea_data: Optional dictionary mapping image IDs to fovea coordinates
+
+    Returns:
+        List of paths to created visualization images
+    """
+    # Create output directory if it doesn't exist
+    output_dir.mkdir(exist_ok=True, parents=True)
+
+    # Get all RGB images
+    rgb_files = list(rgb_dir.glob("*.png"))
+    if not rgb_files:
+        return []
+
+    # 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)
+
+        # 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,
+            fovea_location=fovea_location,
+            output_path=str(output_file),
+        )