{"cells":[{"metadata":{},"cell_type":"markdown","source":"**In this notebook you'll find:**\n* An exploratory data analysis including an optional script which converts masks to bouding boxes (useful for getting some hints about ship sizes)\n* An implementation of various losses for binary segmentation (BCE, Jaccard, Dice and Focal losses)\n* An implementation of Dice and IoU metrics\n* Two U-Net models: original (same number of layers and channels) and an improved one (more layers)\n* When setting the booleans you can train the models from scratch or pulling a pre-trained U-Net from a github\n\n**What you need to know:**\n* This notebook borrows codes from other kernels - kind of local state-of-the-art - but some parts are mine. \n* It intends to show a U-Net training from scratch and how a basic finetuning can help to increase the performances\n* It won't provide the best performance ever but it's a good way to start\n\nPS: check my other kernels for a better accuracy","execution_count":null},{"metadata":{"trusted":true,"scrolled":false},"cell_type":"code","source":"import os\nimport sys\nimport random\nimport warnings\nimport cv2\nimport gc\n\nfrom tqdm import tqdm\nfrom datetime import datetime\nimport json\nfrom pathlib import Path\nimport pandas as pd\nimport numpy as np\nimport scipy as sp\nimport matplotlib.pyplot as plt\n%matplotlib inline\nfrom skimage.io import imread, imshow, concatenate_images\nfrom skimage import io, transform\nfrom skimage.measure import label, regionprops\nfrom sklearn.model_selection import train_test_split\nimport torch \nimport torchvision\nimport torchvision.transforms as transforms\nimport torchvision.models as models\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.optim as optim\nfrom torch.optim.lr_scheduler import StepLR\nfrom torch.utils.data import Dataset\nfrom PIL import Image \nimport math\n\nimport warnings\nwarnings.filterwarnings(\"ignore\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"**Global variables**","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"# Set paths\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\n# Booleans\nSHOW_PIXELS_DIST = False\nSHOW_SHIP_DIAG = False\nSHOW_IMG_LOADER = False\n\n# Training variables\nBATCH_SZ_TRAIN = 16\nBATCH_SZ_VALID = 4\nLR = 1e-4\nN_EPOCHS = 3\n\n# Define loss function\nLOSS = 'BCEWithDigits' # BCEWithDigits | FocalLossWithDigits | BCEDiceWithLogitsLoss | BCEJaccardWithLogitsLoss\n\n# Define model\nMODEL_SEG = 'UNET_RESNET34ImgNet' # UNET | IUNET | UNET_RESNET34ImgNet \nFREEZE_RESNET = False   # if UNET_RESNET34ImgNet\n\n# Fetch U-Net with a pre-trained RESNET34 encoder on imagenet\nif MODEL_SEG == 'UNET_RESNET34ImgNet':\n    !pip install git+https://github.com/qubvel/segmentation_models.pytorch > /dev/null 2>&1 # Install segmentations_models.pytorch, with no bash output.\n    import segmentation_models_pytorch as smp","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Utils","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"\n# Decode masks in CSV \n# Ref: https://www.kaggle.com/paulorzp/run-length-encode-and-decode\ndef rle_decode(mask_rle, shape=(768, 768)):\n    '''\n    mask_rle: run-length as string formated (start length)\n    shape: (height,width) of array to return \n    Returns numpy array, 1 - mask, 0 - background\n\n    '''\n    s = mask_rle.split()\n    starts, lengths = [np.asarray(x, dtype=int) for x in (s[0:][::2], s[1:][::2])]\n    starts -= 1\n    ends = starts + lengths\n    img = np.zeros(shape[0]*shape[1], dtype=np.uint8)\n    for lo, hi in zip(starts, ends):\n        img[lo:hi] = 1\n    return img.reshape(shape).T  # Needed to align to RLE direction\n\n# Convert CSV masks to image for a given image name \ndef maskcsv_to_img(masks, img_name):\n    masks_img = np.zeros((768,768))\n    masks_bin = masks.loc[masks['ImageId'] == img_name, 'EncodedPixels'].tolist()\n    for mask in masks_bin:\n        if isinstance(mask, str):\n            masks_img += rle_decode(mask)\n    return np.expand_dims(masks_img, -1)\n\n# Convert masks in a list to an image\ndef masks_as_image(in_mask_list):\n    # Take the individual ship masks and create a single mask array for all ships\n    all_masks = np.zeros((768, 768), dtype = np.int16)\n    for mask in in_mask_list:\n        if isinstance(mask, str):\n            all_masks += rle_decode(mask)\n    return np.expand_dims(all_masks, -1)\n    \n# Show an image and its corresponding mask\ndef imshow_mask(img, mask):\n    img = img.numpy().transpose((1, 2, 0))\n    mean = np.array([0.485, 0.456, 0.406])\n    std = np.array([0.229, 0.224, 0.225])\n    img = std * img + mean\n    img = np.clip(img, 0, 1)\n    \n    mask = mask.numpy().transpose((1, 2, 0))\n    mask = np.clip(mask, 0, 1)\n    \n    fig, axs = plt.subplots(1,2, figsize=(10,30))\n    axs[0].imshow(img)\n    axs[0].axis('off')\n    axs[1].imshow(mask)\n    axs[1].axis('off')\n\n    \ndef imshow_gt_out(img, mask_gt, mask_out):\n    img = img.numpy().transpose((1, 2, 0))\n    mean = np.array([0.485, 0.456, 0.406])\n    std = np.array([0.229, 0.224, 0.225])\n    img = std * img + mean\n    img = np.clip(img, 0, 1)\n\n    mask_gt = mask_gt.numpy().transpose((1, 2, 0))\n    mask_gt = np.clip(mask_gt, 0, 1)\n\n    mask_out = mask_out.numpy().transpose((1, 2, 0))\n    mask_out = np.clip(mask_out, 0, 1)\n\n    fig, axs = plt.subplots(1,3, figsize=(10,30))\n    axs[0].imshow(img)\n    axs[0].axis('off')\n    axs[0].set_title(\"Input image\")\n    axs[1].imshow(mask_gt)\n    axs[1].axis('off')\n    axs[1].set_title(\"Ground truth\")\n    axs[2].imshow(mask_out)\n    axs[2].axis('off')\n    axs[2].set_title(\"Model output\")\n    plt.subplots_adjust(wspace=0, hspace=0)\n\ndef imshow_overlay(img, mask, title=None):\n    \"\"\"Imshow for Tensor.\"\"\"\n    img = img.numpy().transpose((1, 2, 0))\n    mean = np.array([0.485, 0.456, 0.406])\n    std = np.array([0.229, 0.224, 0.225])\n    img = std * img + mean\n    img = np.clip(img, 0, 1)\n    mask = mask.numpy().transpose((1, 2, 0))\n    mask = np.clip(mask, 0, 1)\n    fig = plt.figure(figsize = (6,6))\n    plt.imshow(mask_overlay(img, mask))\n    if title is not None:\n        plt.title(title)\n    plt.pause(0.001) \n\ndef mask_overlay(image, mask, color=(0, 1, 0)):\n    \"\"\"\n    Helper function to visualize mask on the top of the image\n    \"\"\"\n    mask = np.dstack((mask, mask, mask)) * np.array(color)\n    weighted_sum = cv2.addWeighted(mask, 0.5, image, 0.5, 0.)\n    img = image.copy()\n    ind = mask[:, :, 1] > 0\n    img[ind] = weighted_sum[ind]    \n    return img\n    \n# This function transforms EncodedPixels into a list of pixels\n# Check our previous notebook for a detailed explanation:\n# https://www.kaggle.com/julian3833/2-understanding-and-plotting-rle-bounding-boxes\ndef rle_to_pixels(rle_code):\n    rle_code = [int(i) for i in rle_code.split()]\n    pixels = [(pixel_position % 768, pixel_position // 768) \n                 for start, length in list(zip(rle_code[0:-1:2], rle_code[1:-2:2])) \n                 for pixel_position in range(start, start + length)]\n    return pixels\n\ndef show_pixels_distribution(df):\n    \"\"\"\n    Prints the amount of ship and no-ship pixels in the df\n    \"\"\"\n    # Total images in the df\n    n_images = df['ImageId'].nunique() \n    \n    # Total pixels in the df\n    total_pixels = n_images * 768 * 768 \n\n    # Keep only rows with RLE boxes, transform them into list of pixels, sum the lengths of those lists\n    ship_pixels = df['EncodedPixels'].dropna().apply(rle_to_pixels).str.len().sum() \n\n    ratio = ship_pixels / total_pixels\n    print(f\"Ship: {round(ratio, 3)} ({ship_pixels})\")\n    print(f\"No ship: {round(1 - ratio, 3)} ({total_pixels - ship_pixels})\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"**PART 1 - EXPLORATORY DATA ANALYSIS**","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"# Read CSV as dataframe\nmasks = pd.read_csv(os.path.join(data_root, 'train_ship_segmentations_v2.csv'))\nprint('Total number of images (original): %d' % masks['ImageId'].value_counts().shape[0])\nif SHOW_PIXELS_DIST == True:\n    show_pixels_distribution(masks)\n    show_pixels_distribution(masks.dropna())\n\n# Create a dataframe with unique images id as indexes and number of ships and image sizes as new columns\nmasks = masks[~masks['ImageId'].isin(['6384c3e78.jpg'])] # remove corrupted file\nunique_img_ids = masks.groupby('ImageId').size().reset_index(name='counts')\nprint('Total number of images (after removing corrupted images): %d' % masks['ImageId'].value_counts().shape[0])\n\n# Plot some images with ships\nimg_wships = masks[~masks['EncodedPixels'].isna()].sample(9)\nfig, arr = plt.subplots(3,3, figsize=(10,10), constrained_layout=True)\nfor i, img in enumerate(img_wships['ImageId']):\n    r = int(i / 3)\n    c = i % 3\n    arr[r,c].imshow(imread(os.path.join(path_train, img)))\n    arr[r,c].axis('off')\nplt.show()\n\n# Plot some images without ships\nimg_woships = masks[masks['EncodedPixels'].isna()].sample(9)\nfig, arr = plt.subplots(3,3, figsize=(10,10), constrained_layout=True)\nfor i, img in enumerate(img_woships['ImageId']):\n    r = int(i / 3)\n    c = i % 3\n    arr[r,c].imshow(imread(os.path.join(path_train, img)))\n    arr[r,c].axis('off')\nplt.show()\n\n# Count number of ships per image\ndf_wships = masks.dropna()\ndf_wships = df_wships.groupby('ImageId').size().reset_index(name='counts')\ndf_woships = masks[masks['EncodedPixels'].isna()]\n\n# Make a plot\nplt.bar(['With ships','Without ships'], [len(df_wships), len(df_woships)])\nplt.ylabel('Number of images')\nplt.show()\n\nprint('Number of images with ships : %d | Number of images without ships : %d  (x%0.1f)' \\\n      % (df_wships.shape[0], df_woships.shape[0], df_woships.shape[0] / df_wships.shape[0]))\n## -> Unbalanced dataset\n\n# Remove images without ships to help getting a more balanced dataset\nmasks = masks.dropna()\ndf_woships = masks[masks['EncodedPixels'].isna()]\n\n# Make a plot\nplt.bar(['With ships','Without ships'], [len(df_wships), len(df_woships)])\nplt.ylabel('Number of images')\nplt.show()\n\nprint('Number of images with ships : %d | Number of images without ships : %d  (x%0.1f)' \\\n      % (df_wships.shape[0], df_woships.shape[0], df_woships.shape[0] / df_wships.shape[0]))\n## -> Balanced dataset\n\n# Plot histogram\nhist = df_wships.hist(bins=df_wships['counts'].max())\nplt.title(\"Histogram of ships count\")\nplt.xlabel(\"Number of ships\")\nplt.ylabel(\"Number of images\")\nplt.show(hist)\n\n# Plot images with 15 ships\ndf_w15ships = df_wships.loc[df_wships['counts'] == 5]\nlist_w15ships = df_w15ships.values.tolist()\n\nfig, axarr = plt.subplots(2, 2, figsize=(10, 10), constrained_layout=True)\nfor i in range(4):\n    rd_id = random.randrange(len(list_w15ships))\n    img_masks = masks.loc[masks['ImageId'] == str(list_w15ships[rd_id][0]), 'EncodedPixels'].tolist()\n\n    # Take the individual ship masks and create a single mask array for all ships\n    all_masks = np.zeros((768, 768))\n    for mask in img_masks:\n        all_masks += rle_decode(mask)\n        \n    r = int(i / 2)\n    c = i % 2\n\n    axarr[r][c].imshow(imread(os.path.join(path_train, list_w15ships[rd_id][0])))\n    axarr[r][c].imshow(all_masks, alpha=0.3)\n    axarr[r][c].axis('off')\n                    \nplt.show()\n\n# Split dataset into training and validation sets\n# statritify : same histograms of numbe of ships\nunique_img_ids = masks.groupby('ImageId').size().reset_index(name='counts')\n\ntrain_ids, val_ids = train_test_split(unique_img_ids, test_size=0.05, stratify=unique_img_ids['counts'], random_state=42)\ntrain_df = pd.merge(masks, train_ids)\nvalid_df = pd.merge(masks, val_ids)\n\ntrain_df['counts'] = train_df.apply(lambda c_row: c_row['counts'] if isinstance(c_row['EncodedPixels'], str) else 0, 1)\nvalid_df['counts'] = valid_df.apply(lambda c_row: c_row['counts'] if isinstance(c_row['EncodedPixels'], str) else 0, 1)\n\nprint('Number of training images : %d' % train_df['ImageId'].value_counts().shape[0])\ntrain_df['counts'].hist(bins=train_df['counts'].max())\nplt.title(\"Histogram of ships count (training)\")\nplt.xlabel(\"Number of ships\")\nplt.ylabel(\"Number of images\")\nplt.show()\n\nprint('Number of validation images : %d' % valid_df['ImageId'].value_counts().shape[0])\nvalid_df['counts'].hist(bins=valid_df['counts'].max())\nplt.title(\"Histogram of ships count (validation)\")\nplt.xlabel(\"Number of ships\")\nplt.ylabel(\"Number of images\")\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"*[Optional] : convert segmentation masks into bouding boxes (without rotation) and compute ship sizes (diagonals) -> takes time*","execution_count":null},{"metadata":{"trusted":true,"scrolled":false},"cell_type":"code","source":"if SHOW_SHIP_DIAG == True:\n# Transform masks to bounding boxes\n# Analisys of ship sizes through the dataset\n\n    bboxes = []\n    bboxes_dict = {}\n\n    # Compute bouding boxes coordinates\n    for img_id in tqdm(df_wships['ImageId']):\n        bboxes = []\n\n        # Get binary mask\n        masks_val = masks.loc[masks['ImageId'] == str(img_id), 'EncodedPixels'].tolist()\n\n        # Take the individual ship masks and create a single mask array for all ships\n        bin_mask = np.zeros((768, 768))\n        for mask in masks_val:\n            bin_mask += rle_decode(mask)\n\n        # Extract bounding boxes\n        lbl = label(bin_mask)\n        props = regionprops(lbl)\n\n        for prop in props:\n            bboxes.append(prop.bbox)\n\n        bboxes_dict[img_id] = bboxes.copy()\n    \n    # Plot some images with bouding boxes\n    fig, axarr = plt.subplots(2, 2, figsize = (15, 15), constrained_layout=True)\n    for i in range(4):\n        img_id = df_wships.loc[i*4, 'ImageId']\n        img = imread(os.path.join(path_train, str(img_id)))\n\n        bboxs = bboxes_dict[img_id]\n        for bbox in bboxs:\n            cv2.rectangle(img, (bbox[1], bbox[0]), (bbox[3], bbox[2]), (255, 0, 0), 2)\n\n        r = int(i / 2)\n        c = i % 2\n        axarr[r][c].imshow(img)\n        axarr[r][c].axis('off')\n\n    plt.show()\n\n    # Compute diagonal of each ships and plot histogram\n    diag = []\n    for i in bboxes_dict:\n        for j in bboxes_dict[i]:\n            diag.append(int(np.sqrt((j[0] - j[2]) ** 2 + (j[1] - j[3]) ** 2)))\n\n    df_diag = pd.DataFrame(diag, columns =['Diagonal size'])\n\n    axes = df_diag.hist()\n    plt.title(\"Histogram of ship diagonal sizes\")\n    plt.xlabel(\"Number of ships\")\n    plt.ylabel(\"Diagonal in pixels\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"**PART 2 - DATASET, LOSSES, METRICS **","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"Dataset","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"class AirbusDataset(Dataset):\n    def __init__(self, in_df, transform=None, mode='train'):\n        grp = list(in_df.groupby('ImageId'))\n        self.image_ids =  [_id for _id, _ in grp] \n        self.image_masks = [m['EncodedPixels'].values for _,m in grp]\n        self.transform = transform\n        self.mode = mode\n        self.img_transform = transforms.Compose([\n        transforms.ToTensor(),\n        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])  # use mean and std from ImageNet \n\n    def __len__(self):\n        return len(self.image_ids)\n               \n    def __getitem__(self, idx):\n        img_file_name = self.image_ids[idx]\n        if (self.mode == 'train') | (self.mode == 'validation'):\n            rgb_path = os.path.join(path_train, img_file_name)\n        else:\n            rgb_path = os.path.join(path_test, img_file_name)\n        img = imread(rgb_path)\n        mask = masks_as_image(self.image_masks[idx])\n        \n        if self.transform is not None: \n            img, mask = self.transform(img, mask)\n            \n        if (self.mode == 'train') | (self.mode == 'validation'):\n            return self.img_transform(img), torch.from_numpy(np.moveaxis(mask, -1, 0)).float()  \n        else:\n            return self.img_transform(img), str(img_file_name)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Data augmentation","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"# Implementation from  https://github.com/ternaus/robot-surgery-segmentation\ndef clip(img, dtype, maxval):\n    return np.clip(img, 0, maxval).astype(dtype)\n\nclass DualCompose:\n    def __init__(self, transforms):\n        self.transforms = transforms\n\n    def __call__(self, x, mask=None):\n        for t in self.transforms:\n            x, mask = t(x, mask)\n        return x, mask\n\nclass ImageOnly:\n    def __init__(self, trans):\n        self.trans = trans\n\n    def __call__(self, x, mask=None):\n        return self.trans(x), mask\n\n\nclass VerticalFlip:\n    def __init__(self, prob=0.5):\n        self.prob = prob\n\n    def __call__(self, img, mask=None):\n        if random.random() < self.prob:\n            img = cv2.flip(img, 0)\n            if mask is not None:\n                mask = cv2.flip(mask, 0)\n        return img, mask\n\n\nclass HorizontalFlip:\n    def __init__(self, prob=0.5):\n        self.prob = prob\n\n    def __call__(self, img, mask=None):\n        if random.random() < self.prob:\n            img = cv2.flip(img, 1)\n            if mask is not None:\n                mask = cv2.flip(mask, 1)\n        return img, mask\n\n\nclass RandomFlip:\n    def __init__(self, prob=0.5):\n        self.prob = prob\n\n    def __call__(self, img, mask=None):\n        if random.random() < self.prob:\n            d = random.randint(-1, 1)\n            img = cv2.flip(img, d)\n            if mask is not None:\n                mask = cv2.flip(mask, d)\n        return img, mask\n\n\nclass Rotate:\n    def __init__(self, limit=90, prob=0.5):\n        self.prob = prob\n        self.limit = limit\n\n    def __call__(self, img, mask=None):\n        if random.random() < self.prob:\n            angle = random.uniform(-self.limit, self.limit)\n\n            height, width = img.shape[0:2]\n            mat = cv2.getRotationMatrix2D((width / 2, height / 2), angle, 1.0)\n            img = cv2.warpAffine(img, mat, (height, width),\n                                 flags=cv2.INTER_LINEAR,\n                                 borderMode=cv2.BORDER_REFLECT_101)\n            if mask is not None:\n                mask = cv2.warpAffine(mask, mat, (height, width),\n                                      flags=cv2.INTER_LINEAR,\n                                      borderMode=cv2.BORDER_REFLECT_101)\n\n        return img, mask\n\nclass RandomCrop:\n    def __init__(self, size):\n        self.h = size[0]\n        self.w = size[1]\n\n    def __call__(self, img, mask=None):\n        height, width, _ = img.shape\n\n        h_start = np.random.randint(0, height - self.h)\n        w_start = np.random.randint(0, width - self.w)\n\n        img = img[h_start: h_start + self.h, w_start: w_start + self.w,:]\n\n        assert img.shape[0] == self.h\n        assert img.shape[1] == self.w\n\n        if mask is not None:\n            if mask.ndim == 2:\n                mask = np.expand_dims(mask, axis=2)\n            mask = mask[h_start: h_start + self.h, w_start: w_start + self.w,:]\n\n        return img, mask\n\nclass CenterCrop:\n    def __init__(self, size):\n        self.height = size[0]\n        self.width = size[1]\n\n    def __call__(self, img, mask=None):\n        h, w, c = img.shape\n        dy = (h - self.height) // 2\n        dx = (w - self.width) // 2\n        y1 = dy\n        y2 = y1 + self.height\n        x1 = dx\n        x2 = x1 + self.width\n        img = img[y1:y2, x1:x2,:]\n        if mask is not None:\n            if mask.ndim == 2:\n                mask = np.expand_dims(mask, axis=2)\n            mask = mask[y1:y2, x1:x2,:]\n\n        return img, mask","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Metrics","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"def compute_metrics(pred, true, batch_size=16, threshold=0.5):\n    pred = pred.view(batch_size, -1)\n    true = true.view(batch_size, -1)\n    \n    pred = (pred > threshold).float()\n    true = (true > threshold).float()\n    \n    pred_sum = pred.sum(-1)\n    true_sum = true.sum(-1)\n    \n    neg_index = torch.nonzero(true_sum == 0)\n    pos_index = torch.nonzero(true_sum >= 1)\n    \n    dice_neg = (pred_sum == 0).float()\n    dice_pos = 2 * ((pred * true).sum(-1)) / ((pred + true).sum(-1))\n    \n    dice_neg = dice_neg[neg_index]\n    dice_pos = dice_pos[pos_index]\n    \n    dice = torch.cat([dice_pos, dice_neg])\n    jaccard = dice / (2 - dice)\n    \n    return dice, jaccard\n    \nclass metrics:\n    def __init__(self, batch_size=16, threshold=0.5):\n        self.threshold = threshold\n        self.batchsize = batch_size\n        self.dice = []\n        self.jaccard = []\n    def collect(self, pred, true):\n        pred = torch.sigmoid(pred)\n        dice, jaccard = compute_metrics(pred, true, batch_size=self.batchsize, threshold=self.threshold)\n        self.dice.extend(dice)\n        self.jaccard.extend(jaccard)\n    def get(self):\n        dice = np.nanmean(self.dice)\n        jaccard = np.nanmean(self.jaccard)\n        return dice, jaccard","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Losses","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"class BCEJaccardWithLogitsLoss(nn.Module):\n    def __init__(self, jaccard_weight=1, smooth=1):\n        super().__init__()\n        self.bce = nn.BCEWithLogitsLoss()\n        self.jaccard_weight = jaccard_weight\n        self.smooth = smooth\n\n    def forward(self, outputs, targets):\n        if outputs.size() != targets.size():\n            raise ValueError(\"size mismatch, {} != {}\".format(outputs.size(), targets.size()))\n            \n        loss = self.bce(outputs, targets)\n\n        if self.jaccard_weight:\n            targets = (targets == 1.0).float()\n            targets = targets.view(-1)\n            outputs = torch.sigmoid(outputs)\n            outputs = outputs.view(-1)\n\n            intersection = (targets * outputs).sum()\n            union = outputs.sum() + targets.sum() - intersection\n\n            loss -= self.jaccard_weight * torch.log((intersection + self.smooth ) / (union + self.smooth )) # try with 1-dice\n        return loss\n\nclass BCEDiceWithLogitsLoss(nn.Module):\n    def __init__(self, dice_weight=1, smooth=1):\n        super().__init__()\n        self.bce = nn.BCEWithLogitsLoss()\n        self.dice_weight = dice_weight\n        self.smooth = smooth\n        \n    def __call__(self, outputs, targets):\n        if outputs.size() != targets.size():\n            raise ValueError(\"size mismatch, {} != {}\".format(outputs.size(), targets.size()))\n            \n        loss = self.bce(outputs, targets)\n\n        targets = (targets == 1.0).float()\n        targets = targets.view(-1)\n        outputs = F.sigmoid(outputs)\n        outputs = outputs.view(-1)\n\n        intersection = (outputs * targets).sum()\n        dice = 2.0 * (intersection + self.smooth)  / (targets.sum() + outputs.sum() + self.smooth)\n        \n        loss -= self.dice_weight * torch.log(dice) # try with 1- dice\n\n        return loss\n    \nclass FocalWithLogitsLoss(nn.Module):\n    def __init__(self, alpha=0.25, gamma=2):\n        super().__init__()\n        self.bce = nn.BCEWithLogitsLoss()\n        self.alpha = alpha\n        self.gamma = gamma\n        \n    def __call__(self, outputs, targets):\n        if outputs.size() != targets.size():\n            raise ValueError(\"size mismatch, {} != {}\".format(outputs.size(), targets.size()))\n            \n        loss = self.bce(outputs, targets)\n\n        targets = (targets == 1.0).float()\n        targets = targets.view(-1)\n        outputs = torch.sigmoid(outputs)\n        outputs = outputs.view(-1)\n        outputs = torch.where(targets == 1, outputs, 1 - outputs)\n\n        focal = self.alpha * (1 - outputs) ** (self.gamma)\n        loss *= focal.mean()\n\n        return loss\n\ndef dice_loss(input, target):\n    input = torch.sigmoid(input)\n    smooth = 1.0\n\n    iflat = input.view(-1)\n    tflat = target.view(-1)\n    intersection = (iflat * tflat).sum()\n    \n    return ((2.0 * intersection + smooth) / (iflat.sum() + tflat.sum() + smooth))\n\nclass FocalLoss(nn.Module):\n    def __init__(self, gamma):\n        super().__init__()\n        self.gamma = gamma\n        \n    def forward(self, input, target):\n        if not (target.size() == input.size()):\n            raise ValueError(\"Target size ({}) must be the same as input size ({})\"\n                             .format(target.size(), input.size()))\n\n        max_val = (-input).clamp(min=0)\n        loss = input - input * target + max_val + \\\n            ((-max_val).exp() + (-input - max_val).exp()).log()\n\n        invprobs = F.logsigmoid(-input * (target * 2.0 - 1.0))\n        loss = (invprobs * self.gamma).exp() * loss\n        \n        return loss.mean()\n\nclass MixedLoss(nn.Module):\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        loss = self.alpha*self.focal(input, target) - torch.log(dice_loss(input, target))\n        return loss.mean()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Models","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"\"\"\"Improved U-Net\"\"\" \n# Implementation from https://github.com/timctho/unet-pytorch/\nclass IUNet_down_block(torch.nn.Module):\n    def __init__(self, input_channel, output_channel, down_size):\n        super(IUNet_down_block, self).__init__()\n        self.conv1 = torch.nn.Conv2d(input_channel, output_channel, 3, padding=1)\n        self.bn1 = torch.nn.BatchNorm2d(output_channel)\n        self.conv2 = torch.nn.Conv2d(output_channel, output_channel, 3, padding=1)\n        self.bn2 = torch.nn.BatchNorm2d(output_channel)\n        self.conv3 = torch.nn.Conv2d(output_channel, output_channel, 3, padding=1)\n        self.bn3 = torch.nn.BatchNorm2d(output_channel)\n        self.max_pool = torch.nn.MaxPool2d(2, 2)\n        self.relu = torch.nn.ReLU()\n        self.down_size = down_size\n\n    def forward(self, x):\n        if self.down_size:\n            x = self.max_pool(x)\n        x = self.relu(self.bn1(self.conv1(x)))\n        x = self.relu(self.bn2(self.conv2(x)))\n        x = self.relu(self.bn3(self.conv3(x)))\n        return x\n\nclass IUNet_up_block(torch.nn.Module):\n    def __init__(self, prev_channel, input_channel, output_channel):\n        super(IUNet_up_block, self).__init__()\n        self.up_sampling = torch.nn.Upsample(scale_factor=2, mode='bilinear')\n        self.conv1 = torch.nn.Conv2d(prev_channel + input_channel, output_channel, 3, padding=1)\n        self.bn1 = torch.nn.BatchNorm2d(output_channel)\n        self.conv2 = torch.nn.Conv2d(output_channel, output_channel, 3, padding=1)\n        self.bn2 = torch.nn.BatchNorm2d(output_channel)\n        self.conv3 = torch.nn.Conv2d(output_channel, output_channel, 3, padding=1)\n        self.bn3 = torch.nn.BatchNorm2d(output_channel)\n        self.relu = torch.nn.ReLU()\n\n    def forward(self, prev_feature_map, x):\n        x = self.up_sampling(x)\n        x = torch.cat((x, prev_feature_map), dim=1)\n        x = self.relu(self.bn1(self.conv1(x)))\n        x = self.relu(self.bn2(self.conv2(x)))\n        x = self.relu(self.bn3(self.conv3(x)))\n        return x\n\n\nclass IUNet(torch.nn.Module):\n    def __init__(self):\n        super(IUNet, self).__init__()\n\n        self.down_block1 = IUNet_down_block(3, 16, False)\n        self.down_block2 = IUNet_down_block(16, 32, True)\n        self.down_block3 = IUNet_down_block(32, 64, True)\n        self.down_block4 = IUNet_down_block(64, 128, True)\n        self.down_block5 = IUNet_down_block(128, 256, True)\n        self.down_block6 = IUNet_down_block(256, 512, True)\n        self.down_block7 = IUNet_down_block(512, 1024, True)\n\n        self.mid_conv1 = torch.nn.Conv2d(1024, 1024, 3, padding=1)\n        self.bn1 = torch.nn.BatchNorm2d(1024)\n        self.mid_conv2 = torch.nn.Conv2d(1024, 1024, 3, padding=1)\n        self.bn2 = torch.nn.BatchNorm2d(1024)\n        self.mid_conv3 = torch.nn.Conv2d(1024, 1024, 3, padding=1)\n        self.bn3 = torch.nn.BatchNorm2d(1024)\n\n        self.up_block1 = IUNet_up_block(512, 1024, 512)\n        self.up_block2 = IUNet_up_block(256, 512, 256)\n        self.up_block3 = IUNet_up_block(128, 256, 128)\n        self.up_block4 = IUNet_up_block(64, 128, 64)\n        self.up_block5 = IUNet_up_block(32, 64, 32)\n        self.up_block6 = IUNet_up_block(16, 32, 16)\n\n        self.last_conv1 = torch.nn.Conv2d(16, 16, 3, padding=1)\n        self.last_bn = torch.nn.BatchNorm2d(16)\n        self.last_conv2 = torch.nn.Conv2d(16, 1, 1, padding=0)\n        self.relu = torch.nn.ReLU()\n\n    def forward(self, x):\n        self.x1 = self.down_block1(x)\n        self.x2 = self.down_block2(self.x1)\n        self.x3 = self.down_block3(self.x2)\n        self.x4 = self.down_block4(self.x3)\n        self.x5 = self.down_block5(self.x4)\n        self.x6 = self.down_block6(self.x5)\n        self.x7 = self.down_block7(self.x6)\n        self.x7 = self.relu(self.bn1(self.mid_conv1(self.x7)))\n        self.x7 = self.relu(self.bn2(self.mid_conv2(self.x7)))\n        self.x7 = self.relu(self.bn3(self.mid_conv3(self.x7)))\n        x = self.up_block1(self.x6, self.x7)\n        x = self.up_block2(self.x5, x)\n        x = self.up_block3(self.x4, x)\n        x = self.up_block4(self.x3, x)\n        x = self.up_block5(self.x2, x)\n        x = self.up_block6(self.x1, x)\n        x = self.relu(self.last_bn(self.last_conv1(x)))\n        x = self.last_conv2(x)\n        return x\n        \n        \n\"\"\"Original U-Net\"\"\"\nclass UNet_down_block(torch.nn.Module):\n    def __init__(self, input_channel, output_channel, down_size):\n        super(UNet_down_block, self).__init__()\n        self.conv1 = torch.nn.Conv2d(input_channel, output_channel, 3, padding=1)\n        self.bn1 = torch.nn.BatchNorm2d(output_channel)\n        self.conv2 = torch.nn.Conv2d(output_channel, output_channel, 3, padding=1)\n        self.bn2 = torch.nn.BatchNorm2d(output_channel)\n        self.max_pool = torch.nn.MaxPool2d(2, 2)\n        self.relu = torch.nn.ReLU()\n        self.down_size = down_size\n\n    def forward(self, x):\n        if self.down_size:\n            x = self.max_pool(x)\n        x = self.relu(self.bn1(self.conv1(x)))\n        x = self.relu(self.bn2(self.conv2(x)))\n        return x\n\nclass UNet_up_block(torch.nn.Module):\n    def __init__(self, prev_channel, input_channel, output_channel):\n        super(UNet_up_block, self).__init__()\n        self.up_sampling = torch.nn.Upsample(scale_factor=2, mode='bilinear')\n        self.conv1 = torch.nn.Conv2d(prev_channel + input_channel, output_channel, 3, padding=1)\n        self.bn1 = torch.nn.BatchNorm2d(output_channel)\n        self.conv2 = torch.nn.Conv2d(output_channel, output_channel, 3, padding=1)\n        self.bn2 = torch.nn.BatchNorm2d(output_channel)\n        self.relu = torch.nn.ReLU()\n\n    def forward(self, prev_feature_map, x):\n        x = self.up_sampling(x)\n        x = torch.cat((x, prev_feature_map), dim=1)\n        x = self.relu(self.bn1(self.conv1(x)))\n        x = self.relu(self.bn2(self.conv2(x)))\n        return x\n\n\nclass UNet(torch.nn.Module):\n    def __init__(self):\n        super(UNet, self).__init__()\n\n        self.down_block1 = UNet_down_block(3, 64, False)\n        self.down_block2 = UNet_down_block(64, 128, True)\n        self.down_block3 = UNet_down_block(128, 256, True)\n        self.down_block4 = UNet_down_block(256, 512, True)\n        self.down_block5 = UNet_down_block(512, 1024, True)\n\n        self.up_block1 = UNet_up_block(512, 1024, 512)\n        self.up_block2 = UNet_up_block(256, 512, 256)\n        self.up_block3 = UNet_up_block(128, 256, 128)\n        self.up_block4 = UNet_up_block(64, 128, 64)\n\n        self.last_conv = torch.nn.Conv2d(64, 1, 1, padding=0)\n\n    def forward(self, x):\n        self.x1 = self.down_block1(x)\n        self.x2 = self.down_block2(self.x1)\n        self.x3 = self.down_block3(self.x2)\n        self.x4 = self.down_block4(self.x3)\n        self.x5 = self.down_block5(self.x4)\n        x = self.up_block1(self.x4, self.x5)\n        x = self.up_block2(self.x3, x)\n        x = self.up_block3(self.x2, x)\n        x = self.up_block4(self.x1, x)\n        x = self.last_conv(x)\n        return x","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# PART 3 - CREATE TRAIN AND VALIDATION LOOPS","execution_count":null},{"metadata":{"trusted":true,"scrolled":false},"cell_type":"code","source":"# main train routine\n# Implementation from  https://github.com/ternaus/robot-surgery-segmentation\ndef train(lr, model, criterion, train_loader, valid_loader, init_optimizer, train_batch_sz=16, valid_batch_sz=4, n_epochs=1, fold=1):\n    \n    model_path = Path('model_{fold}.pt'.format(fold=fold))\n    if model_path.exists():\n        state = torch.load(str(model_path))\n        epoch = state['epoch']\n        step = state['step']\n        model.load_state_dict(state['model'])\n        print('Restored model, epoch {}, step {:,}'.format(epoch, step))\n    else:\n        epoch = 1\n        step = 0\n\n    save = lambda ep: torch.save({\n        'model': model.state_dict(),\n        'epoch': ep,\n        'step': step,\n    }, str(model_path))\n\n    report_each = 50\n    log = open('train_{fold}.log'.format(fold=fold),'at', encoding='utf8')\n    \n    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n    model = model.to(device)\n    optimizer = init_optimizer(lr)\n\n    for epoch in range(epoch, n_epochs + 1):\n        model.train()\n        random.seed()\n        tq = tqdm(total=len(train_loader) *  train_batch_sz)\n        tq.set_description('Epoch {}, lr {}'.format(epoch, lr))\n        losses = []\n        valid_metrics = metrics(batch_size=valid_batch_sz)  # for validation\n        tl = train_loader\n        try:\n            mean_loss = 0\n            for i, (inputs, targets) in enumerate(tl):\n                inputs, targets = inputs.to(device), targets.to(device)\n                optimizer.zero_grad()\n                outputs = model.forward(inputs)\n                loss = criterion(outputs, targets)\n                batch_size = inputs.size(0)\n                loss.backward()\n                optimizer.step()\n                step += 1\n                tq.update(batch_size)\n                losses.append(loss.item())\n                mean_loss = np.mean(losses[-report_each:])\n                tq.set_postfix(loss='{:.5f}'.format(mean_loss))\n                if i and i % report_each == 0:\n                    write_event(log, step, loss=mean_loss)\n            write_event(log, step, loss=mean_loss)\n            tq.close()\n            save(epoch + 1)\n            \n            # Validation\n            comb_loss_metrics = validation(model, criterion, valid_loader, valid_metrics)\n            write_event(log, step, **comb_loss_metrics)\n\n        except KeyboardInterrupt:\n            tq.close()\n            print('Ctrl+C, saving snapshot')\n            save(epoch)\n            print('done.')\n            return\n        \ndef validation(model: nn.Module, criterion, valid_loader, metrics):\n    print(\"Validation\")\n    \n    losses = []\n    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n    model.eval()\n    \n    for inputs, targets in valid_loader:\n        inputs, targets = inputs.to(device), targets.to(device)\n        outputs = model.forward(inputs)\n        loss = criterion(outputs, targets)\n        losses.append(loss.item())\n        metrics.collect(outputs.detach().cpu(), targets.detach().cpu()) # get metrics \n    \n    valid_loss = np.mean(losses)  # float\n    valid_dice, valid_jaccard = metrics.get() # float\n\n    print('Valid loss: {:.5f}, Jaccard: {:.5f}, Dice: {:.5f}'.format(valid_loss, valid_jaccard, valid_dice))\n    comb_loss_metrics = {'valid_loss': valid_loss, 'jaccard': valid_jaccard.item(), 'dice': valid_dice.item()}\n    \n    return comb_loss_metrics\n\ndef write_event(log, step: int, **data):\n    data['step'] = step\n    data['dt'] = datetime.now().isoformat()\n    log.write(json.dumps(data, sort_keys=True))\n    log.write('\\n')\n    log.flush()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# PART 4 - TRAINING","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"# Data augmentation\ntrain_transform = DualCompose([HorizontalFlip(), VerticalFlip(), RandomCrop((256,256,3))])\nval_transform = DualCompose([CenterCrop((512,512,3))])\n\n# Initialize dataset\ntrain_dataset = AirbusDataset(train_df, transform=train_transform, mode='train')\nval_dataset = AirbusDataset(valid_df, transform=val_transform, mode='validation')\n\nprint('Train samples : %d | Validation samples : %d' % (len(train_dataset), len(val_dataset)))\n\n# Get loaders\ntrain_loader = torch.utils.data.DataLoader(train_dataset, batch_size=BATCH_SZ_TRAIN, shuffle=True, num_workers=0)\nval_loader = torch.utils.data.DataLoader(val_dataset, batch_size=BATCH_SZ_VALID, shuffle=True, num_workers=0)\n\nif SHOW_IMG_LOADER == True:\n    # Display some images from loader\n    images, mask = next(iter(train_loader))\n    imshow_mask(torchvision.utils.make_grid(images, nrow=1), torchvision.utils.make_grid(mask, nrow=1))\n    plt.show()\n    \n# Train\nrun_id = 1\n\nif MODEL_SEG == 'IUNET':\n    model = IUNet()\nelif MODEL_SEG == 'UNET':\n    model = UNet()\nelif MODEL_SEG == 'UNET_RESNET34ImgNet':\n    model = smp.Unet(\"resnet34\", encoder_weights=\"imagenet\", activation=None)\n    if FREEZE_RESNET == True:\n        for name, p in model.named_parameters():\n            if \"encoder\" in name:\n                p.requires_grad = False\nelse:\n    raise NameError(\"model not supported\")\n    \nif LOSS == 'BCEWithDigits':\n    criterion = nn.BCEWithLogitsLoss()\nelif LOSS == 'FocalLossWithDigits':\n    criterion = MixedLoss(10, 2)\nelif LOSS == 'BCEDiceWithLogitsLoss':\n    criterion = BCEDiceWithLogitsLoss()\nelif LOSS == 'BCEJaccardWithLogitsLoss':\n    criterion = BCEJaccardWithLogitsLoss()\nelse:\n    raise NameError(\"loss not supported\")\n\ntrain(init_optimizer=lambda lr: optim.Adam(model.parameters(), lr=lr),\n        lr = LR,\n        n_epochs = N_EPOCHS,\n        model=model,\n        criterion=criterion,\n        train_loader=train_loader,\n        valid_loader=val_loader,\n        train_batch_sz= BATCH_SZ_TRAIN,\n        valid_batch_sz=BATCH_SZ_VALID,\n        fold=run_id\n        )\n\n\n# Plot losses\nlog_file = 'train_{fold}.log'.format(fold=run_id)\nlogs = pd.read_json(log_file, lines=True)\n\nplt.figure(figsize=(26,6))\nplt.subplot(1, 2, 1)\nplt.plot(logs.step[logs.loss.notnull()],\n            logs.loss[logs.loss.notnull()],\n            label=\"on training set\")\n\nplt.plot(logs.step[logs.valid_loss.notnull()],\n            logs.valid_loss[logs.valid_loss.notnull()],\n            label = \"on validation set\")\n         \nplt.xlabel('step')\nplt.legend(loc='center left')\nplt.tight_layout()\nplt.show();\n\n\n# Model inference\nmodel_path ='model_{fold}.pt'.format(fold=run_id)\nstate = torch.load(str(model_path))\nstate = {key.replace('module.', ''): value for key, value in state['model'].items()}\nmodel.load_state_dict(state)\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nmodel = model.to(device)\nmodel.eval()\n\nval_dataset = AirbusDataset(valid_df, transform=val_transform)\nval_loader = torch.utils.data.DataLoader(val_dataset, batch_size=8, shuffle=True, num_workers=0)\n\n# Display some images from loader\nimages, gt = next(iter(val_loader))\ngt = gt.data.cpu()\nimages = images.to(device)\nout = model.forward(images)\nout = ((out > 0).float()) * 255\nimages = images.data.cpu()\nout = out.data.cpu()\nimshow_gt_out(torchvision.utils.make_grid(images, nrow=1),torchvision.utils.make_grid(gt, nrow=1), torchvision.utils.make_grid(out, nrow=1))\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# PART 5 - SUBMISSION","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"from skimage.morphology import binary_opening, disk\n# ref: https://www.kaggle.com/paulorzp/run-length-encode-and-decode\ndef rle_encode(img):\n    '''\n    img: numpy array, 1 - mask, 0 - background\n    Returns run length as string formated\n    '''\n    pixels = img.flatten()\n    pixels = np.concatenate([[0], pixels, [0]])\n    runs = np.where(pixels[1:] != pixels[:-1])[0] + 1\n    runs[1::2] -= runs[::2]\n    return ' '.join(str(x) for x in runs)\n\ndef multi_rle_encode(img):\n    labels = label(img)\n    return [rle_encode(labels==k) for k in np.unique(labels[labels>0])]\n\nlist_img_test = os.listdir(path_test)\nprint(len(list_img_test), 'test images found')\n\n# Create dataframe\ntest_df = pd.DataFrame({'ImageId': list_img_test, 'EncodedPixels': None})\nloader = torch.utils.data.DataLoader(dataset=AirbusDataset(test_df, transform=None, mode='test'), shuffle=False, batch_size=2, num_workers=0)\n    \nout_pred_rows = []\nfor batch_num, (inputs, paths) in enumerate(tqdm(loader, desc='Test')):\n    inputs = inputs.to(device)\n    outputs = model(inputs)\n    for i, image_name in enumerate(paths):\n        mask = F.sigmoid(outputs[i,0]).data.detach().cpu().numpy()\n        cur_seg = binary_opening(mask>0.5, disk(2))\n        cur_rles = multi_rle_encode(cur_seg)\n        if len(cur_rles)>0:\n            for c_rle in cur_rles:\n                out_pred_rows += [{'ImageId': image_name, 'EncodedPixels': c_rle}]\n        else:\n            out_pred_rows += [{'ImageId': image_name, 'EncodedPixels': None}]\n        \nsubmission_df = pd.DataFrame(out_pred_rows)[['ImageId', 'EncodedPixels']]\nsubmission_df.to_csv('submission.csv', index=False)\nsubmission_df.sample(10)","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}