/bhb10k-dl-benchmark

A Reproducible Benchmark for CNN Models on the BHB-10K Dataset

Primary LanguagePython

Benchmarking CNN Models on Brain MRI BHB-10K dataset

In this repository you can find all the CNN models implemented in Pytorch used as a baseline on the BHB-10K dataset for age and sex prediction. You can also find the baselines for schizophrenia prediction on a clinical dataset. This repository contains also the data augmentation transformations for 3D MRI implemented as well as a convenient Pytorch Dataloader.

All preprocessed data are available through the links below.

Finally, the baselines using simple linear (resp. logistic) regression for age (resp. sex and diagnosis) prediction are also provided.

Data

Alt text

Pre-processings

Quasi-Raw. It consists essentially in brain extraction with BET2 and a linear registration to the MNI template with FLIRT for a final isotropic spatial resolution of 1.5mm.

Voxel-Based Morphometry (VBM). This is an extensive pre-processing performed with CAT12. The brain is segmented into 3 tissues: Gray Matter (GM), White Matter (WM) and Cerebrospinal Fluid (CSF) and it is then re-aligned non-linearly to the MNI template with DARTEL resampled at 1.5mm isotropic. We used only the T1-weighted modulated GM modality.

Alt text

BHB-10K Dataset

We aggregated 13 brain MRI datasets of healthy controls (HC) both pre-processed with VBM and Quasi-Raw. The link to the BHB-10K dataset can be found below.

Available Source # Subjects # Sessions Age Sex (%F) # Sites
❌ HCP 1113 1113 29 ± 4 45 1
❌ IXI 559 559 48 ± 16 55 3
❌ CoRR 1371 2897 26 ± 16 50 19
❌ NPC 65 65 26 ± 4 55 1
❌ NAR 303 323 22 ± 5 58 1
❌ RBP 40 40 23 ± 5 52 1
❌ OASIS 3 597 1262 67 ± 9 62 3
❌ GSP 1570 1639 21 ± 3 58 1
❌ ICBM 622 977 30 ± 12 45 3
❌ ABIDE 1 567 567 17 ± 8 17 20
❌ ABIDE 2 559 580 15 ± 9 30 17
❌ Localizer 82 82 25 ± 7 56 2
❌ MPI-Leipzig 316 316 37 ± 19 40 2
Total 7764 10420 32 ± 19 50 74
Currently 0 0 0 0 0

Clinical Datasets

We focused our analysis only on HC and Schizophrenia for the clinical datasets.

Available Source # Subjects Diagnosis Age Sex (%F) # Sites
❌ BSNIP 394 Schizophrenia
Control		 34 ± 12
38 ± 13		 44
58		 5
❌ SCHIZCONNECT-VIP 605 Schizophrenia
Control		 34 ± 12
32 ± 12		 27
47		 4

Links to the Pre-Processed Datasets

Dataset Pre-Processing # Images Target Link
BHB-10K Quasi-Raw ? Age + Sex ?
BHB-10K VBM ? Age + Sex ?
SCHIZCONNECT-VIP Quasi-Raw 605 SCZ vs CTL Pending
SCHIZCONNECT-VIP VBM 605 SCZ vs CTL Pending

CNN Models

The main CNN models currently available in this repository are:

  • tiny-VGG initially built for age prediction, it remains a competitive network
  • ResNet
  • VGG
  • tiny-DenseNet
  • DenseNet
  • ResNeXt
  • SFCN state-of-the-art for age and sex prediction on UKBioBank

Training/Test Split

We aim at giving results in the real clinical setting where the model is evaluated on new data arriving from different hospitals (different acquisition protocols).

Task Training Set Test Sets
Age BHB-10K BSNIP (only HC)
Sex BHB-10K BSNIP (only HC)
SCZ vs HC SCHIZCONNECT-VIP BSNIP

Main Results

Learning curves for BSNIP

1) VBM

Alt text

2) Quasi-Raw

Alt text Note: linear models give results no better than chance with quasi-raw data.

Best Final Results

We reported the results for 3 (resp. 5) runs on 3 (resp. 5) Stratified Shuffle Splits at N=10K (resp. N=500). We stratified according to the label to predict. The hyperparameters for the linear models are tuned using grid search.

Linear Model Baseline

Test on BSNIP
Task Model # Training Samples Pre-Processing AUC(%) MAE
Age Ridge 10K VBM N/A 4.65±0.02
Sex Logistic 10K VBM 97.05±0.03 N/A
SCZ vs HC Logistic 500 VBM 78.71±0.76 N/A

CNN Models

The models are trained for 300 epochs at N=10K and for 100 epochs at N=500. We use Deep Ensemble with T=5 CNN to have accurate and better calibrated models.

Task Model # Training Samples Deep Ensemble (x5) Pre-Processing Weights
Age DenseNet121 10K ❌ VBM download
Age ResNet34 10K ❌ Quasi-Raw download
Sex DenseNet121 10K ❌️ VBM download
Sex DenseNet121 10K ❌️ Quasi-Raw download
SCZ vs HC tiny-DenseNet 500 ✔️ VBM download
SCZ vs HC DenseNet121 500 ✔️ Quasi-Raw download
Test on BSNIP
Task Model Pre-Processing AUC MAE
Age DenseNet121 VBM N/A 4.03±0.13
Age ResNet34 Quasi-Raw N/A 4.84±0.26
Sex DenseNet121 VBM 97.69±0.21 N/A
Sex DenseNet121 Quasi-Raw 96.58±0.21 N/A
SCZ vs HC tiny-DenseNet ️ VBM 80.92±0.47 N/A
SCZ vs HC DenseNet121 Quasi-Raw 71.98±1.49 N/A