This repository contains a TensorFlow implementation of BayTTA (Bayesian based TTA), from the paper:

[Medical image classification with optimized test-time augmentation using Bayesian model averaging]

by Zeinab Sherkatghanad, Moloud Abdar, Mohammadreza Bakhtyari, Paweł Pławiak, Vladimir Makarenkov

BayTTA (Bayesian based TTA) is a method that optimizes test-time augmentation using Bayesian model averaging. This technique generates a model list associated with different variations of the input data created through TTA. BMA then combines model predictions weighted by their respective posterior probabilities. Such an approach allows one to consider model uncertainty and thus enhance the predictive performance of machine learning models.

In this repo, we implement BayTTA for image classification with several different architectures on three public medical image datasets comprising skin cancer, breast cancer, and chest X-ray images as well as two popular gene editing datasets, CRISPOR and GUIDE-seq.

python setup.py 

https://github.com/Z-Sherkat/BayTTA

See requirements.txt file for requirements that came from our setup. All experiments were run on Compute Canada cluster with NVIDIA Tesla P100 and NVIDIA v100 GPUs (the Cedar cluster).

python train_models.py --base_models 'VGG'\ --learning_rates 0.0005\  --Batch_size 256

python ./predict_BayTTA.py --aug 'flip'

|-- requirements.txt (The requirement for reproducing the environment)
|-- plots (The model and results images)
|-- setup.py (Make this project pip installation with 'pip install -e')
|-- src/ 
|   |-- __init__.py
|   |-- checkpoint/ (Checkpoints for trained models)
|   |-- constants/ (Script to specified constant values for training our models)
|   |-- dataset/ (Folder for description of datasets and the URL link for the datasets)
|   |   |-- dataset_description.txt
|   |-- models/
|    |   |-- predict_BayTTA.py (Used trained model for making the prediction based on BayTTA method) 
|    |   |-- BMA.py (Class definition for BMA)
|    |   |-- train_models.py (Train models)
|   |-- utils/ (Utility functions to process the data and metrics)

Comparison of the baseline models accuracy (%) ± STD performance with TTA and BayTTA versions on the skin cancer dataset.

Comparison of the baseline models accuracy (%) ± STD performance with TTA and BayTTA versions on the breast cancer dataset.

Comparison of the baseline models accuracy (%) and STD performance with TTA and BayTTA versions on the chest X-ray dataset.

Accuracy results of our proposed method on the CRISPOR and GUIDE-seq datasets.