{

 "cells": [

  {

   "cell_type": "markdown",

   "metadata": {},

   "source": [

    "# GI Tract Segmentation Competition"

   ]

  },

  {

   "cell_type": "markdown",

   "metadata": {},

   "source": [

    "# Load And Prepare"

   ]

  },

  {

   "cell_type": "markdown",

   "metadata": {},

   "source": [

    "## Imports"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": 1,

   "metadata": {

    "ExecuteTime": {

     "end_time": "2022-06-28T03:46:47.438209Z",

     "start_time": "2022-06-28T03:46:44.722879Z"

    }

   },

   "outputs": [],

   "source": [

    "from fastai.vision.all import *\n",

    "import matplotlib.patches as mpatches\n",

    "import albumentations as A\n",

    "import cv2\n",

    "import pynvml\n",

    "from scipy.spatial.distance import directed_hausdorff\n",

    "import timm\n",

    "import segmentation_models_pytorch as smp\n",

    "\n",

    "\n",

    "# !cp kaggle.json /home/kgeorgio/.kaggle\n",

    "# !kaggle competitions download -c uw-madison-gi-tract-image-segmentation\n"

   ]

  },

  {

   "cell_type": "markdown",

   "metadata": {},

   "source": [

    "## Load CSV"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": 2,

   "metadata": {

    "ExecuteTime": {

     "end_time": "2022-06-28T03:46:47.672193Z",

     "start_time": "2022-06-28T03:46:47.439755Z"

    },

    "pycharm": {

     "name": "#%%\n"

    }

   },

   "outputs": [],

   "source": [

    "train_df = pd.read_csv('dataset/train.csv', low_memory=False)\n",

    "train_df = train_df.pivot(index='id', columns='class', values='segmentation').reset_index()"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": 3,

   "metadata": {

    "ExecuteTime": {

     "end_time": "2022-06-28T03:46:47.739865Z",

     "start_time": "2022-06-28T03:46:47.730425Z"

    },

    "pycharm": {

     "name": "#%%\n"

    }

   },

   "outputs": [

    {

     "data": {

      "text/html": [

       "<div>\n",

       "<style scoped>\n",

       "    .dataframe tbody tr th:only-of-type {\n",

       "        vertical-align: middle;\n",

       "    }\n",

       "\n",

       "    .dataframe tbody tr th {\n",

       "        vertical-align: top;\n",

       "    }\n",

       "\n",

       "    .dataframe thead th {\n",

       "        text-align: right;\n",

       "    }\n",

       "</style>\n",

       "<table border=\"1\" class=\"dataframe\">\n",

       "  <thead>\n",

       "    <tr style=\"text-align: right;\">\n",

       "      <th>class</th>\n",

       "      <th>id</th>\n",

       "      <th>large_bowel</th>\n",

       "      <th>small_bowel</th>\n",

       "      <th>stomach</th>\n",

       "    </tr>\n",

       "  </thead>\n",

       "  <tbody>\n",

       "    <tr>\n",

       "      <th>0</th>\n",

       "      <td>case101_day20_slice_0001</td>\n",

       "      <td>NaN</td>\n",

       "      <td>NaN</td>\n",

       "      <td>NaN</td>\n",

       "    </tr>\n",

       "    <tr>\n",

       "      <th>1</th>\n",

       "      <td>case101_day20_slice_0002</td>\n",

       "      <td>NaN</td>\n",

       "      <td>NaN</td>\n",

       "      <td>NaN</td>\n",

       "    </tr>\n",

       "    <tr>\n",

       "      <th>2</th>\n",

       "      <td>case101_day20_slice_0003</td>\n",

       "      <td>NaN</td>\n",

       "      <td>NaN</td>\n",

       "      <td>NaN</td>\n",

       "    </tr>\n",

       "    <tr>\n",

       "      <th>3</th>\n",

       "      <td>case101_day20_slice_0004</td>\n",

       "      <td>NaN</td>\n",

       "      <td>NaN</td>\n",

       "      <td>NaN</td>\n",

       "    </tr>\n",

       "    <tr>\n",

       "      <th>4</th>\n",

       "      <td>case101_day20_slice_0005</td>\n",

       "      <td>NaN</td>\n",

       "      <td>NaN</td>\n",

       "      <td>NaN</td>\n",

       "    </tr>\n",

       "  </tbody>\n",

       "</table>\n",

       "</div>"

      ],

      "text/plain": [

       "class                        id large_bowel small_bowel stomach\n",

       "0      case101_day20_slice_0001         NaN         NaN     NaN\n",

       "1      case101_day20_slice_0002         NaN         NaN     NaN\n",

       "2      case101_day20_slice_0003         NaN         NaN     NaN\n",

       "3      case101_day20_slice_0004         NaN         NaN     NaN\n",

       "4      case101_day20_slice_0005         NaN         NaN     NaN"

      ]

     },

     "execution_count": 3,

     "metadata": {},

     "output_type": "execute_result"

    }

   ],

   "source": [

    "train_df.head()"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": 4,

   "metadata": {

    "ExecuteTime": {

     "end_time": "2022-06-28T03:46:47.950258Z",

     "start_time": "2022-06-28T03:46:47.742286Z"

    },

    "pycharm": {

     "name": "#%%\n"

    }

   },

   "outputs": [

    {

     "name": "stdout",

     "output_type": "stream",

     "text": [

      "dataset/train/case101/case101_day26/scans/slice_0121_266_266_1.50_1.50.png\n",

      "dataset/train/case101/case101_day26/scans/slice_0081_266_266_1.50_1.50.png\n",

      "dataset/train/case101/case101_day26/scans/slice_0065_266_266_1.50_1.50.png\n",

      "dataset/train/case101/case101_day26/scans/slice_0068_266_266_1.50_1.50.png\n",

      "dataset/train/case101/case101_day26/scans/slice_0059_266_266_1.50_1.50.png\n",

      "dataset/train/case101/case101_day26/scans/slice_0064_266_266_1.50_1.50.png\n",

      "dataset/train/case101/case101_day26/scans/slice_0002_266_266_1.50_1.50.png\n",

      "dataset/train/case101/case101_day26/scans/slice_0088_266_266_1.50_1.50.png\n",

      "dataset/train/case101/case101_day26/scans/slice_0098_266_266_1.50_1.50.png\n",

      "dataset/train/case101/case101_day26/scans/slice_0074_266_266_1.50_1.50.png\n",

      "dataset/train/case101/case101_day26/scans/slice_0021_266_266_1.50_1.50.png\n",

      "...\n",

      "And 38486 more lines.\n"

     ]

    }

   ],

   "source": [

    "path = Path('dataset/train')\n",

    "\n",

    "fnames = get_image_files(path)\n",

    "\n",

    "for ind, file_name in enumerate(fnames):\n",

    "    print(file_name)\n",

    "    if ind>9:\n",

    "        print('...')\n",

    "        break\n",

    "print(f\"And {len(fnames)-10} more lines.\")"

   ]

  },

  {

   "cell_type": "markdown",

   "metadata": {},

   "source": [

    "# Functions"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": 5,

   "metadata": {},

   "outputs": [],

   "source": [

    "def get_slice_id(fname):\n",

    "    return fname.parts[3] + '_' + fname.parts[5][:10]\n",

    "\n",

    "\n",

    "def rle_decode(mask_rle, shape, color=1):\n",

    "    \"\"\" TBD\n",

    "\n",

    "    Args:\n",

    "        mask_rle (str): run-length as string formated (start length)\n",

    "        shape (tuple of ints): (height,width) of array to return \n",

    "\n",

    "    Returns: \n",

    "        Mask (np.array)\n",

    "            - 1 indicating mask\n",

    "            - 0 indicating background\n",

    "\n",

    "    \"\"\"\n",

    "    # Split the string by space, then convert it into a integer array\n",

    "    s = np.array(mask_rle.split(), dtype=int)\n",

    "\n",

    "    # Every even value is the start, every odd value is the \"run\" length\n",

    "    starts = s[0::2] - 1\n",

    "    lengths = s[1::2]\n",

    "    ends = starts + lengths\n",

    "\n",

    "    # The image image is actually flattened since RLE is a 1D \"run\"\n",

    "    if len(shape) == 3:\n",

    "        h, w, d = shape\n",

    "        img = np.zeros((h * w, d), dtype=np.float32)\n",

    "    else:\n",

    "        h, w = shape\n",

    "        img = np.zeros((h * w,), dtype=np.float32)\n",

    "\n",

    "    # The color here is actually just any integer you want!\n",

    "    for lo, hi in zip(starts, ends):\n",

    "        img[lo: hi] = color\n",

    "\n",

    "    # Don't forget to change the image back to the original shape\n",

    "    return img.reshape(shape)\n",

    "\n",

    "\n",

    "def label_func(fname):\n",

    "    # First we need to get the slice row\n",

    "    slice_id = get_slice_id(fname)\n",

    "    slice_row = train_df.query('id == @slice_id')\n",

    "\n",

    "    # Then we need to extract the slice width and height which are provided in the fname last part\n",

    "    # Typically the height is the first part of a dimension, but for some reason the slices have\n",

    "    # widths provided first\n",

    "    w, h = map(lambda x: int(x), fname.parts[-1].split('_')[2:4])\n",

    "\n",

    "    # Create mask array (It needs to have 3 channels but fastai will only keep the first one anyways)\n",

    "    mask = np.zeros((h, w, 3), dtype=np.uint8)\n",

    "\n",

    "    # If the segmentation mask is str\n",

    "    # Each mask should have it's own code (color) where fastai will use them for identification\n",

    "    if isinstance(slice_row['large_bowel'].item(), str):\n",

    "        mask[:, :, 0] = rle_decode(slice_row['large_bowel'].item(), shape=(h, w), color=255)\n",

    "\n",

    "    if isinstance(slice_row['small_bowel'].item(), str):\n",

    "        mask[:, :, 1] = rle_decode(slice_row['small_bowel'].item(), shape=(h, w), color=255)\n",

    "\n",

    "    if isinstance(slice_row['stomach'].item(), str):\n",

    "        mask[:, :, 2] = rle_decode(slice_row['stomach'].item(), shape=(h, w), color=255)\n",

    "\n",

    "    return mask\n",

    "\n",

    "\n",

    "# This was the code available in fastai\n",

    "@ToTensor\n",

    "def encodes(self, o: PILMask): return o._tensor_cls(image2tensor(o)[0])\n",

    "\n",

    "\n",

    "# And this is how we customize it to suit our needs\n",

    "@ToTensor\n",

    "def encodes(self, o: PILMask): return o._tensor_cls(image2tensor(o))\n",

    "\n",

    "\n",

    "@typedispatch\n",

    "def show_batch(x: TensorImage, y: TensorMask, samples, ctxs=None, max_n=6, nrows=None, ncols=2,\n",

    "               figsize=None, **kwargs):\n",

    "    if figsize is None: figsize = (ncols * 3, max_n // ncols * 3)\n",

    "    if ctxs is None: ctxs = get_grid(max_n, nrows=nrows, ncols=ncols, figsize=figsize)\n",

    "    for i, ctx in enumerate(ctxs):\n",

    "        x_i = x[i] / x[i].max()\n",

    "        show_image(x_i, ctx=ctx, cmap='gray', **kwargs)\n",

    "        show_image(y[i], ctx=ctx, cmap='Spectral_r', alpha=0.35, **kwargs)\n",

    "        red_patch = mpatches.Patch(color='red', label='large_bowel')\n",

    "        green_patch = mpatches.Patch(color='green', label='small_bowel')\n",

    "        blue_patch = mpatches.Patch(color='blue', label='stomach')\n",

    "        ctx.legend(handles=[red_patch, green_patch, blue_patch], fontsize=figsize[0] / 2)\n",

    "\n",

    "\n",

    "def pad_img(img, up_size=None):\n",

    "    if up_size is None:\n",

    "        return img\n",

    "    shape0 = np.array(img.shape[:2])\n",

    "    resize = np.array(up_size)\n",

    "    if np.any(shape0 != resize):\n",

    "        diff = resize - shape0\n",

    "        pad0 = diff[0]\n",

    "        pad1 = diff[1]\n",

    "        pady = [pad0 // 2, pad0 // 2 + pad0 % 2]\n",

    "        padx = [pad1 // 2, pad1 // 2 + pad1 % 2]\n",

    "        img = np.pad(img, [pady, padx])\n",

    "        img = img.reshape((*resize))\n",

    "    return img\n",

    "\n",

    "\n",

    "def unpad_img(img, up_size, org_size):\n",

    "    shape0 = np.array(org_size)\n",

    "    resize = np.array(up_size)\n",

    "    if np.any(shape0 != resize):\n",

    "        diff = resize - shape0\n",

    "        pad0 = diff[0]\n",

    "        pad1 = diff[1]\n",

    "        pady = [pad0 // 2, pad0 // 2 + pad0 % 2]\n",

    "        padx = [pad1 // 2, pad1 // 2 + pad1 % 2]\n",

    "        img = img[pady[0]:-pady[1], padx[0]:-padx[1], :]\n",

    "        img = img.reshape((*shape0, 3))\n",

    "    return img\n",

    "\n",

    "\n",

    "def load_image(fname, up_size=None):\n",

    "    img = np.array(Image.open(fname))\n",

    "    img = np.interp(img, [np.min(img), np.max(img)], [0, 255])\n",

    "    return pad_img(img, up_size)\n",

    "\n",

    "\n",

    "def get_25D_image(row, up_size=None):\n",

    "    if up_size:\n",

    "        imgs = np.zeros((*up_size, len(row['fnames'])))\n",

    "    else:\n",

    "        imgs = np.zeros((row['slice_h'], row['slice_w'], len(row['fnames'])))\n",

    "\n",

    "    for i, fname in enumerate(row['fnames']):\n",

    "        img = load_image(fname, up_size)\n",

    "        imgs[..., i] += img\n",

    "    return imgs.astype(np.uint8)\n",

    "\n",

    "\n",

    "def get_mask(row, up_size=None):\n",

    "    if up_size:\n",

    "        mask = np.zeros((*up_size, 3))\n",

    "    else:\n",

    "        mask = np.zeros((row['slice_h'], row['slice_w'], 3))\n",

    "\n",

    "    if isinstance(row['large_bowel'], str):\n",

    "        mask[..., 0] += pad_img(\n",

    "            rle_decode(row['large_bowel'], shape=(row['slice_h'], row['slice_w']), color=255), up_size)\n",

    "    if isinstance(row['small_bowel'], str):\n",

    "        mask[..., 1] += pad_img(\n",

    "            rle_decode(row['small_bowel'], shape=(row['slice_h'], row['slice_w']), color=255), up_size)\n",

    "    if isinstance(row['stomach'], str):\n",

    "        mask[..., 2] += pad_img(\n",

    "            rle_decode(row['stomach'], shape=(row['slice_h'], row['slice_w']), color=255), up_size)\n",

    "\n",

    "    return mask.astype(np.uint8)\n",

    "\n",

    "\n",

    "def get_train_aug(img_size, crop=0.9, p=0.4):\n",

    "    crop_size = round(img_size[0] * crop)\n",

    "    return A.Compose([\n",

    "        A.RandomCrop(height=crop_size, width=crop_size, always_apply=True),\n",

    "        A.HorizontalFlip(p=p),\n",

    "        A.OneOf([\n",

    "            A.GridDistortion(num_steps=5, distort_limit=0.05, p=1.0),\n",

    "            A.ElasticTransform(alpha=1, sigma=50, alpha_affine=50, p=1.0)\n",

    "        ], p=p),\n",

    "        A.CoarseDropout(\n",

    "            max_holes=8, min_holes=8,\n",

    "            max_height=crop_size // 10, max_width=crop_size // 10,\n",

    "            min_height=4, min_width=4, mask_fill_value=0, p=0.2 * p),\n",

    "        A.ShiftScaleRotate(\n",

    "            shift_limit=0.0625, scale_limit=0.2, rotate_limit=25,\n",

    "            interpolation=cv2.INTER_AREA, p=p),\n",

    "        A.HorizontalFlip(p=0.5 * p),\n",

    "        A.OneOf([\n",

    "            A.MotionBlur(p=0.2 * p),\n",

    "            A.MedianBlur(blur_limit=3, p=0.1 * p),\n",

    "            A.Blur(blur_limit=3, p=0.1 * p),\n",

    "        ], p=0.2 * p),\n",

    "        A.GaussNoise(var_limit=0.001, p=0.2 * p),\n",

    "        A.OneOf([\n",

    "            A.OpticalDistortion(p=0.3 * p),\n",

    "            A.GridDistortion(p=0.1 * p),\n",

    "            A.PiecewiseAffine(p=0.3 * p),\n",

    "        ], p=0.2 * p),\n",

    "        A.OneOf([\n",

    "            A.Sharpen(p=0.2 * p),\n",

    "            A.Emboss(p=0.2 * p),\n",

    "            A.RandomBrightnessContrast(p=0.2 * p),\n",

    "        ]),\n",

    "    ])\n",

    "\n",

    "\n",

    "def get_test_aug(img_size, crop=0.9):\n",

    "    crop_size = round(crop * img_size[0])\n",

    "    return A.Compose([\n",

    "        A.CenterCrop(height=crop_size, width=crop_size),\n",

    "    ])\n",

    "\n",

    "\n",

    "class AlbumentationsTransform(ItemTransform, RandTransform):\n",

    "    split_idx, order = None, 2\n",

    "\n",

    "    def __init__(self, train_aug, valid_aug):\n",

    "        store_attr()\n",

    "\n",

    "    def before_call(self, b, split_idx):\n",

    "        self.idx = split_idx\n",

    "\n",

    "    def encodes(self, x):\n",

    "        if len(x) > 1:\n",

    "            img, mask = x\n",

    "            if self.idx == 0:\n",

    "                aug = self.train_aug(image=np.array(img), mask=np.array(mask))\n",

    "            else:\n",

    "                aug = self.valid_aug(image=np.array(img), mask=np.array(mask))\n",

    "            return PILImage.create(aug[\"image\"]), PILMask.create(aug[\"mask\"])\n",

    "        else:\n",

    "            img = x[0]\n",

    "            aug = self.valid_aug(image=np.array(img))\n",

    "            return PILImage.create(aug[\"image\"])\n"

   ]

  },

  {

   "cell_type": "markdown",

   "metadata": {},

   "source": [

    "# Preprocess"

   ]

  },

  {

   "cell_type": "markdown",

   "metadata": {},

   "source": [

    "## Split name metada into different columns"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": 6,

   "metadata": {

    "ExecuteTime": {

     "end_time": "2022-06-28T03:46:48.134095Z",

     "start_time": "2022-06-28T03:46:47.951449Z"

    },

    "pycharm": {

     "name": "#%%\n"

    }

   },

   "outputs": [],

   "source": [

    "train_df['partial_fname'] = train_df.id\n",

    "fname_df = pd.DataFrame({'partial_fname': [f'{fname.parts[-3]}_slice_{fname.parts[-1][6:10]}' for fname in fnames],\n",

    "                         'fname': fnames})\n",

    "\n",

    "train_df = train_df.merge(fname_df, on='partial_fname').drop('partial_fname', axis=1)\n",

    "\n",

    "train_df['case_id'] = train_df.id.apply(lambda x: x.split('_')[0])\n",

    "train_df['day_num'] = train_df.id.apply(lambda x: x.split('_')[1])\n",

    "\n",

    "train_df['slice_w'] = train_df[\"fname\"].apply(lambda x: int(str(x)[:-4].rsplit(\"_\",4)[1]))\n",

    "train_df['slice_h'] = train_df[\"fname\"].apply(lambda x: int(str(x)[:-4].rsplit(\"_\",4)[2]))"

   ]

  },

  {

   "cell_type": "markdown",

   "metadata": {},

   "source": [

    "## Add Fnames"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": 7,

   "metadata": {

    "ExecuteTime": {

     "end_time": "2022-06-28T03:46:48.186750Z",

     "start_time": "2022-06-28T03:46:48.135151Z"

    },

    "pycharm": {

     "name": "#%%\n"

    }

   },

   "outputs": [],

   "source": [

    "channels = 3\n",

    "stride = 2\n",

    "for j, i in enumerate(range(-1*(channels-channels//2-1), channels//2+1)):\n",

    "    method = 'ffill'\n",

    "    if i <= 0: method = 'bfill'\n",

    "    train_df[f'fname_{j:02}'] = train_df.groupby(['case_id', 'day_num'])['fname'].shift(stride*-i).fillna(method=method)\n",

    "    \n",

    "train_df['fnames'] = train_df[[f'fname_{j:02d}' for j in range(channels)]].values.tolist()"

   ]

  },

  {

   "cell_type": "markdown",

   "metadata": {

    "pycharm": {

     "name": "#%%\n"

    }

   },

   "source": [

    "# Training"

   ]

  },

  {

   "cell_type": "markdown",

   "metadata": {

    "pycharm": {

     "name": "#%%\n"

    }

   },

   "source": [

    "## Metrics"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": 8,

   "metadata": {

    "ExecuteTime": {

     "end_time": "2022-06-28T03:46:49.086943Z",

     "start_time": "2022-06-28T03:46:49.086930Z"

    }

   },

   "outputs": [],

   "source": [

    "def dice_coeff_adj(inp, targ):\n",

    "    inp = np.where(sigmoid(inp).cpu().detach().numpy() > 0.5, 1, 0)\n",

    "    targ = targ.cpu().detach().numpy()\n",

    "    eps = 1e-5\n",

    "    dice_scores = []\n",

    "    for i in range(targ.shape[0]):\n",

    "        dice_i = []\n",

    "        for j in range(targ.shape[1]):\n",

    "            if inp[i, j].sum() == targ[i, j].sum() == 0:\n",

    "                continue\n",

    "            I = (targ[i, j] * inp[i, j]).sum()\n",

    "            U =  targ[i, j].sum() + inp[i, j].sum()\n",

    "            dice_i.append((2.*I)/(U+eps))\n",

    "        if dice_i:\n",

    "            dice_scores.append(np.mean(dice_i))\n",

    "    \n",

    "    if dice_scores:\n",

    "        return np.mean(dice_scores)\n",

    "    else:\n",

    "        return 0\n",

    "    \n",

    "    \n",

    "def hd_dist_per_slice(inp, targ, seed):    \n",

    "    inp = np.argwhere(inp) / np.array(inp.shape)\n",

    "    targ = np.argwhere(targ) / np.array(targ.shape)\n",

    "    haussdorf_dist = 1 - directed_hausdorff(inp, targ, seed)[0]\n",

    "    return haussdorf_dist if haussdorf_dist > 0 else 0\n",

    "\n",

    "def hd_dist_adj(inp, targ, seed=42):\n",

    "    inp = np.where(sigmoid(inp).cpu().detach().numpy() > 0.5, 1, 0)\n",

    "    targ = targ.cpu().detach().numpy()\n",

    "    hd_scores = []\n",

    "    for i in range(targ.shape[0]):\n",

    "        hd_i = []\n",

    "        for j in range(targ.shape[1]):\n",

    "            if inp[i, j].sum() == targ[i, j].sum() == 0:\n",

    "                continue\n",

    "            hd_i.append(hd_dist_per_slice(inp[i, j], targ[i, j], seed))\n",

    "        if hd_i:\n",

    "            hd_scores.append(np.mean(hd_i))\n",

    "    if hd_scores:\n",

    "        return np.mean(hd_scores)\n",

    "    else:\n",

    "        return 0\n",

    "\n",

    "def custom_metric_adj(inp, targ, seed=42):\n",

    "    hd_score_per_batch = hd_dist_adj(inp, targ, seed)\n",

    "    dice_score_per_batch = dice_coeff_adj(inp, targ)\n",

    "    \n",

    "    return 0.4*dice_score_per_batch + 0.6*hd_score_per_batch"

   ]

  },

  {

   "cell_type": "markdown",

   "metadata": {},

   "source": [

    "## Loss Functions"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": 9,

   "metadata": {

    "ExecuteTime": {

     "end_time": "2022-06-28T03:46:49.088279Z",

     "start_time": "2022-06-28T03:46:49.088271Z"

    }

   },

   "outputs": [],

   "source": [

    "class DiceBCEModule(Module):\n",

    "    def __init__(self, eps:float=1e-5, from_logits=True):\n",

    "        store_attr()\n",

    "        \n",

    "    def forward(self, inp:Tensor, targ:Tensor) -> Tensor:\n",

    "        inp = inp.view(-1)\n",

    "        targ = targ.view(-1)\n",

    "        \n",

    "        if self.from_logits: \n",

    "            bce_loss = nn.BCEWithLogitsLoss()(inp, targ)\n",

    "            inp = torch.sigmoid(inp)\n",

    "            \n",

    "            \n",

    "        intersection = (inp * targ).sum()                            \n",

    "        dice = (2.*intersection + self.eps)/(inp.sum() + targ.sum() + self.eps)  \n",

    "        \n",

    "        return 0.5*(1 - dice) + 0.5*bce_loss\n",

    "\n",

    "\n",

    "class DiceBCELoss(BaseLoss):\n",

    "    def __init__(self, *args, eps:float=1e-5, from_logits=True, thresh=0.5, **kwargs):\n",

    "        super().__init__(DiceBCEModule, *args, eps=eps, from_logits=from_logits, flatten=False, is_2d=True, floatify=True, **kwargs)\n",

    "        self.thresh = thresh\n",

    "    \n",

    "    def decodes(self, x:Tensor) -> Tensor:\n",

    "        \"Converts model output to target format\"\n",

    "        return (x>self.thresh).long()\n",

    "\n",

    "    def activation(self, x:Tensor) -> Tensor:\n",

    "        \"`nn.BCEWithLogitsLoss`'s fused activation function applied to model output\"\n",

    "        return torch.sigmoid(x)\n",

    "\n",

    "# Source: https://www.kaggle.com/code/thedrcat/focal-multilabel-loss-in-pytorch-explained/notebook\n",

    "def focal_binary_cross_entropy(logits, targets, gamma=2, n=3):\n",

    "    p = torch.sigmoid(logits)\n",

    "    p = torch.where(targets >= 0.5, p, 1-p)\n",

    "    logp = - torch.log(torch.clamp(p, 1e-4, 1-1e-4))\n",

    "    loss = logp*((1-p)**gamma)\n",

    "    loss = n*loss.mean()\n",

    "    return loss\n",

    "\n",

    "class DiceFocalModule(Module):\n",

    "    def __init__(self, eps:float=1e-5, from_logits=True, ws=[0.5, 0.5], gamma=2, n=3):\n",

    "        store_attr()\n",

    "        \n",

    "    def forward(self, inp:Tensor, targ:Tensor) -> Tensor:\n",

    "        inp = inp.view(-1)\n",

    "        targ = targ.view(-1)\n",

    "        \n",

    "        if self.from_logits: \n",

    "            focal_loss = focal_binary_cross_entropy(inp, targ, self.gamma, self.n)\n",

    "            inp = torch.sigmoid(inp)\n",

    "            \n",

    "            \n",

    "        intersection = (inp * targ).sum()                            \n",

    "        dice = (2.*intersection + self.eps)/(inp.sum() + targ.sum() + self.eps)  \n",

    "        \n",

    "        return self.ws[0]*(1 - dice) + self.ws[1]*focal_loss\n",

    "    \n",

    "class DiceFocalLoss(BaseLoss):\n",

    "    def __init__(self, *args, eps:float=1e-5, from_logits=True, ws=[0.5, 0.5], gamma=2, n=3, thresh=0.5, **kwargs):\n",

    "        super().__init__(DiceFocalModule, *args, eps=eps, from_logits=from_logits, ws=ws, gamma=gamma, n=n, flatten=False, is_2d=True, floatify=True, **kwargs)\n",

    "        self.thresh = thresh\n",

    "    \n",

    "    def decodes(self, x:Tensor) -> Tensor:\n",

    "        \"Converts model output to target format\"\n",

    "        return (x>self.thresh).long()\n",

    "\n",

    "    def activation(self, x:Tensor) -> Tensor:\n",

    "        \"`nn.BCEWithLogitsLoss`'s fused activation function applied to model output\"\n",

    "        return torch.sigmoid(x)\n",

    "\n",

    "class FocalTverskyLossModule(Module):\n",

    "    def __init__(self, eps:float=1e-5, from_logits=True, alpha=0.3, beta=0.7, gamma=3/4):\n",

    "        store_attr()\n",

    "        \n",

    "    def forward(self, inp:Tensor, targ:Tensor) -> Tensor:\n",

    "        inp = inp.view(-1)\n",

    "        targ = targ.view(-1)\n",

    "        \n",

    "        if self.from_logits: \n",

    "            inp = torch.sigmoid(inp)\n",

    "            \n",

    "        inp_0, inp_1 = inp, 1 - inp\n",

    "        targ_0, targ_1 = targ, 1 - targ\n",

    "            \n",

    "        num = (inp_0 * targ_0).sum() \n",

    "        denom = num + (self.alpha * (inp_0 * targ_1).sum()) + (self.beta * (inp_1 * targ_0).sum()) + self.eps\n",

    "        loss = 1 - (num / denom)\n",

    "        return loss**self.gamma \n",

    "    \n",

    "class FocalTverskyLoss(BaseLoss):\n",

    "    def __init__(self, *args, eps:float=1e-5, from_logits=True, alpha=0.3, beta=0.7, gamma=3/4, thresh=0.5, **kwargs):\n",

    "        super().__init__(FocalTverskyLossModule, *args, eps=eps, from_logits=from_logits, alpha=alpha, beta=beta, gamma=gamma, flatten=False, is_2d=True, floatify=True, **kwargs)\n",

    "        self.thresh = thresh\n",

    "    \n",

    "    def decodes(self, x:Tensor) -> Tensor:\n",

    "        \"Converts model output to target format\"\n",

    "        return (x>self.thresh).long()\n",

    "\n",

    "    def activation(self, x:Tensor) -> Tensor:\n",

    "        \"`nn.BCEWithLogitsLoss`'s fused activation function applied to model output\"\n",

    "        return torch.sigmoid(x)\n",

    "\n",

    "def focal_binary_cross_entropy(logits, targets, gamma=2, n=3):\n",

    "    p = torch.sigmoid(logits)\n",

    "    p = torch.where(targets >= 0.5, p, 1-p)\n",

    "    logp = - torch.log(torch.clamp(p, 1e-4, 1-1e-4))\n",

    "    loss = logp*((1-p)**gamma)\n",

    "    loss = n*loss.mean()\n",

    "    return loss\n",

    "\n",

    "class ComboModule(Module):\n",

    "    def __init__(self, eps:float=1e-5, from_logits=True, ws=[2, 3, 1], gamma=2, n=3):\n",

    "        store_attr()\n",

    "        \n",

    "    def forward(self, inp:Tensor, targ:Tensor) -> Tensor:\n",

    "        inp = inp.view(-1)\n",

    "        targ = targ.view(-1)\n",

    "        \n",

    "        if self.from_logits: \n",

    "            focal_loss = focal_binary_cross_entropy(inp, targ, self.gamma, self.n)\n",

    "            bce_loss = nn.BCEWithLogitsLoss()(inp, targ)\n",

    "            inp = torch.sigmoid(inp)\n",

    "                \n",

    "        intersection = (inp * targ).sum()                            \n",

    "        dice = (2.*intersection + self.eps)/(inp.sum() + targ.sum() + self.eps)  \n",

    "        \n",

    "        return self.ws[0]*(1 - dice) + self.ws[1]*focal_loss + self.ws[2]*bce_loss\n",

    "    \n",

    "class ComboLoss(BaseLoss):\n",

    "    def __init__(self, *args, eps:float=1e-5, from_logits=True, ws=[2, 3, 1], gamma=2, n=3, thresh=0.5, **kwargs):\n",

    "        super().__init__(ComboModule, *args, eps=eps, from_logits=from_logits, ws=ws, gamma=gamma, n=n, flatten=False, is_2d=True, floatify=True, **kwargs)\n",

    "        self.thresh = thresh\n",

    "    \n",

    "    def decodes(self, x:Tensor) -> Tensor:\n",

    "        \"Converts model output to target format\"\n",

    "        return (x>self.thresh).long()\n",

    "\n",

    "    def activation(self, x:Tensor) -> Tensor:\n",

    "        \"`nn.BCEWithLogitsLoss`'s fused activation function applied to model output\"\n",

    "        return torch.sigmoid(x)"

   ]

  },

  {

   "cell_type": "markdown",

   "metadata": {},

   "source": [

    "## Set up CUDA"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": 10,

   "metadata": {},

   "outputs": [

    {

     "name": "stdout",

     "output_type": "stream",

     "text": [

      "device_name: GeForce RTX 2080 Ti\n",

      "device_capability: (7, 5)\n",

      "device_properties: _CudaDeviceProperties(name='GeForce RTX 2080 Ti', major=7, minor=5, total_memory=11019MB, multi_processor_count=68)\n",

      "current_device: 0\n",

      "\n",

      "Available Device IDs:  (0, 1, 2, 3)\n",

      "Device 0: 10620 MB free\n",

      "Device 1: 7899 MB free\n",

      "Device 2: 16274 MB free\n",

      "Device 3: 9241 MB free\n"

     ]

    }

   ],

   "source": [

    "# Select Cuda GPU device\n",

    "def get_memory_free_MiB(gpu_index):\n",

    "    pynvml.nvmlInit()\n",

    "    handle = pynvml.nvmlDeviceGetHandleByIndex(int(gpu_index))\n",

    "    mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle)\n",

    "    return mem_info.free // 1024 ** 2\n",

    "\n",

    "device_count = torch.cuda.device_count()\n",

    "current_device = torch.cuda.current_device()\n",

    "device_name = torch.cuda.get_device_name(current_device)\n",

    "device_capability = torch.cuda.get_device_capability(current_device)\n",

    "device_properties = torch.cuda.get_device_properties(current_device)\n",

    "is_available = torch.cuda.is_available()\n",

    "device_cuda = torch.device(\"cuda\")\n",

    "devices_tup = tuple(range(device_count))\n",

    "\n",

    "print('device_name: {device_name}'.format(device_name=device_name))\n",

    "print('device_capability: {device_capability}'.format(device_capability=device_capability))\n",

    "print('device_properties: {device_properties}'.format(device_properties=device_properties))\n",

    "print('current_device: {current_device}'.format(current_device=current_device))\n",

    "print(\"\\nAvailable Device IDs: \", devices_tup)\n",

    "\n",

    "for device_id in devices_tup:\n",

    "    print(f\"Device {device_id}: {get_memory_free_MiB(device_id)} MB free\")"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": 11,

   "metadata": {},

   "outputs": [

    {

     "name": "stdout",

     "output_type": "stream",

     "text": [

      "\n",

      "Successfully selected device 2\n"

     ]

    }

   ],

   "source": [

    "device_id = 2\n",

    "torch.backends.cudnn.benchmark = True\n",

    "\n",

    "if device_id is not None:\n",

    "    torch.cuda.set_device(device_id)\n",

    "    current_device = torch.cuda.current_device()\n",

    "    if current_device == device_id:\n",

    "        print(f\"\\nSuccessfully selected device {device_id}\")\n",

    "    else:\n",

    "        print(f\"Error: Couldn't change device from {current_device} to {device_id}\")\n"

   ]

  },

  {

   "cell_type": "markdown",

   "metadata": {},

   "source": [

    "## Baseline train using SMP"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": 12,

   "metadata": {

    "ExecuteTime": {

     "end_time": "2022-06-28T03:46:49.089093Z",

     "start_time": "2022-06-28T03:46:49.089085Z"

    }

   },

   "outputs": [],

   "source": [

    "def build_model(encoder_name, in_c=3, classes=3, weights=\"imagenet\"):\n",

    "    model = smp.Unet(\n",

    "        encoder_name=encoder_name,      \n",

    "        encoder_weights=weights,     \n",

    "        in_channels=in_c,                \n",

    "        classes=classes,        \n",

    "        activation=None\n",

    "    )\n",

    "    return model\n",

    "\n",

    "# Split any of model parameters from smp into encoder and decoder, \n",

    "# so that we can freeze and unfreeze encoder layers.\n",

    "def smp_splitter(model):\n",

    "    model_layers = list(model.children())\n",

    "    encoder_params = params(model_layers[0])\n",

    "    decoder_params = params(model_layers[1]) + params(model_layers[2])\n",

    "    return L(encoder_params, decoder_params)"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": 13,

   "metadata": {},

   "outputs": [

    {

     "name": "stderr",

     "output_type": "stream",

     "text": [

      "/home/kgeorgio/miniconda3/envs/gi-tract/lib/python3.7/site-packages/torch/_tensor.py:575: UserWarning: floor_divide is deprecated, and will be removed in a future version of pytorch. It currently rounds toward 0 (like the 'trunc' function NOT 'floor'). This results in incorrect rounding for negative values.\n",

      "To keep the current behavior, use torch.div(a, b, rounding_mode='trunc'), or for actual floor division, use torch.div(a, b, rounding_mode='floor'). (Triggered internally at  /opt/conda/conda-bld/pytorch_1631630742027/work/aten/src/ATen/native/BinaryOps.cpp:467.)\n",

      "  return torch.floor_divide(self, other)\n"

     ]

    }

   ],

   "source": [

    "dev_df = train_df.sample(frac=0.2)\n",

    "up_size = (320, 384)\n",

    "tfms = [[partial(get_25D_image, up_size=up_size), PILImage.create],\n",

    "        [partial(get_mask, up_size=up_size), PILMask.create]]  # the pipeline\n",

    "\n",

    "splits = RandomSplitter()(dev_df)\n",

    "# https://docs.fast.ai/tutorial.albumentations.html\n",

    "albu_aug = AlbumentationsTransform(get_train_aug(up_size), \n",

    "                                   get_test_aug(up_size))\n",

    "                                                                \n",

    "# https://docs.fast.ai/data.core.html#Datasets\n",

    "dsets = Datasets(train_df, tfms, splits=splits)\n",

    "# https://docs.fast.ai/data.load.html#DataLoader\n",

    "dls = dsets.dataloaders(bs=16, \n",

    "                        after_item=[albu_aug, ToTensor], \n",

    "                        after_batch=[IntToFloatTensor(div_mask=255),                                                                             \n",

    "                                     Normalize.from_stats(*imagenet_stats)],\n",

    "                        device=device_id)\n",

    "# https://smp.readthedocs.io/en/latest/encoders.html\n",

    "model = build_model('efficientnet-b0', in_c=3, classes=3, weights=\"imagenet\")\n",

    "model = model.cuda(device_id)\n",

    "metrics = [\n",

    "           dice_coeff_adj, \n",

    "           hd_dist_adj, \n",

    "           custom_metric_adj\n",

    "          ]\n",

    "loss_func = ComboLoss()\n",

    "splitter = smp_splitter"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": 14,

   "metadata": {

    "ExecuteTime": {

     "end_time": "2022-06-28T03:46:49.091707Z",

     "start_time": "2022-06-28T03:46:49.091699Z"

    }

   },

   "outputs": [

    {

     "data": {

      "text/html": [

       "\n",

       "<style>\n",

       "    /* Turns off some styling */\n",

       "    progress {\n",

       "        /* gets rid of default border in Firefox and Opera. */\n",

       "        border: none;\n",

       "        /* Needs to be in here for Safari polyfill so background images work as expected. */\n",

       "        background-size: auto;\n",

       "    }\n",

       "    .progress-bar-interrupted, .progress-bar-interrupted::-webkit-progress-bar {\n",

       "        background: #F44336;\n",

       "    }\n",

       "</style>\n"

      ],

      "text/plain": [

       "<IPython.core.display.HTML object>"

      ]

     },

     "metadata": {},

     "output_type": "display_data"

    },

    {

     "data": {

      "text/html": [

       "<table border=\"1\" class=\"dataframe\">\n",

       "  <thead>\n",

       "    <tr style=\"text-align: left;\">\n",

       "      <th>epoch</th>\n",

       "      <th>train_loss</th>\n",

       "      <th>valid_loss</th>\n",

       "      <th>dice_coeff_adj</th>\n",

       "      <th>hd_dist_adj</th>\n",

       "      <th>custom_metric_adj</th>\n",

       "      <th>time</th>\n",

       "    </tr>\n",

       "  </thead>\n",

       "  <tbody>\n",

       "    <tr>\n",

       "      <td>0</td>\n",

       "      <td>1.383521</td>\n",

       "      <td>1.162622</td>\n",

       "      <td>0.383626</td>\n",

       "      <td>0.797801</td>\n",

       "      <td>0.632131</td>\n",

       "      <td>01:41</td>\n",

       "    </tr>\n",

       "  </tbody>\n",

       "</table>"

      ],

      "text/plain": [

       "<IPython.core.display.HTML object>"

      ]

     },

     "metadata": {},

     "output_type": "display_data"

    },

    {

     "data": {

      "text/html": [

       "\n",

       "<style>\n",

       "    /* Turns off some styling */\n",

       "    progress {\n",

       "        /* gets rid of default border in Firefox and Opera. */\n",

       "        border: none;\n",

       "        /* Needs to be in here for Safari polyfill so background images work as expected. */\n",

       "        background-size: auto;\n",

       "    }\n",

       "    .progress-bar-interrupted, .progress-bar-interrupted::-webkit-progress-bar {\n",

       "        background: #F44336;\n",

       "    }\n",

       "</style>\n"

      ],

      "text/plain": [

       "<IPython.core.display.HTML object>"

      ]

     },

     "metadata": {},

     "output_type": "display_data"

    },

    {

     "data": {

      "text/html": [

       "<table border=\"1\" class=\"dataframe\">\n",

       "  <thead>\n",

       "    <tr style=\"text-align: left;\">\n",

       "      <th>epoch</th>\n",

       "      <th>train_loss</th>\n",

       "      <th>valid_loss</th>\n",

       "      <th>dice_coeff_adj</th>\n",

       "      <th>hd_dist_adj</th>\n",

       "      <th>custom_metric_adj</th>\n",

       "      <th>time</th>\n",

       "    </tr>\n",

       "  </thead>\n",

       "  <tbody>\n",

       "    <tr>\n",

       "      <td>0</td>\n",

       "      <td>0.814299</td>\n",

       "      <td>0.605600</td>\n",

       "      <td>0.621061</td>\n",

       "      <td>0.814855</td>\n",

       "      <td>0.737338</td>\n",

       "      <td>01:59</td>\n",

       "    </tr>\n",

       "  </tbody>\n",

       "</table>"

      ],

      "text/plain": [

       "<IPython.core.display.HTML object>"

      ]

     },

     "metadata": {},

     "output_type": "display_data"

    }

   ],

   "source": [

    "# https://docs.fast.ai/learner.html#Learner\n",

    "learn = Learner(dls, \n",

    "                model,\n",

    "                metrics=metrics, \n",

    "                loss_func=loss_func, \n",

    "                splitter=splitter).to_fp16()\n",

    "\n",

    "#https://docs.fast.ai/callback.schedule.html#Learner.fine_tune\n",

    "learn.freeze()\n",

    "learn.fine_tune(1, 1e-2)\n",

    "learn.export('test_model.pkl')"

   ]

  },

  {

   "cell_type": "markdown",

   "metadata": {

    "pycharm": {

     "name": "#%%\n"

    }

   },

   "source": [

    "## Use DynamicUnet"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": 15,

   "metadata": {},

   "outputs": [],

   "source": [

    "def timm_model_sizes(encoder, img_size):\n",

    "    sizes = []\n",

    "    for layer in encoder.feature_info:\n",

    "        sizes.append(torch.Size([1, layer['num_chs'], img_size[0]//layer['reduction'], img_size[1]//layer['reduction']]))\n",

    "    return sizes\n",

    "\n",

    "\n",

    "def get_timm_output_layers(encoder):\n",

    "    outputs = []\n",

    "    for layer in encoder.feature_info:\n",

    "        # Converts 'blocks.0.0' to ['blocks', '0', '0']\n",

    "        attrs = layer['module'].split('.')\n",

    "        output_layer = getattr(encoder, attrs[0])[int(attrs[1])][int(attrs[2])]\n",

    "        outputs.append(output_layer)\n",

    "    return outputs\n",

    "\n",

    "\n",

    "class DynamicTimmUnet(SequentialEx):\n",

    "    \"Create a U-Net from a given architecture in timm.\"\n",

    "    def __init__(self, encoder, n_out, img_size, blur=False, blur_final=True, self_attention=False,\n",

    "                 y_range=None, last_cross=True, bottle=False, act_cls=defaults.activation,\n",

    "                 init=nn.init.kaiming_normal_, norm_type=None, **kwargs):\n",

    "        imsize = img_size\n",