mirabilis-ringboundary-seg-minimal.ipynb

{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.7.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"\nThis is the minimal code for training a VGG-11 U-Net to segment ring boundaries on images of macroscopic colonies formed by motile microbes (e.g. the swarming bacterium, *Proteus mirabilis*).\n\n\n**References/Acknowledgments:**\nOur work in bacterial ring boundary segmentation greatly benefited from [Segmentation Models: Python library with Neural Networks for Image Segmentation based on PyTorch](https://github.com/qubvel/segmentation_models.pytorch) (SMP for short). The model (its architecture and pretrained encoder) and many utility functions used below came from SMP. Overall, this script closely follows [SMP's car segmentation example](https://github.com/qubvel/segmentation_models.pytorch/tree/master/examples), particularly the functions needed for data loading, augmentation, and model training, with slight modifications. ","metadata":{}},{"cell_type":"markdown","source":"# Imports","metadata":{}},{"cell_type":"code","source":"# Install PyTorch segmentation models \n!pip install git+https://github.com/qubvel/segmentation_models.pytorch","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport cv2\nimport csv\nimport copy\nimport time\nfrom tqdm import tqdm\nimport os\nimport matplotlib.pyplot as plt\n\nimport torch\nimport torchvision\nfrom torchvision import transforms\nfrom torch import nn\nfrom torch.nn import functional as F\nfrom torchvision import models\nfrom torch.utils.data import Dataset, DataLoader\n\nimport segmentation_models_pytorch as smp\nfrom segmentation_models_pytorch import losses\nfrom segmentation_models_pytorch.encoders import get_preprocessing_fn\nimport albumentations as albu","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Dataset","metadata":{}},{"cell_type":"code","source":"# set the paths to image and ring boundary mask folders\nimg_dir = './images'\nmask_dir = '../masks'","metadata":{"execution":{"iopub.status.busy":"2021-12-12T23:21:27.988274Z","iopub.execute_input":"2021-12-12T23:21:27.98918Z","iopub.status.idle":"2021-12-12T23:21:27.997724Z","shell.execute_reply.started":"2021-12-12T23:21:27.989126Z","shell.execute_reply":"2021-12-12T23:21:27.996854Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# lists of filenames of train, val, and test images\n# (these lists could be read in from an Excel file, for example)\n# note: filenames of corresponding images and masks should be the same\ntrain_IDs =\nval_IDs = \ntest_IDs = ","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Dataset class\nclass BacteriaDataset(Dataset):\n    \n    CLASSES = ['boundaries']\n    \n    def __init__(self, img_IDs, img_dir, mask_dir, \n                 classes=None,augmentation=None, preprocessing=None):\n        self.img_IDs = img_IDs\n        self.img_dir = img_dir\n        self.mask_dir = mask_dir  \n        self.augmentation = augmentation         # for augmentations\n        self.preprocessing = preprocessing       # preprocessing to normalize images\n        \n         # convert str names to class values on masks\n        self.class_values = [self.CLASSES.index(cls.lower()) for cls in classes]\n        \n    def __len__(self):\n        return len(self.img_IDs)\n\n    def __getitem__(self, i):\n        \n        # read data\n        img_path = os.path.join(self.img_dir, self.img_IDs[i])\n        mask_path = os.path.join(self.mask_dir, self.img_IDs[i])\n        \n        img = cv2.imread(img_path)\n        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n\n        mask = cv2.imread(mask_path, cv2.IMREAD_UNCHANGED)\n        mask = (mask >= 1).astype('float32')\n        mask = np.expand_dims(mask, axis=2) \n        \n        # apply augmentations\n        if self.augmentation:\n            sample = self.augmentation(image=img, mask=mask)\n            img, mask = sample['image'], sample['mask']\n        \n        # apply preprocessing\n        if self.preprocessing:\n            sample = self.preprocessing(image=img, mask=mask)\n            img, mask = sample['image'], sample['mask']\n            \n        return img, mask","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Transformations definitions \ndef get_training_augmentation():\n    train_transform = [albu.PadIfNeeded(min_height=1024, min_width=1024, always_apply=True, border_mode=cv2.BORDER_REFLECT_101),\n                       albu.Rotate(limit=(-10,10), border_mode=cv2.BORDER_REFLECT_101, p=0.5),\n                       albu.HorizontalFlip(p=0.5),\n                       albu.VerticalFlip(p=0.5),\n                       albu.ShiftScaleRotate(shift_limit=0.05, scale_limit=0, rotate_limit=0,\n                                          border_mode=cv2.BORDER_REFLECT_101, p=0.5), # translate\n                       albu.ShiftScaleRotate(shift_limit=0, scale_limit=0.5, rotate_limit=0,\n                                          border_mode=cv2.BORDER_REFLECT_101, p=0.5), # zoom\n                      ]\n    return albu.Compose(train_transform)\n\ndef get_val_test_augmentation():\n    val_test_transform = [\n                       albu.PadIfNeeded(min_height=1024, min_width=1024, always_apply=True, border_mode=cv2.BORDER_REFLECT_101),\n                      ]\n    return albu.Compose(val_test_transform)\n\ndef to_tensor(x, **kwargs):\n    return x.transpose(2, 0, 1).astype('float32')\n\ndef get_preprocessing(preprocessing_fn):\n    _transform = [\n        albu.Lambda(image=preprocessing_fn),\n        albu.Lambda(image=to_tensor, mask=to_tensor),\n    ]\n    return albu.Compose(_transform)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model architecture and hyperparameters","metadata":{}},{"cell_type":"code","source":"# Set some variables \nmodel_name  = 'vgg11unet' # for saving the model\nEncoder = 'vgg11'\nAttention = None # None / 'scse'\nWeights = 'imagenet' # if initializing model with pretrained weights \nACTIVATION = 'sigmoid'\nCLASSES = ['boundaries']\npreprocess_input = get_preprocessing_fn(Encoder, Weights)\npatience = 2 # for early stopping\nto_augment = False # False / True (whether or not to augment training data)\ntrain_batch_size = 3\nval_batch_size = 1\ntest_batch_size = 1","metadata":{"execution":{"iopub.status.busy":"2021-12-12T23:04:44.146933Z","iopub.execute_input":"2021-12-12T23:04:44.147593Z","iopub.status.idle":"2021-12-12T23:04:44.235598Z","shell.execute_reply.started":"2021-12-12T23:04:44.147487Z","shell.execute_reply":"2021-12-12T23:04:44.23458Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Create segmentation model with pretrained encoder\n# https://github.com/qubvel/segmentation_models.pytorch\nmodel = smp.Unet(\n    encoder_name=Encoder, \n    encoder_weights=Weights, \n    decoder_attention_type=Attention,\n    in_channels=3, \n    classes=len(CLASSES), \n    activation=ACTIVATION,\n)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# initialize loss, metrics, % optimizer:\nloss = smp.utils.losses.DiceLoss()\n\nmetrics = [\n    smp.utils.metrics.IoU(threshold=0.5),\n    smp.utils.metrics.Fscore(),\n    smp.utils.metrics.Accuracy(),\n    smp.utils.metrics.Recall(),\n    smp.utils.metrics.Precision()\n]\n\noptimizer = torch.optim.Adam([ \n    dict(params=model.parameters(), lr=0.0001),\n])","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Training & Validation","metadata":{}},{"cell_type":"code","source":"# Whether or not we're augmenting training data\nif to_augment is True:\n    training_aug = get_training_augmentation()\nelse: \n    training_aug = get_val_test_augmentation() # if not augmenting training data, just use the function for validation & test data","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Create transformed & preprocessed datasets\ntrain_dataset = BacteriaDataset(train_IDs, img_dir, mask_dir, classes=['boundaries'],\n                                augmentation=training_aug,\n                                preprocessing=get_preprocessing(preprocess_input),\n                               )\n\nval_dataset = BacteriaDataset(val_IDs, img_dir, mask_dir,classes=['boundaries'],\n                              augmentation=get_val_test_augmentation(),\n                              preprocessing=get_preprocessing(preprocess_input),\n                             )\n\ntest_dataset = BacteriaDataset(test_IDs, img_dir, mask_dir,classes=['boundaries'],\n                              augmentation=get_val_test_augmentation(),\n                              preprocessing=get_preprocessing(preprocess_input),\n                              )\n\n# Create dataloaders\ntrain_loader = DataLoader(train_dataset, batch_size=train_batch_size, shuffle=True, num_workers=12)\nval_loader = DataLoader(val_dataset, batch_size=val_batch_size, shuffle=True, num_workers=4)\ntest_loader = DataLoader(test_dataset, batch_size=test_batch_size, shuffle=False, num_workers=4)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Create epoch runners, as done in https://github.com/qubvel/segmentation_models.pytorch\ndevice = 'cuda' if torch.cuda.is_available() else 'cpu'\n\ntrain_epoch = smp.utils.train.TrainEpoch(\n    model, \n    loss=loss, \n    metrics=metrics, \n    optimizer=optimizer,\n    device=device,\n    verbose=True,\n)\n\nval_epoch = smp.utils.train.ValidEpoch(\n    model, \n    loss=loss, \n    metrics=metrics, \n    device=device,\n    verbose=True,\n)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Create dataframe for storing metrics\ndfTrainVal = pd.DataFrame(columns=['Epoch', \n                                   'Train Loss','Val Loss', \n                                   'Train Accuracy','Val Accuracy', \n                                   'Train Precision','Val Precision', \n                                   'Train Recall','Val Recall', \n                                   'Train IoU','Val IoU', \n                                   'Train Fscore','Val Fscore',\n                                   'Train Dice','Val Dice'])\n\n# For saving the dataframe:\ntrainvalCSVname = model_name + '_TrainValcsv.csv'","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Train & Validate Model\nEPOCHS = 35 # Set max number of epochs to train/validate for\nes = 0 # initiliaze early stopping counter\n\n\nfor epoch in range(0, EPOCHS):\n    \n    print('\\nEpoch: {}'.format(epoch))\n    train_logs = train_epoch.run(train_loader)\n    val_logs = val_epoch.run(val_loader)\n    \n    # Determine what the previous min val loss was \n    if epoch == 0:\n        min_val_loss = 1\n    else:\n        min_val_loss = dfTrainVal['Val Loss'].min()\n    \n    # Update the dataframe with scores from this epoch\n    dfTrainVal.loc[epoch, ['Epoch']] = epoch\n    dfTrainVal.loc[epoch, ['Train Loss']] = train_logs['dice_loss']\n    dfTrainVal.loc[epoch, ['Val Loss']] = val_logs['dice_loss']\n    dfTrainVal.loc[epoch, ['Train Accuracy']] = train_logs['accuracy']\n    dfTrainVal.loc[epoch, ['Val Accuracy']] = val_logs['accuracy']\n    dfTrainVal.loc[epoch, ['Train Precision']] = train_logs['precision']\n    dfTrainVal.loc[epoch, ['Val Precision']] = val_logs['precision']\n    dfTrainVal.loc[epoch, ['Train Recall']] = train_logs['recall']\n    dfTrainVal.loc[epoch, ['Val Recall']] = val_logs['recall']\n    dfTrainVal.loc[epoch, ['Train IoU']] = train_logs['iou_score']\n    dfTrainVal.loc[epoch, ['Val IoU']] = val_logs['iou_score']\n    dfTrainVal.loc[epoch, ['Train Fscore']] = train_logs['fscore']\n    dfTrainVal.loc[epoch, ['Val Fscore']] = val_logs['fscore']\n    dfTrainVal.loc[epoch, ['Train Dice']] = 1 - train_logs['dice_loss']\n    dfTrainVal.loc[epoch, ['Val Dice']] = 1 - val_logs['dice_loss']\n    \n    # Save the dataframe\n    dfTrainVal.to_csv(trainvalCSVname,index=False)\n    \n    # Save model checkpoints\n    checkpoint = {'epoch': epoch,\n                  'model_state_dict': model.state_dict(),\n                  'optimizer_state_dict': optimizer.state_dict(),\n                  'loss': loss}\n    checkpoint_path = './'+model_name+'_epoch_'+str(epoch)+'.pth'\n    torch.save(checkpoint, checkpoint_path)\n    \n    # Early stopping: check if val loss has decreased/increased from the previous min_val_loss\n    val_loss = val_logs['dice_loss']\n    if val_loss < min_val_loss:\n        es = 0 # Early stopping not considered\n    else: \n        es += 1 # Start counting\n        print(\"EarlyStopping Counter {} of {}\".format(es,patience))\n        \n        if es >= patience:\n            print(\"Early stopping with min_val_loss: \", min_val_loss, \"and val_loss for this epoch: \", val_loss, \"...\")\n            break","metadata":{},"execution_count":null,"outputs":[]}]}