MSIAug

MSIAug is a Data Augmentation framework for Mass Spectrometry Imaging classification. Users need to provide their dataset in the imzML format and a corresponding segmentation mask in the npy format. MSIAug tests different spectral augmentation methods with varying magnitude to automatically determine the best augmentation approach. In addition, MSIAug features guided augmentations and model souping¹.

 

Installation

1. Clone the repository:

   git clone https://github.com/CeMOS-IS/MSIAug
   cd MSIAug

2. Setup conda environment:

   conda create --name msiaug python=3.11 --yes
   conda activate msiaug

3. Install torch and torchvision depending on your CUDA version (see instructions here):

   # e.g. CUDA 12.4
   pip install torch==2.6.0 torchvision==0.21.0 --index-url https://download.pytorch.org/whl/cu124

3. Install other dependencies:

   pip install -e .

Configuration

Copy the template config in /configs and edit it:

   cp configs/template.yaml configs/my_config.yaml

Required specifications in the YAML config:

folderpath: Location where the imzML, ibd and npy files are located. The npy file corresponds to the segmentation mask. All three files need to have the same name:

├── folderpath
    ├── training1.imzML
    ├── training1.ibd
    ├── training1.npy
    ├── training2.imzML
    ├── training2.ibd
    ├── training2.npy
    ├── test.imzML
    ├── test.ibd
    ├── test.npy
    └── ... 

train_val_test: Files that are used for training and testing. Validation files can be included but are not required. If no validation files are present, 20% of training spectra are randomly sampled for validation.

Run MSIAug

Start the experiment:

   python run.py --config configs/my_config.yaml

Start a dry run to check whether everything works properly:

   python run.py --config configs/my_config.yaml --dry_run true

Reuse trained models

To utilize a trained model for evaluation, use the /configs/evaluation_template.yaml file and adapt the weights.pt checkpoint of that trained model.

   python evaluate.py --config configs/my_evaluation_config.yaml

Augmentations

• Gridmask²
• Random Gridmask
• Cutout³
• CutMix⁴
• Jittering
• Scaling
• Max. Masking
• m/z shift
• Gaussian Blur

Guided Augmentations

• Guided Gridmask
• Guided Cutout
• Guided CutMix
• Guided Jittering
• Guided Scaling
• Guided Gaussian Blur

Usage of single modules

Augmentation Pipeline:

   from msiaug.augmentations.pipeline import AugmentationPipeline
   from msiaug.augmentations import GaussianNoise

   augmentation = AugmentationPipeline([
      GaussianNoise(std=0.05)
   ])

   augmentation.setup(mz_values=575)

   dataset = ...

   # Custom training loop #
   for epoch in range(100):
      for x, y in dataset:
         x, y = augmentation(x, y, augmentation_rate=0.3)
         ...

Trainer module:

See the example script use_trainer_module.py for usage of the trainer module.

Citing

If you use MSIAug in your research, please cite our work (paper under review). 🎓

@misc{WeigandMSIAug,
  title = {MSIAug: Systematic Data Augmentation for Mass Spectrometry Imaging},
  author = {Weigand, Philipp and Ebert, Nikolas and Abu Sammour, Denis and Mohammed, Shad A and Hopf, Carsten and Wasenm{\"u}ller, Oliver},
  publisher = {GitHub},
  howpublished = {\url{https://github.com/CeMOS-IS/MSIAug}}
}

References

¹ Wortsman, Mitchell, et al. "Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time." International conference on machine learning. PMLR (2022).

² Chen, Pengguang, et al. "Gridmask data augmentation." arXiv preprint arXiv:2001.04086 (2020).

³ DeVries, Terrance, and Graham W. Taylor. "Improved regularization of convolutional neural networks with cutout." arXiv preprint arXiv:1708.04552 (2017).

⁴ Yun, Sangdoo, et al. "Cutmix: Regularization strategy to train strong classifiers with localizable features." Proceedings of the IEEE/CVF international conference on computer vision. (2019).