strong-augment 0.0.1

StrongAugment implementation.

Exposing and fixing the fragility of neural networks in digital pathology [arXiv]

This repository contains implementations for StrongAugment and creating distribution-shifted datasets for shifted evaluation of neural networks intended for clinical practice. Simply put, without strong augmentation, your neural networks will be extremely fragile to even small distribution shifts. This is beautifully shown in the figure below, where StrongAugment is compared to RandAugment and TrivialAugment on histological samples where the haematoxylin and eosin stains have been computationally modified. All cells have been annotated with AUROC values.

Installation

pip3 install strong_augment

Training with strong augmentation.

To train your neural networks with strong augmentatiom simply include StrongAugment to your image transformation pipeline!

import torchvision.transforms as T
from strong_augment import StrongAugment

trnsf = T.Compose(
    T.RandomResizedCrop(224),
    T.RandomVerticalFlip(0.5),
    T.RandomHorizontalFlip(0.5),
    StrongAugment(min_ops=2, max_ops=4, p=0.4), # Just one line!
    T.ToTensor(),
    T.Normalize(mean=[0.5, 0.5, 0.5], std=[0.2, 0.2, 0.2])
    T.RandomErase(0.2)
)

We can also visualize the augmented images!

from strong_augment import augmentation_collage, EXAMPLE_IMAGE, StrongAugment

augmentation_collage(EXAMPLE_IMAGE, StrongAugment())

Creating shifted datasets.

Function shift_dataset can be used create the distribution-shifted datasets for shifted evaluation.

from functools import partial
import torchvision.transforms.functional as F
from strong_augment import shift_dataset

# Let's define the distribution shift function.
shift_fn = partial(F.adjust_gamma, gamma=0.2)

# Now we can shift the dataset!
shift_dataset(
    paths=paths_to_dataset_images,
    output_dir="/data/shifted_datasets/gamma_02",
    function=shift_fn,
    num_workers=20,
)

Processing images |##########| 100000/100000 [00:49<00:00]

Citation

If you use StrongAugment or shifted evaluation, please cite us!

@paper{strong_augment2022,
    title = {Exposing and addressing the fragility of neural networks in digital pathology},
    author = {Pohjonen, Joona and Stürenberg, Carolin and Föhr, Atte and Rannikko, Antti and Mirtti, Tuomas and Pitkänen, Esa},
    url = {being uploaded...},
    publisher = {arXiv},
    year = {2022},
}