{"cells":[{"metadata":{},"cell_type":"markdown","source":"Fine-tuning UNet with fastai v1"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import os\nfrom tqdm import tqdm\nfrom skimage.morphology import label, binary_opening, disk\nfrom fastai import vision, basic_data, layers, metrics\nfrom fastai.callbacks import hooks\nimport torch \nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.autograd import Variable\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\ndata_root = '../input/airbus-ship-detection/'\npath_train = os.path.join(data_root,'train_v2')\npath_test = os.path.join(data_root,'test_v2')\n\nimg_shape = (768, 768)\n\n# Booleans\nUSE_SELF_ATTENTION = True\nUSE_UNET34_AIRBUS = False\nUSE_FULL_RES_PRED = False","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Get groundtruth and remove empty images"},{"metadata":{"trusted":true},"cell_type":"code","source":"# Get dataframe with label\nmasks_df = pd.read_csv(os.path.join(data_root, 'train_ship_segmentations_v2.csv'))\nmasks_df = masks_df[~masks_df['ImageId'].isin(['6384c3e78.jpg'])]  # remove corrupted image\nmasks_df = masks_df.dropna() # remove images withtout ships\nunique_img_ids_df = masks_df.groupby('ImageId').size().reset_index(name='counts')\nunique_img_ids_df = unique_img_ids_df.drop(columns=['counts'])","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Losses"},{"metadata":{"trusted":true},"cell_type":"code","source":"# https://www.kaggle.com/iafoss/unet34-dice-0-87\n# https://becominghuman.ai/investigating-focal-and-dice-loss-for-the-kaggle-2018-data-science-bowl-65fb9af4f36c\ndef cuda(x): \n    return x.cuda() if torch.cuda.is_available() else x\n\ndef make_one_hot(labels, c=2):\n    one_hot = cuda(torch.FloatTensor(labels.size(0), c, labels.size(2), labels.size(3)).zero_())\n    target = one_hot.scatter_(1, labels.data, 1)\n    target = cuda(Variable(target))\n    return target\n\ndef dice_loss(input, target):\n    \"\"\"Soft dice loss function.\n    https://github.com/pytorch/pytorch/issues/1249\"\"\"\n    # Input is of shape N,C,H,W\n    smooth = 1\n    batch_size = input.size(0)\n    input = F.softmax(input, dim=1)\n    # Since we have only 2 classes transform it to N,H,W and treat as sigmoid\n    input = input.view(batch_size, 2, -1)[:, 1, :]\n    target = make_one_hot(target).view(batch_size, 2, -1)[:, 1, :]\n\n    inter = torch.sum(input * target) + smooth\n    union = torch.sum(input) + torch.sum(target) + smooth\n\n    return -torch.log(2.0 * inter / union)\n\nclass FocalLoss(nn.Module):\n    \"\"\"Focal loss function.\"\"\"\n    def __init__(self, gamma):\n        super(FocalLoss, self).__init__()\n        self.gamma = gamma\n\n    def forward(self, input, target):\n        # One-hot encode target\n        target = target.squeeze(1)\n    \n        input = input.view(input.size(0),input.size(1),-1) # N,C,H,W => N,C,H*W\n        input = input.transpose(1,2)                       # N,C,H*W => N,H*W,C\n        input = input.contiguous().view(-1,input.size(2))  # N,H*W,C => N*H*W,C\n        target = target.view(-1,1)\n\n        logpt = F.log_softmax(input, dim=1)\n        logpt = logpt.gather(1,target)\n        logpt = logpt.view(-1)\n        pt = Variable(logpt.data.exp())\n\n        loss = -1 * (1-pt)**self.gamma * logpt\n        return loss.mean()\n    \nclass MixedLoss(nn.Module):\n    \"\"\"Combine two losses and bring them to similar scale.\"\"\"\n    def __init__(self, alpha, gamma):\n        super().__init__()\n        self.alpha = alpha\n        self.focal = FocalLoss(gamma)\n        \n    def forward(self, input, target):\n        return self.alpha * self.focal(input, target) + dice_loss(input, target)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Metrics"},{"metadata":{"trusted":true},"cell_type":"code","source":"# https://nbviewer.jupyter.org/github/polakowo/mlprojects/blob/master/airbus-ship-segmentation/ShipDetection.ipynb\ndef dice(input, target):\n    \"\"\"Dice metric.\"\"\"\n    input = input.argmax(dim=1).float()\n    target = target.squeeze(1).float()\n    \n    smooth = 1.\n    numerator = 2. * (input * target).sum()\n    denumerator = (input + target).sum()\n    return (numerator + smooth) / (denumerator + smooth)\n\ndef IoU(input, target):\n    \"\"\"Intersection over Union (IoU) metric.\"\"\"\n    input = input.argmax(dim=1).float()\n    target = target.squeeze(1).float()\n    \n    smooth = 1.\n    intersection = (input * target).sum()\n    union = (input + target).sum() - intersection\n    return (intersection + smooth) / (union + smooth)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Data loader"},{"metadata":{"trusted":true},"cell_type":"code","source":"def open_mask(fn):\n    masks = masks_df[masks_df['ImageId'] == str(os.path.split(fn)[1])]['EncodedPixels'].tolist()\n    masks = \" \".join(str(x) for x in masks) # convert list to string\n    mask_img = vision.image.open_mask_rle(masks, shape=(768, 768))\n    return vision.ImageSegment(mask_img.data.T.permute(2,0,1).float())\n     \nclass SegmentationLabelList(vision.ImageList):\n    \"\"\"Our labels are created from encodings in masks_df, no disk I/O operations required.\"\"\"\n    _processor=vision.data.SegmentationProcessor\n    def __init__(self, items:basic_data.Iterator, classes:basic_data.Collection=None, **kwargs):\n        super().__init__(items, **kwargs)\n        self.copy_new.append('classes')\n        self.classes, self.loss_func = classes, layers.CrossEntropyFlat(axis=1)\n\n    def open(self, fn): return open_mask(fn)\n    def analyze_pred(self, pred, thresh:float=0.5): return pred.argmax(dim=0)[None]\n    def reconstruct(self, t:basic_data.Tensor): return vision.ImageSegment(t)\n\nclass SegmentationItemList(vision.ImageList):\n    \"`ItemList` suitable for segmentation tasks.\"\n    _label_cls, _square_show_res = SegmentationLabelList, False\n    \ndef get_data(df, bs=16, img_size=(256, 256)):\n    # Do not augment since we have a large dataset anyway, only resize\n    tfms = ([vision.transform.crop_pad()], [vision.transform.crop_pad()])\n    \n    # Build DataBunch\n    return (SegmentationItemList.from_df(unique_img_ids_df, path=path_train)\n            .split_by_rand_pct(0.2)\n            .label_from_func(lambda x: x, classes=[0, 1])\n            .transform(tfms, size=img_size, tfm_y=True)\n            .add_test(vision.Path(path_test).ls(), tfm_y=False)\n            .databunch(path=data_root, bs=bs)\n            .normalize(vision.imagenet_stats))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Build model"},{"metadata":{"trusted":true},"cell_type":"code","source":"# Build databunch\ndata = get_data(masks_df)\n\nif USE_UNET34_AIRBUS == True:\n    model = torch.load('../input/unet-resnet34-airbus/UNET34Airbus_2cl.pth')\n    learner = vision.Learner(data, model, loss_func=MixedLoss(10., 2.), metrics=[dice, IoU])\nelse:\n    model = vision.models.resnet34\n    learner = vision.unet_learner(data, model, loss_func=MixedLoss(10., 2.), metrics=[dice, IoU], self_attention=USE_SELF_ATTENTION)\n    \nlearner.model_dir = '/kaggle/working'","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"data.show_batch()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Training"},{"metadata":{"trusted":true},"cell_type":"code","source":"# Find optimal LR\nlearner.lr_find()\nlearner.recorder.plot(suggestion=True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Train on 3 epoch\nlearner.fit_one_cycle(3, max_lr=5e-4)\nlearner.recorder.plot_losses()\nlearner.recorder.plot_lr(show_moms=True)\nlearner.save('unet_3ep')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Train on 5 more epoch with model unfreezed\nlearner.unfreeze()\nlearner.lr_find()\nlearner.recorder.plot(suggestion=True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"learner.fit_one_cycle(2, max_lr=slice(1e-6, 1e-5))\nlearner.recorder.plot_losses()\nlearner.recorder.plot_metrics()\nlearner.recorder.plot_lr(show_moms=True)\nlearner.save('unet_5ep')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"learner.show_results()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Testing"},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_test_masks():\n    \n    if USE_FULL_RES_PRED == True:\n        learner.data = get_data(masks_df, img_size=(768, 768))\n\n    pred_masks = []\n    for x, y in tqdm(learner.data.test_ds):\n        _, _, output = learner.predict(x) # network output 2x256x256 or 2x768x768\n        \n        if USE_FULL_RES_PRED == False:\n            upsampler = torch.nn.Upsample(scale_factor=3, mode='bilinear', align_corners=False) # 768/256 = factor of 3\n            output = upsampler(output.unsqueeze(0)).squeeze(0)  # 2x256x256\n        \n        probs = F.softmax(output, dim=0)  # 2x256x256 or 2x768x768\n        mask_tensor = probs.argmax(dim=0)   # 256x256 or 768x768 (hot tensor)\n        \n        labels = label(mask_tensor)\n        pred_masks.append([vision.image.rle_encode((labels.T)==k) for k in np.unique(labels[labels>0])])\n    return pred_masks\n\ndef get_test_masks_opening():\n    pred_masks = []\n    for x, y in tqdm(learner.data.test_ds):\n        _, _, output = learner.predict(x) # network output 2x256x256 or 2x768x768\n\n        if USE_FULL_RES_PRED == False:\n            upsampler = torch.nn.Upsample(scale_factor=3, mode='bilinear', align_corners=False) # 768/256 = factor of 3\n            output = upsampler(output.unsqueeze(0)).squeeze(0)  # 2x256x256\n\n        probs = F.softmax(output, dim=0)  # 2x256x256 or 2x768x768\n        mask_tensor = probs.argmax(dim=0)   # 256x256 or 768x768 (hot tensor)\n\n        mask_tensor = binary_opening(mask_tensor, disk(2))\n        labels = label(mask_tensor)\n        pred_masks.append([vision.image.rle_encode((labels.T)==k) for k in np.unique(labels[labels>0])])\n    return pred_masks\n    \ndef create_submission_df(test_masks):\n    \"\"\"Create submission dataframe.\"\"\"\n    test_ids = list(map(lambda x: x.name, learner.data.test_dl.dataset.items))\n    img_masks = list(zip(test_ids, test_masks))\n    flat_img_masks = [] \n    for img, masks in img_masks:\n        if len(masks) > 0:\n            for mask in masks:\n                flat_img_masks.append([img, mask])\n        else:\n            flat_img_masks.append([img, None])\n    df = pd.DataFrame(flat_img_masks, columns=['ImageId', 'EncodedPixels'])\n    return df","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_masks = get_test_masks_opening()\ndf_submission = create_submission_df(test_masks)\ndf_submission.to_csv('submission_wo_clf.csv', header=True, index=False)\n\nfrom IPython.display import FileLink\nFileLink('submission_wo_clf.csv')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Remove false positives with a pre-trained classifier"},{"metadata":{"trusted":true},"cell_type":"code","source":"# Get dataframe with label\nclf_df = pd.read_csv(os.path.join('../input', 'clfairbus/clf_256_test_preds.csv'))\n\nfor i, row in clf_df.iterrows():\n    if row['Label'] == 0:\n        df_submission.loc[df_submission['ImageId'] == row['ImageId'], 'EncodedPixels'] = None","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"df_submission.to_csv('submission_w_clf.csv', header=True, index=False)\nFileLink('submission_w_clf.csv')","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}