{"cells":[{"metadata":{},"cell_type":"markdown","source":" # <div align = 'center'><u> Understanding TTA in Image Segmentation</u></div>"},{"metadata":{},"cell_type":"markdown","source":"This notebook will explain you the idea behind TTA in segmentation in a simple way."},{"metadata":{},"cell_type":"markdown","source":"# Table of contents <a id='0.1'></a>\n\n1. [What is TTA ?](#1)\n2. [TTA in Image Classification](#2)\n3. [TTA in Image Segmentation](#3)"},{"metadata":{"trusted":true,"_kg_hide-input":true,"_kg_hide-output":true},"cell_type":"code","source":"!pip install segmentation-models-pytorch\n\n# Copying diagrams to the current working directory\n!cp ../input/tta-segmentation-diagram/TTA_classification.png .\n!cp ../input/tta-segmentation-diagram/TTA_not_right_way.png .\n!cp ../input/tta-segmentation-diagram/TTA_segmentation.png .","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_kg_hide-input":true,"_kg_hide-output":false},"cell_type":"code","source":"import os\nimport numpy as np\nimport matplotlib.pyplot as plt\nfrom albumentations import *\nimport torch\nimport torch.nn as nn\nfrom segmentation_models_pytorch.unet import Unet\nfrom segmentation_models_pytorch.encoders import get_preprocessing_fn\n\nPATH_DATA = '../input/hubmap-256x256/'\nPATH_TRAINED_MODEL = '../input/tta-segmentation-diagram/hubmap_8core_model_V25.pth'","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# 1. <a id='1'>What is TTA ?</a>\n[Table of contents](#0.1)"},{"metadata":{},"cell_type":"markdown","source":"TTA stands for **Test Time Augmentation**. \n\nIn TTA, predictions of the augmented test samples (A single test sample augmented multiple times) are averaged to get the final prediction. TTA boost the performance of the model.\n\nRefer to this blog, if you are new to TTA : \n[How to Use Test-Time Augmentation to Make Better Predictions](https://machinelearningmastery.com/how-to-use-test-time-augmentation-to-improve-model-performance-for-image-classification/)"},{"metadata":{},"cell_type":"markdown","source":"# 2. <a id='2'>TTA in Image Classification</a>\n[Table of contents](#0.1)\n\nWe will understand TTA in segmentation, by understanding TTA in classification.\n\nIn image classification, model outputs a probability (score) given an input image. \nWe simply average the predicted probability of the different augmented images of a test image.\n\nLet's take an example of Melanoma Detection. Below diagram explains it well.\n\n<center><img align=\"right\" src=\"TTA_classification.png\"></center>"},{"metadata":{},"cell_type":"markdown","source":"If we would have only used orignal image for the inference, we would have got 0.88 probability. But by applying TTA, we got 0.913 score. \n\nSimilarly, TTA in segmentation will give boost in a performance. \n**But in segmentation, we have to be careful and in some cases, we need to do a extra step while applying TTA.**"},{"metadata":{},"cell_type":"markdown","source":"# 3. <a id='3'>TTA in Image Segmentation</a>\n[Table of contents](#0.1)\n\nIn image segmentation, model outputs a grid of probability (score). \n\n*Note:- We are talking about binary segmentation, where we need to detect a blob in a input image.*\n\nLet's take an example, where we are using Horizontal Flip (HF) and Vertical FLip (VF) as our augmentations in TTA.\n\n`im` is our test image whose predicted mask needs to be generated. `im_hf` and `im_vf` are the augmented images of `im` by applying HF and VF respectively.\n\n`pred1`, `pred2_hf` and `pred3_vf` are the predicted masks of the `im`, `im_hf` and `im_vf` images respectively.\n\nHere's the catch, we cannot directly average this predicted masks, becasue two of these masks are the predictions of the augmented images, so these predicted masks are augmented too. (By averaging means, we are averaging the corresponding pixels of these masks) \n\nIf we average these augmented masks, we will be averaging wrong pixels of these masks and we will end up like this, where we will get three detected blobs instead of one in the final prediction.  \n\n<center><img height = 500, width = 500, src=\"TTA_not_right_way.png\"></center>\n\n\n\nSo, we need to **deaugment the augmented predicted masks**, so that their pixels aligns with the original predicted mask and then we can average them to get the correct final prediction."},{"metadata":{},"cell_type":"markdown","source":"Corect procedure is demonstrated in the diagram below.\n\n<center><img align=\"right\" src=\"TTA_segmentation.png\"></center>"},{"metadata":{},"cell_type":"markdown","source":"Now, we will code this."},{"metadata":{},"cell_type":"markdown","source":"Utilities (Hidden)"},{"metadata":{"trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"def get_pred_mask(image):\n    preprocessing_fn = Lambda(image = get_preprocessing_fn(encoder_name = ENCODER_NAME,\n                                                       pretrained = 'imagenet'))\n    im = preprocessing_fn(image=image)['image']\n    im = np.moveaxis(im, -1, 0)\n    im = torch.from_numpy(im)\n    pred = model(im.float()[None])\n    \n    return pred.detach().numpy()[0][0]\n\n\ndef show(xs, titles, cmap = None):\n    _, axs = plt.subplots(1, len(xs))\n    \n    for i, x in enumerate(xs):\n        axs[i].imshow(x, cmap = cmap)\n        axs[i].set_title(titles[i])\n\n    plt.show()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"We will be using a model trained in this [notebook](https://www.kaggle.com/joshi98kishan/training-pytorch-tpu-8-cores)."},{"metadata":{"trusted":true,"_kg_hide-output":true},"cell_type":"code","source":"ENCODER_NAME = 'se_resnext50_32x4d'\n\nclass HuBMAPModel(nn.Module):\n    def __init__(self):\n        super(HuBMAPModel, self).__init__()\n        self.model = Unet(encoder_name = ENCODER_NAME, \n                          encoder_weights = None,\n                          classes = 1,\n                          activation = None)\n        \n    def forward(self, images):\n        img_masks = self.model(images)\n        return img_masks\n\nstate_dict = torch.load(PATH_TRAINED_MODEL)\nmodel = HuBMAPModel()\nmodel.load_state_dict(state_dict)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Reading a sample image. \n# Consider 'im' as a test image and 'mask' as ground truth.\n\nim = plt.imread(os.path.join(PATH_DATA, 'train/095bf7a1f_861.png'))\nmask = plt.imread(os.path.join(PATH_DATA, 'masks/095bf7a1f_861.png'))\n\nshow([im, mask], ['im', 'mask'], cmap = 'gray')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Defining the augmentations\n\nhorizontal_flip = HorizontalFlip(p = 1.0)\nvertical_flip = VerticalFlip(p = 1.0)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Augmentation\n\nim_hf = horizontal_flip(image = im)['image']\nim_vf = vertical_flip(image = im)['image']\n\nshow([im, im_hf, im_vf], ['im', 'im_hf', 'im_vf'])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Prediction\n\npred1 = get_pred_mask(im)\npred2_hf = get_pred_mask(im_hf)\npred3_vf = get_pred_mask(im_vf)\n\nshow([pred1, pred2_hf, pred3_vf], \n     ['pred1', 'pred2_hf', 'pred3_vf'],\n     cmap = 'gray')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Deaugmentation\n\npred2 = horizontal_flip(image = np.zeros((256, 256, 3), dtype = np.int), \n                          mask = pred2_hf)['mask']\n\npred3 = vertical_flip(image = np.zeros((256, 256, 3), dtype = np.int), \n                          mask = pred3_vf)['mask']\n\n\nshow([pred1, pred2, pred3], \n     ['pred1', 'pred2', 'pred3'],\n     cmap = 'gray')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Averaging\n\npred = (pred1 + pred2 + pred3)/3\n\nplt.imshow(pred, cmap = 'gray')\nplt.title('pred')\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Here, we have used such augmentations ([Affine transformation](https://en.wikipedia.org/wiki/Affine_transformation)) which change the location of the original pixels. But, we can also use [pixel level augmentations](https://github.com/albumentations-team/albumentations#pixel-level-transforms) in TTA, where we don't need to do deaugmentation because such augmentations will not do any affine transformation.\n\nYou can find whether you need to do deaugmentation or not, by printing the predicted mask of the augmented test image."},{"metadata":{},"cell_type":"markdown","source":"TTA in segmentaion is implemented in this notebook with a clear code:\n\n### [[Inference] PyTorch-TTA-Sub-0.84](https://www.kaggle.com/joshi98kishan/inference-pytorch-tta-sub-0-84)"},{"metadata":{},"cell_type":"markdown","source":"So, this is it. We have understood the idea.\n\nIf you liked it, then please upvote it. :)"}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}