Multimodal Contrastive Learning with Tabular and Imaging Data

Please cite our CVPR paper, Best of Both Worlds: Multimodal Contrastive Learning with Tabular and Imaging Data, if this code was helpful.

@InProceedings{Hager_2023_CVPR,
    author    = {Hager, Paul and Menten, Martin J. and Rueckert, Daniel},
    title     = {Best of Both Worlds: Multimodal Contrastive Learning With Tabular and Imaging Data},
    booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
    month     = {June},
    year      = {2023},
    pages     = {23924-23935}
}

If you want an overview of the paper checkout:

• this podcast episode
• the video abstract

Instructions

Install environment using conda env create --file environment.yaml.

To run, execute python run.py.

Arguments - Command Line

If pretraining, pass pretrain=True and datatype={imaging|multimodal|tabular} for the desired pretraining type. multimodal uses our strategy from the paper, tabular uses SCARF, and imaging can be specified with the loss argument. Default is SimCLR, other options are byol, simsiam, and barlowtwins.

If you do not pass pretrain=True, the model will train fully supervised with the data modality specified in datatype, either tabular or imaging.

You can evaluate a model by passing the path to the final pretraining checkpoint with the argument checkpoint={PATH_TO_CKPT}. After pretraining, a model will be evaluated with the default settings (frozen eval, lr=1e-3).

Arguments - Hydra

All argument defaults can be set in hydra yaml files found in the configs folder.

Most arguments are set to those in the paper and work well out of the box. Default model is ResNet50.

Code is integrated with weights and biases, so set wandb_project and wandb_entity in config.yaml.

Path to folder containing data is set through the data_base argument and then joined with filenames set in the dataset yamls. Best strategy is to take dvm_all_server.yaml as a template and fill in the appropriate filenames.

• For the images, provide a .pt with a list of your images or a list of the paths to your images.
  • If providing a list of paths, set live_loading=True.
• delete_segmentation deletes the first channel of a three channel image (historical reasons) and should typically be left to false.
• If weights is set, during finetuning a weighted sampled will be used instead of assuming the evaluation train data has been properly balanced
• eval_metric supports acc for accuracy (top-1) and auc (for unbalanced data)
• If doing multimodal pretraining or tabular pretraining (SCARF), the tabular data should be provided as NOT one-hot encoded so the sampling from the empirical marginal distribution works correctly. You must provide a file field_lengths_tabular which is an array that in the order of your tabular columns specifies how many options there are for that field. Continuous fields should thus be set to 1 (i.e. no one-hot encoding necessary), while categorical fields should specify how many columns should be created for the one_hot encoding

Data

The UKBB data is semi-private. You can apply for access here.

The DVM cars dataset is open-access and can be found here.

Processing steps for the DVM dataset can be found here.

The exact data splits used in the paper are saved in the data folder.

Additional information

MMCL

To get the data simply run mmcl_dataprocessing.slurm otherwise:

OAI 1-Year Dataset Pipeline

This repo was extended to run multimodal contrastive learning on the OAI 1-year X-ray dataset.

1. Patient List

The starting point is a text file containing the OAI patient IDs to use: data/patients_to_use.txt

Example: 9000099 9000798 9001400

2. Create Image Inventory

Script: python datasets/make_oai_1year_image_inventory.py

Inputs: data/patients_to_use.txt /vol/miltank/projects/practical_sose26/Osteoarthritis/Xray_1year/download/

Output: data/oai_1year_image_inventory.csv

This script finds the local 1-year X-ray file for each patient ID.

3. Create Clinical Manifest

Script:

python datasets/make_oai_1year_manifest.py

Inputs: data/oai_1year_image_inventory.csv AllClinical01.txt

Output:

data/oai_1year_manifest.csv

This script merges image paths with OAI clinical/tabular variables, drops incomplete rows, and creates the temporary binary label target_high_pain.

4. Create MMCL Files

Script: python datasets/prepare_oai_1year_mmcl.py

Input: data/oai_1year_manifest.csv

Output folder:

data/oai_1year_mmcl/

Created files:

oai_1year_train_images.pt oai_1year_val_images.pt oai_1year_test_images.pt

oai_1year_train_tabular.csv oai_1year_val_tabular.csv oai_1year_test_tabular.csv

oai_1year_train_labels.pt oai_1year_val_labels.pt oai_1year_test_labels.pt

oai_1year_tabular_lengths.pt

These are the files expected by the existing MMCL dataset loader.

5. Dataset Config

The Hydra dataset config is:

configs/dataset/oai_1year.yaml

It points the trainer to the OAI image tensors, tabular CSVs, labels, and tabular field lengths.

6. Run MMCL Pretraining

Example command:

python run.py
dataset=oai_1year
models=resnet18
data_base=/vol/miltank/users/wden/ContLearn/MMCL-Tabular-Imaging/data/oai_1year_mmcl
datatype=multimodal
pretrain=True
evaluate=False
test=False
use_wandb=False
offline=True
max_epochs=50
batch_size=64
num_workers=4
one_hot=True
img_size=192

What The Model Learns

Currently only aligns image and tabular representations:

X-ray image embedding <-> tabular row embedding

During training, the model receives matching image-tabular pairs. In each batch, it tries to identify which tabular row belongs to which X-ray. So the contrastive task is:

Given an X-ray, find the matching tabular record among the other records in the batch.

The label target_high_pain is included for later supervised evaluation or fine-tuning, but basic loss=clip MMCL pretraining mainly optimizes image-tabular alignment.

Data used currently:

womac_pain womac_total koos_pain nrs_pain bmi age

Predictor

The predictor is separate from MMCL for now. It trains a small MLP on clinical/tabular OAI data and predicts whether pain gets worse in later visits.

The current target is:

max(nrs_pain_y2, nrs_pain_y3) - nrs_pain_y1 > 1

So the model predicts a binary label:

0 = pain does not get worse by more than 1 point 1 = pain gets worse by more than 1 point

The predictor does not directly predict WOMAC pain. WOMAC pain, WOMAC total, KOOS pain, NRS pain, BMI, and age are input features from year 1.

1. Create Predictor Manifest

Script:

python datasets/make_oai_pain_progression_manifest.py

Input:

data/oai_1year_manifest.csv AllClinical01.txt AllClinical02.txt AllClinical03.txt

Output:

data/oai_pain_progression_manifest.csv

This script uses the same OAI cohort as the MMCL 1-year data, then adds year-2 and year-3 clinical values. It creates the binary target target_pain_worse from future NRS pain.

The current manifest is strict: patients must have the required year-1, year-2, and year-3 clinical fields. Needs to be updated as a lot is missing

Slurm script:

sbatch scripts/mlp_create_dataset.slurm

2. Train Predictor MLP

Script:

python predictors/train_pain_progression.py

Input:

data/oai_pain_progression_manifest.csv

The train script splits the manifest internally into train, validation, and test sets using a stratified split. The current split is approximately:

70% train 15% validation 15% test

The model is a small MLP with hidden layers defined by --hidden-dims, for example:

--hidden-dims 64 32

This means two hidden layers: first 64 neurons, then 32 neurons.

Each hidden block is:

Linear -> BatchNorm1d -> ReLU -> Dropout

By default, each training run saves under:

runs/predictor/<current_run>/

Important outputs:

best_model.pt metrics.csv pain_progression_splits.csv summary.json

Slurm script:

sbatch scripts/mlp_train.slurm

3. Single Tabular Prediction

A trained predictor checkpoint can be used on one patient row or a small CSV of tabular values.

Script:

python predictors/predict_pain_progression.py

Example:

python predictors/predict_pain_progression.py
--checkpoint runs/predictor/<run_name>/best_model.pt
--input-csv data/oai_pain_progression_manifest.csv
--output-csv runs/predictor/single_patient_prediction.csv

Slurm script:

sbatch scripts/mlp_run_single_patient.slurm

The output contains the predicted probability and binary prediction for target_pain_worse.

4. Combined MMCL + Predictor Pipeline

The combined pipeline takes one X-ray image and runs:

X-ray image -> MMCL retrieval -> retrieved tabular row -> predictor MLP -> pain worsening prediction

Script:

python predict.py

The MMCL part retrieves the most similar tabular representation for the input image. The predictor then uses that retrieved tabular data as input.

Example:

python predict.py
--image-path /path/to/single/xray
--mmcl-checkpoint runs/multimodal/2wclumnv/checkpoints/checkpoint_last_epoch_49.ckpt
--predictor-checkpoint runs/predictor/pain_progression_mlp_162331/best_model.pt
--top-k 10
--aggregation top1
--output-csv runs/pipeline/prediction.csv
--matches-output-csv runs/pipeline/retrieval_matches.csv

aggregation is not included yet! only top1 works for now Slurm script:

sbatch scripts/predict_from_image.slurm

Outputs:

runs/pipeline/prediction_<job_id>.csv runs/pipeline/retrieval_matches_<job_id>.csv

prediction_<job_id>.csv contains the final predictor result.

retrieval_matches_<job_id>.csv contains the top MMCL tabular matches for the input image.

5. Aggregation Options

The combined pipeline supports two ways to convert MMCL retrieval results into predictor input:

top1: Use the single most similar retrieved tabular row.

weighted: Use a weighted average of the top-k retrieved tabular rows based on MMCL similarity.

Current default:

--aggregation top1

6. Current Interpretation

The full pipeline does not yet diagnose pain directly from the image.

It does this:

1. Find a tabular patient representation that is close to the input image in MMCL embedding space.
2. Use that retrieved tabular data to predict whether future NRS pain becomes worse by more than 1 point.