MedAugment: Automatic Medical Augmenter

Official Pytorch Implementation for Paper “MedAugment: Universal Automatic Data Augmentation Plug-in for Medical Image Analysis”

Update

We have released the online version of MedAugment (OAA) here. The OAA allows online data augmentation as training progresses

Preparation

To use MedAugment as a plug-in for your project, you should have a "baseline" folder for your custom dataset at

./datasets/classification/your_dataset_name/baseline

The organization of the "baseline" folder should follow

├── classification
    ├── your_dataset_name
        ├── baseline
            ├── training
            |   ├── class_1
            |   |   ├── img_1.jpg  # .png
            |   │   ├── img_2.jpg
            |   │   ├── ...
            |   ├── class_2
            |   |   ├── img_a.jpg
            |   │   ├── img_b.jpg
            |   │   ├── ...
            |   ├── ...
            └── validation
            └── test

├── segmentation
    ├── your_dataset_name
        ├── baseline
            ├── training
            |   ├── img_1.jpg
            │   ├── img_2.jpg
            │   ├── ...
            ├── training_mask
            |   ├── img_1_mask.jpg  # suffix
            │   ├── img_2_mask.jpg
            │   ├── ...
            └── validation
            └── validation_mask
            └── test
            └── test_mask

You can then move the plug-in to your project utils and run

python ./utils/medaugment.py --dataset=your_dataset_name

This will produce an augmented dataset named "medaugment" at the same level as the "baseline" folder

To run MedAugment and other augmentation methods and train the model in the paper

python ./utils/generation.py
python classification.py

You should have a "recording" folder at the root with two subfolders named "classification" and "segmentation"

Citation

If you find MedAugment useful for your research, please cite our paper as

@misc{liu2023medaugment,
      title={MedAugment: Universal Automatic Data Augmentation Plug-in for Medical Image Analysis}, 
      author={Zhaoshan Liu and Qiujie Lv and Yifan Li and Ziduo Yang and Lei Shen},
      year={2023},
      eprint={2306.17466},
      archivePrefix={arXiv},
      primaryClass={eess.IV}
}