{"cells":[{"metadata":{},"cell_type":"markdown","source":"# <div align = 'center'><u> PyTorch Inference with TTA </u></div>"},{"metadata":{},"cell_type":"markdown","source":"Our ultimate PyTorch pipeline for this competition:\n* Training on TPU - **[[FoldTraining] PyTorch-TPU🔥-8-Cores](https://www.kaggle.com/joshi98kishan/training-pytorch-tpu-8-cores)**\n* Inference with TTA (This notebook)"},{"metadata":{"trusted":true,"_kg_hide-output":true,"_kg_hide-input":true},"cell_type":"code","source":"# https://www.kaggle.com/vineeth1999/hubmap-pytorch-efficientunet-offline\n\n!mkdir -p /tmp/pip/cache/\n!cp ../input/segmentationmodelspytorch/segmentation_models/efficientnet_pytorch-0.6.3.xyz /tmp/pip/cache/efficientnet_pytorch-0.6.3.tar.gz\n!cp ../input/segmentationmodelspytorch/segmentation_models/pretrainedmodels-0.7.4.xyz /tmp/pip/cache/pretrainedmodels-0.7.4.tar.gz\n!cp ../input/segmentationmodelspytorch/segmentation_models/segmentation-models-pytorch-0.1.2.xyz /tmp/pip/cache/segmentation_models_pytorch-0.1.2.tar.gz\n!cp ../input/segmentationmodelspytorch/segmentation_models/timm-0.1.20-py3-none-any.whl /tmp/pip/cache/\n!cp ../input/segmentationmodelspytorch/segmentation_models/timm-0.2.1-py3-none-any.whl /tmp/pip/cache/\n!pip install --no-index --find-links /tmp/pip/cache/ efficientnet-pytorch\n!pip install --no-index --find-links /tmp/pip/cache/ segmentation-models-pytorch","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_kg_hide-input":true,"_kg_hide-output":false},"cell_type":"code","source":"import numpy as np\nimport pathlib\nimport pandas as pd\nimport numba, cv2, gc, os, glob\nfrom tqdm.notebook import tqdm\nimport matplotlib.pyplot as plt\n\nfrom albumentations import *\n\nimport torch\nimport torch.nn as nn\n\nfrom segmentation_models_pytorch.unet import Unet\nfrom segmentation_models_pytorch.encoders import get_preprocessing_fn\n\nimport rasterio\nfrom rasterio.windows import Window\n\nDEVICE = 'cuda' if torch.cuda.is_available() else 'cpu' \n\nimport warnings\nwarnings.filterwarnings('ignore')\n\nprint(DEVICE)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"I have forked the training nb, and trained the fold models for 60 epochs on each fold."},{"metadata":{"trusted":true},"cell_type":"code","source":"DATA_PATH = '../input/hubmap-kidney-segmentation'\n\n# path to our training notebook.\nPATH_FOLD_MODELS = '../input/test-training-pytorch-tpu-8-cores'","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Utilities (Hidden)"},{"metadata":{"trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"@numba.njit()\ndef rle_numba(pixels):\n    size = len(pixels)\n    points = []\n    if pixels[0] == 1: points.append(0)\n    flag = True\n    for i in range(1, size):\n        if pixels[i] != pixels[i-1]:\n            if flag:\n                points.append(i+1)\n                flag = False\n            else:\n                points.append(i+1 - points[-1])\n                flag = True\n    if pixels[-1] == 1: points.append(size-points[-1]+1)    \n    return points\n\ndef rle_numba_encode(image):\n    pixels = image.flatten(order = 'F')\n    points = rle_numba(pixels)\n    return ' '.join(str(x) for x in points)\n\ndef make_grid(shape, window=256, min_overlap=32):\n    \"\"\"\n        Return Array of size (N,4), where N - number of tiles,\n        2nd axis represente slices: x1,x2,y1,y2 \n    \"\"\"\n    x, y = shape\n    nx = x // (window - min_overlap) + 1\n    x1 = np.linspace(0, x, num=nx, endpoint=False, dtype=np.int64)\n    x1[-1] = x - window\n    x2 = (x1 + window).clip(0, x)\n    ny = y // (window - min_overlap) + 1\n    y1 = np.linspace(0, y, num=ny, endpoint=False, dtype=np.int64)\n    y1[-1] = y - window\n    y2 = (y1 + window).clip(0, y)\n    slices = np.zeros((nx,ny, 4), dtype=np.int64)\n    \n    for i in range(nx):\n        for j in range(ny):\n            slices[i,j] = x1[i], x2[i], y1[j], y2[j]    \n    return slices.reshape(nx*ny,4)","execution_count":null,"outputs":[]},{"metadata":{"trusted":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","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"fold_models_paths = glob.glob(os.path.join(PATH_FOLD_MODELS, '*.pth'))\nfold_models = []\n\nfor path in fold_models_paths:\n    state_dict = torch.load(path)\n    model = HuBMAPModel()\n    model.load_state_dict(state_dict)\n    model.float()\n    model.to(DEVICE)\n    model.eval()\n    \n    fold_models.append(model)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"len(fold_models)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"preprocess_input = Lambda(image = get_preprocessing_fn(encoder_name = ENCODER_NAME,\n                                                       pretrained = 'imagenet'))\n\nidentity_trfm = Lambda(image = lambda x,cols=None,rows=None : x)\n\n# Affine transforms\nhorizontal_flip = HorizontalFlip(p = 1.0)\nvertical_flip = VerticalFlip(p = 1.0)\nrotate_cw = Rotate(limit = (-90, -90), p = 1.0)\nrotate_acw = Rotate(limit = (90, 90), p = 1.0)\n\n# Pixel level transformations\npixel_level_trfms = OneOf([\n                    HueSaturationValue(10,15,10),\n                    CLAHE(clip_limit=2),\n                    RandomBrightnessContrast(),            \n                   ], p = 1.0)\n\n# List of augmentations for TTA\ntta_augs = [identity_trfm,\n            horizontal_flip,\n            vertical_flip,\n            rotate_cw,\n            pixel_level_trfms]\n\n# List of deaugmentations corresponding to the above aug list\ntta_deaugs = [None,\n              horizontal_flip,\n              vertical_flip,\n              rotate_acw,\n              None]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"WINDOW=1024\nMIN_OVERLAP=32\nNEW_SIZE=256","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"We will be doing TTA here. \n\nI have explained TTA in this notebook \"[Let's Understand TTA in Segmentation](https://www.kaggle.com/joshi98kishan/let-s-understand-tta-in-segmentation)\" in a simplest way possible."},{"metadata":{"trusted":true},"cell_type":"code","source":"identity = rasterio.Affine(1, 0, 0, 0, 1, 0)\np = pathlib.Path(DATA_PATH)\nsubm = {}\n\nfor i, filename in tqdm(enumerate(p.glob('test/*.tiff')), \n                        total = len(list(p.glob('test/*.tiff')))):\n    \n    print(f'{i+1} Predicting {filename.stem}')\n    \n    dataset = rasterio.open(filename.as_posix(), transform = identity)\n    slices = make_grid(dataset.shape, window=WINDOW, min_overlap=MIN_OVERLAP)\n    preds = np.zeros(dataset.shape, dtype=np.uint8)\n    \n    for (x1,x2,y1,y2) in slices:\n        image = dataset.read([1,2,3],\n                    window=Window.from_slices((x1,x2),(y1,y2)))\n        image = np.moveaxis(image, 0, -1)\n        pred = 0\n        \n        for fold_model in fold_models:  \n            tta_pred = None\n            \n            for j, tta_aug in enumerate(tta_augs):\n                # Augmentation\n                aug_img = tta_aug(image = image)['image']\n                aug_img = preprocess_input(image = aug_img)['image']\n                aug_img = cv2.resize(aug_img, (NEW_SIZE, NEW_SIZE))\n                aug_img = np.moveaxis(aug_img, -1, 0)\n                aug_img = torch.from_numpy(aug_img)\n        \n                with torch.no_grad():\n                    score = fold_model(aug_img.float().to(DEVICE)[None])\n                    score = score.cpu().numpy()[0][0]\n                    \n                    # Deaugmentation\n                    if tta_deaugs[j] is not None:\n                        score = tta_deaugs[j](image = image, \n                                              mask = score)['mask']\n\n                    score = cv2.resize(score, (WINDOW, WINDOW))            \n\n                    if tta_pred is None:\n                        tta_pred = score\n                    else:       \n                        tta_pred += score\n             \n            tta_pred = tta_pred / len(tta_augs) \n            pred += tta_pred\n            \n        pred = pred / len(fold_models)\n        preds[x1:x2,y1:y2] = (pred > 0).astype(np.uint8)\n            \n    subm[i] = {'id':filename.stem, 'predicted': rle_numba_encode(preds)}\n    del preds\n    gc.collect();","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"submission = pd.DataFrame.from_dict(subm, orient='index')\nsubmission.to_csv('submission.csv', index=False)\n\nsubmission.head()","execution_count":null,"outputs":[]}],"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}