Interpretable-by-design deep survival analysis for disease progression modeling

Accepted at MICCAI 2024

This repository contains the code for the paper "Interpretable-by-design Deep Survival Analysis for Disease Progression Modeling". The code is based on PyTorch and is used to train and evaluate a deep learning model for survival analysis on fundus images. The model offers built-in interpretability by yielding an evidence map of local disease risk which is subsequently simply aggregated to a final risk prediction. The model is based on a combination of a Sparse BagNet and a Cox proportional hazards model. As a sample use case, it is trained to predict the risk of conversion to age-related macular degeneration at different time points and its performance is compared to that of black-box SOTA baseline models.

Figure 1: Model architecture.

Pre-requisites

• Obtain the AREDS data (see Data)
• Install the conda environment using conda env create -f requirements.yml (or use requirements_without_R.yml which installs much faster but lacks AUPRC metric support) and activate it with conda activate amd-surv.
• Log in to wandb with the command wandb login. Create an account if you don't have one.

Load and evaluate the pre-trained interpretable-by-design survival model on AREDS data

• Retrieve the model weights 7ufjvvnz_best.pth from here and place it in the checkpoints directory.
• Follow the instructions in evaluate_survival.ipynb to load the model and evaluate it.
• To evaluate the pre-trained baseline models, retrieve the other weights files and repeat the procedure.

Alternatively, train the model yourself

The model can be trained and evaluated using the following command:
python train_and_evaluate_survival.py --config configs/sparsebagnet_cox.yml

Similarly, the baseline models can be trained and evaluated using the files in configs/babenko/ and configs/yan/. The model configs serve to train one model for each inquired point in time (years 1 to 5).

Modify the configuration

The training is configured using a yaml file. The following parameters can be set:

The model takes below 12h to train and evaluate on a single NVIDIA GeForce RTX 2080 Ti GPU.

Data

Upon request, the AREDS data can be accessed from the dbGaP repository. The tabular data needs to be parsed to one metadata file with the rows representing individual macula-centered (F2) images. <!--Throughout the study, two identifiers for participants have been used. Match participants to ID2 identifiers using the mapping file that comes with the AREDS data.
-->

The AREDS data directory as defined in configs/dirs.yml should contain:
- a folder as set the model config (image_dir), the directory of image data. The images should be organized such that the metadata CSV image paths map to the image files, relative to the image_dir.

The <project>/data directory should contain the following files:
- metadata_surv.csv: Survival metadata table that maps screening data (event, duration) to image records. An example entry is provided in data/metadata_surv example.csv with important columns as follows:

• metadatasurvstereo.csv: Image metadata table with the same structure, but of records where both views (LS and RS) exist (for Babenko et al. baseline model).

Both datasets can be created from a parsed AREDS metadata file using save_survival_metadata.ipynb.

Illustrative figures

Figure 2: Example risk evidence maps (a) and the resulting survival curve predictions (b).

Figure 3: Examples of most predictive image patches which could be provided to clinicians (purple: higher risk, green: lower risk).

Credits

This work is mainly based on

Deep CoxPH modeling - Katzman, J. et al. "DeepSurv: personalized treatment recommender system using a Cox proportional hazards deep neural network." BMC Med Res Methodol 18, 24 (2018).

The Sparse BagNet for classification - Djoumessi, K. et al. "Sparse activations for interpretable disease grading." Medical Imaging with Deep Learning (2023).

Baselines: SOTA end-to-end AMD progression models - Babenko, B. et al. "Predicting progression of age-related macular degeneration from fundus images using deep learning." arXiv preprint (2019). - Yan, Q. et al. "Deep-learning-based prediction of late age-related macular degeneration progression." Nat Mach Intell 2, 141–150 (2020). - Note: We used the authors' model variant that does not rely on gene data and is trained end-to-end on fundus images.

This work includes code adaptations from

• scikit-survival: We adapted the Breslow estimator to store the baseline hazard function and survival function and subsequently init the estimator from the saved data.
• auton-survival: We adapted the CoxPH loss implementation from the authors' code.
• BagNet (Brendel and Bethge, 2019) and Sparse BagNet (Djoumessi et al., 2023): We adapted the BagNet implementations to survival modelling.
• Huang et al., 2023: We adapted the data augmentation from the authors' code.

Cite

bibtex @InProceedings{gervelmeyer2024interpretable, author = { Gervelmeyer, Julius and Müller, Sarah and Djoumessi, Kerol and Merle, David and Clark, Simon J. and Koch, Lisa and Berens, Philipp}, title = { { Interpretable-by-design Deep Survival Analysis for Disease Progression Modeling } }, booktitle = {Proceedings of Medical Image Computing and Computer Assisted Intervention -- MICCAI 2024}, year = {2024}, publisher = {Springer Nature Switzerland}, volume = {LNCS 15010}, month = {October}, pages = {502--512}, URL = {https://papers.miccai.org/miccai-2024/421-Paper1325.html} }

Back to top