{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Change Log\n\nVersion 1: augmentations, added UNET model with SeResNext-50_32x4d encoder, predictions visualization <br>\nVersion 2: added predictions by one slide, padded and resized images/masks concatenation <br>\nVersion 5: added inference loop and submission <br>"},{"metadata":{},"cell_type":"markdown","source":"This kernel is based on my previous one [HuBMAP: train/test patches generation](https://www.kaggle.com/mariazorkaltseva/hubmap-train-test-patches-generation)\n\n**Configuration:**\n\nTile size: 256 without overlaping, contain artifacts\n\nreduce rate: 4\n\nBatch_size: 16 on train/val set, 1 on test set\n\nAugmentations: Yes\n\nModel: UNET architecture with SeResNext-50_32x4d encoder\n\nOptimizer: Adam with lr=0.0001\n\nScoring function: Dice loss\n\nMetric: IoU with 0.5 threshold\n"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"_kg_hide-output":true},"cell_type":"code","source":"!pip install segmentation_models_pytorch","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"import os\nimport gc\nimport json\nimport random\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nimport tifffile as tiff\nimport cv2 as cv\nimport albumentations as albu\n\nfrom tqdm.notebook import tqdm\nfrom PIL import Image\nfrom sklearn.model_selection import train_test_split\n\nimport torch\nimport torch.nn as nn\nimport torch.utils.data as D\nimport torch.nn.functional as F\nimport segmentation_models_pytorch as smp\n        \n%matplotlib inline","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"TRAIN_DATA_DIR = '../input/hubmap-256-tiles/train_tiles_256/train_tiles'\nTEST_DATA_DIR = '../input/hubmap-256-tiles/test_tiles_256/test_tiles'\nMODEL_SAVE_DIR = \"/kaggle/working/\"\nTILE_SIZE = 256\nREDUCE_RATE = 4\nSEED = 42\nBATCH_SIZE = 16\nNUM_EPOCHS = 20\n\ntorch.cuda.empty_cache()\n\n\nrandom.seed(SEED)\nnp.random.seed(SEED)\ntorch.manual_seed(SEED)\ntorch.cuda.manual_seed(SEED)\ntorch.backends.cudnn.deterministic = True\n\nDEVICE = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Helpers"},{"metadata":{"trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"def display_pil_images(\n    images, \n    masks=None,\n    labels=None,\n    columns=5, width=20, height=8, max_images=15, \n    label_wrap_length=50, label_font_size=9):\n\n    if len(images) > max_images:\n        print(f\"Showing {max_images} images of {len(images)}:\")\n        images=images[0:max_images]\n        if masks is not None:\n            masks= masks[0:max_images]\n\n    height = max(height, int(len(images)/columns) * height)\n    plt.figure(figsize=(width, height))\n    \n    if masks is not None:\n        for i, (image, mask) in enumerate(zip(images,masks)):\n            plt.subplot(len(images) / columns + 1, columns, i + 1)\n            plt.imshow(image)\n            plt.imshow(mask, cmap='coolwarm', alpha=0.5)\n            \n            if labels is not None:\n                plt.title(labels[i], fontsize=label_font_size); \n            \n    else:\n        for i, image in enumerate(images):\n            plt.subplot(len(images) / columns + 1, columns, i + 1)\n            plt.imshow(image)\n        \n            if labels is not None:\n                plt.title(labels[i], fontsize=label_font_size);\n    \n\ndef visualize(**images):\n    \"\"\"PLot images in one row.\"\"\"\n    n = len(images)\n    plt.figure(figsize=(16, 5))\n    for i, (name, image) in enumerate(images.items()):\n        plt.subplot(1, n, i + 1)\n        plt.xticks([])\n        plt.yticks([])\n        plt.title(' '.join(name.split('_')).title())\n        plt.imshow(image)\n    plt.show()\n    \n#https://www.kaggle.com/bguberfain/memory-aware-rle-encoding\n#with transposed mask\ndef rle_encode_less_memory(img):\n    #the image should be transposed\n    pixels = img.T.flatten()\n    \n    # This simplified method requires first and last pixel to be zero\n    pixels[0] = 0\n    pixels[-1] = 0\n    runs = np.where(pixels[1:] != pixels[:-1])[0] + 2\n    runs[1::2] -= runs[::2]\n    \n    return ' '.join(str(x) for x in runs)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"sample_df = pd.read_csv('../input/hubmap-kidney-segmentation/sample_submission.csv')\nsample_df","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Image paths generation"},{"metadata":{"trusted":true},"cell_type":"code","source":"train_img_paths, train_mask_paths = [], []\nfor item1 in os.listdir(TRAIN_DATA_DIR):\n    for item2 in os.listdir(os.path.join(TRAIN_DATA_DIR, item1)):\n        if item2.startswith('img_'):\n            train_img_paths.append(os.path.join(TRAIN_DATA_DIR, item1, item2))\n        else:\n            train_mask_paths.append(os.path.join(TRAIN_DATA_DIR, item1, item2))\n            \n\ntest_img_paths = []\nfor item1 in os.listdir(TEST_DATA_DIR):\n    for item2 in os.listdir(os.path.join(TEST_DATA_DIR, item1)):\n        test_img_paths.append(os.path.join(TEST_DATA_DIR, item1, item2))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def sort_by(x):\n    return x.rsplit('/', 1)[1]\n\ntrain_img_paths = sorted(train_img_paths, key = sort_by)\ntrain_mask_paths = sorted(train_mask_paths, key = sort_by)\ntest_img_paths = sorted(test_img_paths, key = sort_by)\n        \nprint(train_img_paths[:5])\nprint(train_mask_paths[:5])\nprint(test_img_paths[:5])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print(f\"Amount of samples in the train set {len(train_img_paths)}\")\nprint(f\"Amount of samples in the test set {len(test_img_paths)}\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# View images and masks"},{"metadata":{"trusted":true},"cell_type":"code","source":"imgs = [Image.open(img_path) for img_path in train_img_paths]\nmasks = [Image.open(mask_path) for mask_path in train_mask_paths]\ndisplay_pil_images(imgs[300:500], masks[300:500])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"imgs = [Image.open(img_path) for img_path in test_img_paths]\ndisplay_pil_images(imgs[450:500])","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Dataset class"},{"metadata":{"trusted":true},"cell_type":"code","source":"class HuBMAPDataset(D.Dataset):\n    \n    def __init__(\n            self, \n            paths, \n            mode,\n            augmentation=None,\n            preprocessing=None,\n    ):\n\n        self.paths = paths\n        self.mode = mode\n        self.augmentation = augmentation\n        self.preprocessing = preprocessing\n        \n    def __getitem__(self, i):       \n        if self.mode in ['train', 'val']:\n            image = cv.imread(self.paths[i][0])\n            image = cv.cvtColor(image, cv.COLOR_BGR2RGB)\n            mask = cv.imread(self.paths[i][1], 0)\n            mask = np.expand_dims(mask, axis=2)\n        else:\n            image = cv.imread(self.paths[i])\n            image = cv.cvtColor(image, cv.COLOR_BGR2RGB)\n\n        if self.augmentation:\n            if self.mode in ['train', 'val']:\n                sample = self.augmentation(image=image, mask=mask)\n                image, mask = sample['image'], sample['mask']\n            else:\n                sample = self.augmentation(image=image)\n                image = sample['image']\n\n        if self.preprocessing:\n            if self.mode in ['train', 'val']:\n                sample = self.preprocessing(image=image, mask=mask)\n                image, mask = sample['image'], sample['mask']\n            else:\n                sample = self.preprocessing(image=image)\n                image = sample['image']\n\n        if self.mode in ['train', 'val']:\n            return image, mask\n        \n        return image\n        \n    def __len__(self):\n        return len(self.paths)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Augmentations"},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_training_augmentation():\n    train_transform = [\n\n        albu.HorizontalFlip(p=0.5),\n        albu.VerticalFlip(p=0.5),\n        albu.RandomRotate90(p=0.5),\n        albu.Transpose(p=0.5),\n        \n        albu.ShiftScaleRotate(scale_limit=0.2, rotate_limit=0, shift_limit=0.2, p=0.2, border_mode=0),\n\n        albu.IAAAdditiveGaussianNoise(p=0.2),\n        albu.IAAPerspective(p=0.5),\n\n        albu.OneOf(\n            [\n                albu.CLAHE(p=1),\n                albu.RandomBrightness(p=1),\n                albu.RandomGamma(p=1),\n            ],\n            p=0.9,\n        ),\n\n        albu.OneOf(\n            [\n                albu.IAASharpen(p=1),\n                albu.Blur(blur_limit=3, p=1),\n                albu.MotionBlur(blur_limit=3, p=1),\n            ],\n            p=0.9,\n        ),\n\n        albu.OneOf(\n            [\n                albu.RandomContrast(p=1),\n                albu.HueSaturationValue(p=1),\n            ],\n            p=0.9,\n        ),\n        \n        albu.Compose([\n            albu.VerticalFlip(p=0.5),              \n            albu.RandomRotate90(p=0.5)]\n        )\n    ]\n    return albu.Compose(train_transform)\n\n\ndef get_validation_augmentation():\n    test_transform = [\n    ]\n    return albu.Compose(test_transform)\n\n\ndef to_tensor(x, **kwargs):\n    return x.transpose(2, 0, 1).astype('float32')\n\n\ndef get_preprocessing():\n    _transform = [\n        albu.Normalize(mean=(0.65459856,0.48386562,0.69428385), \n                       std=(0.15167958,0.23584107,0.13146145), \n                       max_pixel_value=255.0, always_apply=True, p=1.0),\n        albu.Lambda(image=to_tensor, mask=to_tensor),\n    ]\n    return albu.Compose(_transform)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_paths, val_paths = train_test_split(list(zip(train_img_paths, train_mask_paths)),\n                                        test_size=0.2,  \n                                        random_state=SEED,\n                                        shuffle=True)\n\nprint(\"Amount of train samples:\", len(train_paths))\nprint(\"Amount of val samples:\", len(val_paths))","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-output":false,"trusted":true},"cell_type":"code","source":"augmented_dataset = HuBMAPDataset(train_paths,\n                                  'train',\n                                  augmentation=get_training_augmentation(),\n                                )\n\nfor i in range(20):\n    image, mask = augmented_dataset[1]\n    visualize(image=image, mask=mask.squeeze(-1))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Model"},{"metadata":{"trusted":true},"cell_type":"code","source":"ENCODER = 'se_resnext50_32x4d'\nENCODER_WEIGHTS = 'imagenet'\nACTIVATION = 'sigmoid' \n\nmodel = smp.Unet(encoder_name=ENCODER, encoder_weights=ENCODER_WEIGHTS, activation=ACTIVATION)\n\nloss = smp.utils.losses.DiceLoss()\nmetrics = [\n    smp.utils.metrics.IoU(threshold=0.5),\n]\n\noptimizer = torch.optim.Adam([ \n    dict(params=model.parameters(), lr=0.0001),\n])","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Train/val dataloaders"},{"metadata":{"trusted":true},"cell_type":"code","source":"train_dataset = HuBMAPDataset(train_paths, \n                              'train',\n                              augmentation=get_training_augmentation(), \n                              preprocessing=get_preprocessing()\n                            )\n\nvalid_dataset = HuBMAPDataset(val_paths,\n                              'val',\n                              augmentation=get_validation_augmentation(), \n                              preprocessing=get_preprocessing()\n                            )\n\ntrain_loader = D.DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=4, pin_memory=True)\nvalid_loader = D.DataLoader(valid_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=4, pin_memory=True)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Train loop"},{"metadata":{"trusted":true},"cell_type":"code","source":"train_epoch = smp.utils.train.TrainEpoch(\n    model, \n    loss=loss, \n    metrics=metrics, \n    optimizer=optimizer,\n    device=DEVICE,\n    verbose=True,\n)\n\nvalid_epoch = smp.utils.train.ValidEpoch(\n    model, \n    loss=loss, \n    metrics=metrics, \n    device=DEVICE,\n    verbose=True,\n)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"best_loss = 1.0\n\ntrain_losses, val_losses = [], []\ntrain_scores, val_scores = [], []\n\nfor i in range(0, NUM_EPOCHS):\n    \n    print('\\nEpoch: {}'.format(i))\n    train_logs = train_epoch.run(train_loader)\n    valid_logs = valid_epoch.run(valid_loader)\n    \n    train_losses.append(train_logs['dice_loss'])\n    val_losses.append(valid_logs['dice_loss'])\n    train_scores.append(train_logs['iou_score'])\n    val_scores.append(valid_logs['iou_score'])\n    \n    if best_loss > valid_logs['dice_loss']:\n        best_loss = valid_logs['dice_loss']\n        torch.save(model, os.path.join(MODEL_SAVE_DIR, 'best_model.pth'))\n        print('Model saved!')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plt.figure()\nplt.plot(train_losses, label='train loss')\nplt.plot(val_losses, label='val loss')\nplt.legend()\nplt.show();","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plt.figure()\nplt.plot(train_scores, label='train score')\nplt.plot(val_scores, label='val score')\nplt.legend()\nplt.show();","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Results visualization"},{"metadata":{"trusted":true},"cell_type":"code","source":"best_model = torch.load(os.path.join(MODEL_SAVE_DIR, 'best_model.pth'), map_location=torch.device(DEVICE))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_dataset = HuBMAPDataset(test_img_paths, \n                             'test',\n                             preprocessing=get_preprocessing()\n                        )\n\ntest_dataloader = D.DataLoader(test_dataset, batch_size=1, shuffle=False, num_workers=4)\n\ntest_dataset_vis = HuBMAPDataset(test_img_paths, 'test')\n\nbest_model.eval()\n\nfor i in range(30):\n    n = np.random.choice(len(test_dataset))\n    \n    image_vis = test_dataset_vis[n].astype('uint8')\n    image = test_dataset[n]\n    \n    x_tensor = torch.from_numpy(image).to(DEVICE).unsqueeze(0)\n    with torch.set_grad_enabled(False):\n        pr_mask = best_model.predict(x_tensor)\n        pr_mask = (pr_mask.squeeze().cpu().numpy().round().astype('uint8'))\n        \n    visualize(\n        image=image_vis, \n        predicted_mask=pr_mask\n    )","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Сombining prediction masks"},{"metadata":{"trusted":true},"cell_type":"code","source":"def extract_slide_tiles(file_paths, id):\n    def sort_by(x):\n        # extract tile numbers\n        return int(x.rsplit('/', 1)[1].split('_')[1].split('.')[0])\n    return sorted(list(filter(lambda x: id in x, file_paths)), key=sort_by)\n\nIDX = 'b2dc8411c'\nslide_paths = extract_slide_tiles(test_img_paths, IDX)\n\n# padded and reduced shape of corresponding image\nheight, width = (3840, 7936)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# one slide predictions\ntest_dataset = HuBMAPDataset(slide_paths, \n                             'test',\n                             preprocessing=get_preprocessing()\n                        )\n\ntest_dataloader = D.DataLoader(test_dataset, batch_size=1, shuffle=False, num_workers=4)\n\ntest_dataset_vis = HuBMAPDataset(slide_paths, 'test')\n\nbest_model.eval()\n\nmask_preds = []\nfor i, image in enumerate(tqdm(test_dataloader)):\n    image_vis = test_dataset_vis[i].astype('uint8')\n    \n    image = image.to(DEVICE)\n    with torch.set_grad_enabled(False):\n        mask_pred = best_model(image)\n        mask_pred = mask_pred.squeeze().cpu().numpy().round().astype('uint8')\n    mask_preds.append(np.expand_dims(mask_pred, axis=0))\n    \n    if i % 20 == 0:\n        visualize(\n            image=image_vis, \n            predicted_mask=mask_pred\n        )\n\nmask_preds = np.concatenate(mask_preds)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"merge_image = np.zeros((height, width, 3))\n\nk = 0\nfor i in range(0, height // TILE_SIZE):\n    for j in range(0, width // TILE_SIZE):\n        image = cv.imread(slide_paths[k])\n        image = cv.cvtColor(image, cv.COLOR_BGR2RGB)\n        merge_image[i*TILE_SIZE:i*TILE_SIZE + TILE_SIZE, j*TILE_SIZE:j*TILE_SIZE + TILE_SIZE, :] = image\n        k += 1","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"merge_mask = np.zeros((height, width))\n\nk = 0\nfor i in range(0, height // TILE_SIZE):\n    for j in range(0, width // TILE_SIZE):\n        merge_mask[i*TILE_SIZE:i*TILE_SIZE + TILE_SIZE, j*TILE_SIZE:j*TILE_SIZE + TILE_SIZE] = mask_preds[k]\n        k += 1","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plt.figure(figsize=(16, 16))\nplt.imshow(merge_image.astype('uint8'))\nplt.imshow(merge_mask.astype('uint8'), cmap='coolwarm', alpha=0.5);","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Inference loop"},{"metadata":{"trusted":true},"cell_type":"code","source":"imgs_size_df = pd.read_csv(\"../input/hubmap-img-sizes/imgs_size.csv\")\nimgs_size_df.set_index(imgs_size_df['id'], inplace=True)\n\nbest_model.eval()\n\nrle_preds = []\nfor idx in sample_df['id']:\n    print(idx)\n    \n    slide_paths = extract_slide_tiles(test_img_paths, idx)\n    \n    test_dataset = HuBMAPDataset(slide_paths, \n                                 'test',\n                                 preprocessing=get_preprocessing()\n                                )\n\n    test_dataloader = D.DataLoader(test_dataset, batch_size=1, shuffle=False, num_workers=4, pin_memory=True)\n\n    mask_preds = []\n    for i, image in enumerate(tqdm(test_dataloader)):\n\n        image = image.to(DEVICE)\n        with torch.set_grad_enabled(False):\n            mask_pred = best_model(image)\n            mask_pred = mask_pred.squeeze().cpu().numpy().round().astype('uint8')\n        mask_preds.append(np.expand_dims(mask_pred, axis=0))\n\n    mask_preds = np.concatenate(mask_preds)\n    \n    original_shape = imgs_size_df['orig_height'].loc[idx], imgs_size_df['orig_width'].loc[idx]\n    height, width = imgs_size_df['reduced_height'].loc[idx], imgs_size_df['reduced_width'].loc[idx]\n    \n    merge_mask = np.zeros((height, width), dtype=np.uint8)\n\n    k = 0\n    for i in range(0, height // TILE_SIZE):\n        for j in range(0, width // TILE_SIZE):\n            merge_mask[i*TILE_SIZE:i*TILE_SIZE + TILE_SIZE, j*TILE_SIZE:j*TILE_SIZE + TILE_SIZE] = mask_preds[k]\n            k += 1\n\n    pad0 = (REDUCE_RATE*TILE_SIZE - original_shape[0]%(REDUCE_RATE*TILE_SIZE))%(REDUCE_RATE*TILE_SIZE)\n    pad1 = (REDUCE_RATE*TILE_SIZE - original_shape[1]%(REDUCE_RATE*TILE_SIZE))%(REDUCE_RATE*TILE_SIZE)\n    \n    merge_mask = cv.resize(merge_mask,(merge_mask.shape[1]*REDUCE_RATE, merge_mask.shape[0]*REDUCE_RATE), \n                    interpolation = cv.INTER_AREA)\n\n    merge_mask = merge_mask[pad0//2:-(pad0-pad0//2) if pad0>0 else merge_mask.shape[0], \\\n                            pad1//2:-(pad1-pad1//2) if pad1>0 else merge_mask.shape[1]]\n        \n    rle_preds.append(rle_encode_less_memory(merge_mask))\n    \n    del slide_paths, test_dataset, test_dataloader, mask_preds, merge_mask\n    gc.collect()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"sample_df['predicted'] = rle_preds\nsample_df.to_csv('submission.csv')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","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}