bayesmedaug: Bayesian Optimization Library for Medical Image Segmentation.

bayesmedaug optimizes your data augmentation hyperparameters for medical image segmentation tasks by using Bayesian Optimization.

bayesmedaug is currently in beta release and still in development.

Optimizing Magnitudes

import torch
import bayesmedaug
from bayesmedaug import VanillaUNet, Trainer, Listed, BOMed
from bayesmedaug import Rotate, ZoomOut, Gamma, RandomCrop, Sharpen

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

auglist = [
    Rotate,
    ZoomOut,
    Gamma,
    RandomCrop,
    Sharpen
]

params = {
    'angle': (0.2,2.8),
    'zoom_amount': (0.5,0.9),
    'gamma': (0.5,1.5),
    'crop_height': (0.1, 0.4),
    'alpha_sharpen': (0.1, 1),
    'lightness': (0.1, 1.5),

}

auglist = Listed(augmentations = auglist)

trainer = Trainer(
    model = VanillaUNet,
    model_args = {"n_channels": 1, "n_classes": 2},
    optimizer = torch.optim.Adam,
    optimizer_args = {"lr": 0.0005},
    device = device,
    epochs = 1,
    train_dir = "~/Desktop/drive/train/",
    eval_dir = "~/Desktop/drive/test/",
    augmentations = auglist,
    batch_size = 1
)

Optimizing Probabilities and Magnitudes

import torch
import bayesmedaug
from bayesmedaug import VanillaUNet, Trainer, Listed, BOMed
from bayesmedaug import Rotate, ZoomOut, Gamma

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

auglist = [
    Rotate,
    ZoomOut,
    Gamma,
    RandomCrop,
    Sharpen
]

params = {
    'angle': (0.2,2.8),                                                         
    'zoom_amount': (0.5,0.9),                                                   
    'gamma': (0.5,1.5),                                                         
    'crop_height': (0.1, 0.4),                                                  
    'alpha_sharpen': (0.1, 1),                                                          
    'lightness': (0.1, 1.5),                                                    
    'p_'+Rotate.__name__: (0, 0.7),                                             
    'p_'+ZoomOut.__name__: (0, 0.7),                                            
    'p_'+Gamma.__name__: (0, 0.7),                                              
    'p_'+RandomCrop.__name__: (0, 0.7),                                         
    'p_'+Sharpen.__name__: (0, 0.7)
}

auglist = Listed(augmentations = auglist, optimize_p = True, randomized = False)

trainer = Trainer(
    model = VanillaUNet,
    model_args = {"n_channels": 1, "n_classes": 2},
    optimizer = torch.optim.Adam,
    optimizer_args = {"lr": 0.0005},
    device = device,
    epochs = 1,
    train_dir = "~/Desktop/drive/train/",
    eval_dir = "~/Desktop/drive/test/",
    augmentations = auglist,
    batch_size = 1
)

Optimizing Binary Probabilities and Magnitudes

import torch
import bayesmedaug
from bayesmedaug import VanillaUNet, Trainer, BinaryListed, BOMed
from bayesmedaug import Rotate, ZoomOut, Gamma

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
tmp_path = "./tmp"

auglist = [
    Rotate,
    ZoomOut,
    Gamma,
    RandomCrop,
    Sharpen
]

params = {
    'angle': (0.2,2.8),                                                         
    'zoom_amount': (0.5,0.9),                                                   
    'gamma': (0.5,1.5),                                                         
    'crop_height': (0.1, 0.4),                                                  
    'alpha_sharpen': (0.1, 1),                                                          
    'lightness': (0.1, 1.5),                                                    
    'p_'+Rotate.__name__: (0, 1),                                             
    'p_'+ZoomOut.__name__: (0, 1),                                            
    'p_'+Gamma.__name__: (0, 1),                                              
    'p_'+RandomCrop.__name__: (0, 1),                                         
    'p_'+Sharpen.__name__: (0, 1)
}

auglist = BinaryListed(
    augmentations = auglist,
    train_image_path = "~/Desktop/drive/train/images",
    train_mask_path = "~/Desktop/drive/train/labels",
    tmp_path = "./tmp",
)

trainer = Trainer(
    model = VanillaUNet,
    model_args = {"n_channels": 1, "n_classes": 2},
    optimizer = torch.optim.Adam,
    optimizer_args = {"lr": 0.0005},
    device = device,
    epochs = 1,
    train_dir = tmp_path,
    eval_dir = "~/Desktop/drive/test/",
    augmentations = auglist,
    batch_size = 1
)

Then:

optimizer = BOMed(
    f = trainer.train,
    pbounds = params,
    random_state = 1,
)

optimizer.maximize(
    init_points = 15,
    n_iter = 15,
)

Authors

• M. Şafak Bilici
• Onur Boyar
• Enes Sadi Uysal
• Alara Hergün