ref_b5.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## General information\n",
    "\n",
    "In this kernel I work with the data from Understanding Clouds from Satellite Images competition.\n",
    "```\n",
    "Shallow clouds play a huge role in determining the Earth's climate. They’re also difficult to understand and to represent in climate models. By classifying different types of cloud organization, researchers at Max Planck hope to improve our physical understanding of these clouds, which in turn will help us build better climate models.\n",
    "```\n",
    "\n",
    "So in this competition we are tasked with multiclass segmentation task: finding 4 different cloud patterns in the images. On the other hand, we make predictions for each pair of image and label separately, so this could be treated as 4 binary segmentation tasks.\n",
    "It is important to notice that images (and masks) are `1400 x 2100`, but predicted masks should be `350 x 525`.\n",
    "\n",
    "In this kernel I'll use (or will use in next versions) the following notable libraries:\n",
    "- [albumentations](https://github.com/albu/albumentations): this is a great library for image augmentation which makes it easier and more convenient\n",
    "- [catalyst](https://github.com/catalyst-team/catalyst): this is a great library which makes using PyTorch easier, helps with reprodicibility and contains a lot of useful utils\n",
    "- [segmentation_models_pytorch](https://github.com/qubvel/segmentation_models.pytorch): this is a great library with convenient wrappers for models, losses and other useful things\n",
    "- [pytorch-toolbelt](https://github.com/BloodAxe/pytorch-toolbelt): this is a great library with many useful shortcuts for building pytorch models\n",
    "\n",
    "\n",
    "UPD: Version 35 - changed calculation of optimal threshold and min size\n",
    "![](https://i.imgur.com/EOvz5kd.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "_kg_hide-output": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: catalyst in /home/nelson/miniconda3/lib/python3.7/site-packages (19.10.2)\n",
      "Requirement already satisfied: tqdm>=4.29.1 in /home/nelson/miniconda3/lib/python3.7/site-packages (from catalyst) (4.36.1)\n",
      "Requirement already satisfied: torchvision>=0.2.1 in /home/nelson/miniconda3/lib/python3.7/site-packages (from catalyst) (0.4.0)\n",
      "Requirement already satisfied: plotly>=4.1.0 in /home/nelson/miniconda3/lib/python3.7/site-packages (from catalyst) (4.2.1)\n",
      "Requirement already satisfied: tensorboardX>=1.8 in /home/nelson/miniconda3/lib/python3.7/site-packages (from catalyst) (1.9)\n",
      "Requirement already satisfied: matplotlib in /home/nelson/miniconda3/lib/python3.7/site-packages (from catalyst) (3.1.1)\n",
      "Requirement already satisfied: seaborn in /home/nelson/miniconda3/lib/python3.7/site-packages (from catalyst) (0.9.0)\n",
      "Requirement already satisfied: safitty>=1.2.3 in /home/nelson/miniconda3/lib/python3.7/site-packages (from catalyst) (1.3)\n",
      "Requirement already satisfied: opencv-python in /home/nelson/miniconda3/lib/python3.7/site-packages (from catalyst) (4.1.1.26)\n",
      "Requirement already satisfied: imageio in /home/nelson/miniconda3/lib/python3.7/site-packages (from catalyst) (2.6.1)\n",
      "Requirement already satisfied: pandas>=0.22 in /home/nelson/miniconda3/lib/python3.7/site-packages (from catalyst) (0.25.2)\n",
      "Requirement already satisfied: GitPython>=2.1.11 in /home/nelson/miniconda3/lib/python3.7/site-packages (from catalyst) (3.0.4)\n",
      "Requirement already satisfied: numpy>=1.16.4 in /home/nelson/miniconda3/lib/python3.7/site-packages (from catalyst) (1.17.2)\n",
      "Requirement already satisfied: ipython in /home/nelson/miniconda3/lib/python3.7/site-packages (from catalyst) (7.9.0)\n",
      "Requirement already satisfied: torch>=0.4.1 in /home/nelson/miniconda3/lib/python3.7/site-packages (from catalyst) (1.2.0)\n",
      "Requirement already satisfied: PyYAML in /home/nelson/miniconda3/lib/python3.7/site-packages (from catalyst) (5.1.2)\n",
      "Requirement already satisfied: crc32c>=1.7 in /home/nelson/miniconda3/lib/python3.7/site-packages (from catalyst) (1.7)\n",
      "Requirement already satisfied: scikit-learn>=0.20 in /home/nelson/miniconda3/lib/python3.7/site-packages (from catalyst) (0.21.3)\n",
      "Requirement already satisfied: scikit-image>=0.14.2 in /home/nelson/miniconda3/lib/python3.7/site-packages (from catalyst) (0.16.2)\n",
      "Requirement already satisfied: six in /home/nelson/miniconda3/lib/python3.7/site-packages (from torchvision>=0.2.1->catalyst) (1.12.0)\n",
      "Requirement already satisfied: pillow>=4.1.1 in /home/nelson/miniconda3/lib/python3.7/site-packages (from torchvision>=0.2.1->catalyst) (6.2.1)\n",
      "Requirement already satisfied: retrying>=1.3.3 in /home/nelson/miniconda3/lib/python3.7/site-packages (from plotly>=4.1.0->catalyst) (1.3.3)\n",
      "Requirement already satisfied: protobuf>=3.8.0 in /home/nelson/miniconda3/lib/python3.7/site-packages (from tensorboardX>=1.8->catalyst) (3.10.0)\n",
      "Requirement already satisfied: cycler>=0.10 in /home/nelson/miniconda3/lib/python3.7/site-packages (from matplotlib->catalyst) (0.10.0)\n",
      "Requirement already satisfied: kiwisolver>=1.0.1 in /home/nelson/miniconda3/lib/python3.7/site-packages (from matplotlib->catalyst) (1.1.0)\n",
      "Requirement already satisfied: pyparsing!=2.0.4,!=2.1.2,!=2.1.6,>=2.0.1 in /home/nelson/miniconda3/lib/python3.7/site-packages (from matplotlib->catalyst) (2.4.2)\n",
      "Requirement already satisfied: python-dateutil>=2.1 in /home/nelson/miniconda3/lib/python3.7/site-packages (from matplotlib->catalyst) (2.8.0)\n",
      "Requirement already satisfied: scipy>=0.14.0 in /home/nelson/miniconda3/lib/python3.7/site-packages (from seaborn->catalyst) (1.3.1)\n",
      "Requirement already satisfied: pytz>=2017.2 in /home/nelson/miniconda3/lib/python3.7/site-packages (from pandas>=0.22->catalyst) (2019.3)\n",
      "Requirement already satisfied: gitdb2>=2.0.0 in /home/nelson/miniconda3/lib/python3.7/site-packages (from GitPython>=2.1.11->catalyst) (2.0.6)\n",
      "Requirement already satisfied: decorator in /home/nelson/miniconda3/lib/python3.7/site-packages (from ipython->catalyst) (4.4.1)\n",
      "Requirement already satisfied: pygments in /home/nelson/miniconda3/lib/python3.7/site-packages (from ipython->catalyst) (2.4.2)\n",
      "Requirement already satisfied: prompt-toolkit<2.1.0,>=2.0.0 in /home/nelson/miniconda3/lib/python3.7/site-packages (from ipython->catalyst) (2.0.10)\n",
      "Requirement already satisfied: traitlets>=4.2 in /home/nelson/miniconda3/lib/python3.7/site-packages (from ipython->catalyst) (4.3.3)\n",
      "Requirement already satisfied: jedi>=0.10 in /home/nelson/miniconda3/lib/python3.7/site-packages (from ipython->catalyst) (0.15.1)\n",
      "Requirement already satisfied: backcall in /home/nelson/miniconda3/lib/python3.7/site-packages (from ipython->catalyst) (0.1.0)\n",
      "Requirement already satisfied: pexpect; sys_platform != \"win32\" in /home/nelson/miniconda3/lib/python3.7/site-packages (from ipython->catalyst) (4.7.0)\n",
      "Requirement already satisfied: pickleshare in /home/nelson/miniconda3/lib/python3.7/site-packages (from ipython->catalyst) (0.7.5)\n",
      "Requirement already satisfied: setuptools>=18.5 in /home/nelson/miniconda3/lib/python3.7/site-packages (from ipython->catalyst) (41.4.0)\n",
      "Requirement already satisfied: joblib>=0.11 in /home/nelson/miniconda3/lib/python3.7/site-packages (from scikit-learn>=0.20->catalyst) (0.13.2)\n",
      "Requirement already satisfied: networkx>=2.0 in /home/nelson/miniconda3/lib/python3.7/site-packages (from scikit-image>=0.14.2->catalyst) (2.4)\n",
      "Requirement already satisfied: PyWavelets>=0.4.0 in /home/nelson/miniconda3/lib/python3.7/site-packages (from scikit-image>=0.14.2->catalyst) (1.1.1)\n",
      "Requirement already satisfied: smmap2>=2.0.0 in /home/nelson/miniconda3/lib/python3.7/site-packages (from gitdb2>=2.0.0->GitPython>=2.1.11->catalyst) (2.0.5)\n",
      "Requirement already satisfied: wcwidth in /home/nelson/miniconda3/lib/python3.7/site-packages (from prompt-toolkit<2.1.0,>=2.0.0->ipython->catalyst) (0.1.7)\n",
      "Requirement already satisfied: ipython-genutils in /home/nelson/miniconda3/lib/python3.7/site-packages (from traitlets>=4.2->ipython->catalyst) (0.2.0)\n",
      "Requirement already satisfied: parso>=0.5.0 in /home/nelson/miniconda3/lib/python3.7/site-packages (from jedi>=0.10->ipython->catalyst) (0.5.1)\n",
      "Requirement already satisfied: ptyprocess>=0.5 in /home/nelson/miniconda3/lib/python3.7/site-packages (from pexpect; sys_platform != \"win32\"->ipython->catalyst) (0.6.0)\n",
      "Requirement already satisfied: pretrainedmodels in /home/nelson/miniconda3/lib/python3.7/site-packages (0.7.4)\n",
      "Requirement already satisfied: munch in /home/nelson/miniconda3/lib/python3.7/site-packages (from pretrainedmodels) (2.5.0)\n",
      "Requirement already satisfied: torchvision in /home/nelson/miniconda3/lib/python3.7/site-packages (from pretrainedmodels) (0.4.0)\n",
      "Requirement already satisfied: torch in /home/nelson/miniconda3/lib/python3.7/site-packages (from pretrainedmodels) (1.2.0)\n",
      "Requirement already satisfied: tqdm in /home/nelson/miniconda3/lib/python3.7/site-packages (from pretrainedmodels) (4.36.1)\n",
      "Requirement already satisfied: six in /home/nelson/miniconda3/lib/python3.7/site-packages (from munch->pretrainedmodels) (1.12.0)\n",
      "Requirement already satisfied: pillow>=4.1.1 in /home/nelson/miniconda3/lib/python3.7/site-packages (from torchvision->pretrainedmodels) (6.2.1)\n",
      "Requirement already satisfied: numpy in /home/nelson/miniconda3/lib/python3.7/site-packages (from torchvision->pretrainedmodels) (1.17.2)\n",
      "Collecting git+https://github.com/qubvel/segmentation_models.pytorch\n",
      "  Cloning https://github.com/qubvel/segmentation_models.pytorch to /tmp/pip-req-build-cu46wm6c\n",
      "  Running command git clone -q https://github.com/qubvel/segmentation_models.pytorch /tmp/pip-req-build-cu46wm6c\n",
      "Requirement already satisfied (use --upgrade to upgrade): segmentation-models-pytorch==0.0.3 from git+https://github.com/qubvel/segmentation_models.pytorch in /home/nelson/miniconda3/lib/python3.7/site-packages\n",
      "Requirement already satisfied: torchvision<=0.4.0,>=0.2.2 in /home/nelson/miniconda3/lib/python3.7/site-packages (from segmentation-models-pytorch==0.0.3) (0.4.0)\n",
      "Requirement already satisfied: pretrainedmodels==0.7.4 in /home/nelson/miniconda3/lib/python3.7/site-packages (from segmentation-models-pytorch==0.0.3) (0.7.4)\n",
      "Requirement already satisfied: efficientnet-pytorch==0.4.0 in /home/nelson/miniconda3/lib/python3.7/site-packages (from segmentation-models-pytorch==0.0.3) (0.4.0)\n",
      "Requirement already satisfied: torch==1.2.0 in /home/nelson/miniconda3/lib/python3.7/site-packages (from torchvision<=0.4.0,>=0.2.2->segmentation-models-pytorch==0.0.3) (1.2.0)\n",
      "Requirement already satisfied: numpy in /home/nelson/miniconda3/lib/python3.7/site-packages (from torchvision<=0.4.0,>=0.2.2->segmentation-models-pytorch==0.0.3) (1.17.2)\n",
      "Requirement already satisfied: six in /home/nelson/miniconda3/lib/python3.7/site-packages (from torchvision<=0.4.0,>=0.2.2->segmentation-models-pytorch==0.0.3) (1.12.0)\n",
      "Requirement already satisfied: pillow>=4.1.1 in /home/nelson/miniconda3/lib/python3.7/site-packages (from torchvision<=0.4.0,>=0.2.2->segmentation-models-pytorch==0.0.3) (6.2.1)\n",
      "Requirement already satisfied: munch in /home/nelson/miniconda3/lib/python3.7/site-packages (from pretrainedmodels==0.7.4->segmentation-models-pytorch==0.0.3) (2.5.0)\n",
      "Requirement already satisfied: tqdm in /home/nelson/miniconda3/lib/python3.7/site-packages (from pretrainedmodels==0.7.4->segmentation-models-pytorch==0.0.3) (4.36.1)\n",
      "Building wheels for collected packages: segmentation-models-pytorch\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  Building wheel for segmentation-models-pytorch (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25h  Created wheel for segmentation-models-pytorch: filename=segmentation_models_pytorch-0.0.3-cp37-none-any.whl size=29991 sha256=bd588180606f0e4c2dc60758a17f0d71a752be162781520ff755f0890fa4e34d\n",
      "  Stored in directory: /tmp/pip-ephem-wheel-cache-ogknrxfl/wheels/79/3f/09/1587a252e0314d26ad242d6d2e165622ab95c95e5cfe4b942c\n",
      "Successfully built segmentation-models-pytorch\n",
      "Requirement already satisfied: pytorch_toolbelt in /home/nelson/miniconda3/lib/python3.7/site-packages (0.2.1)\n",
      "Requirement already satisfied: opencv-python>=4.0 in /home/nelson/miniconda3/lib/python3.7/site-packages (from pytorch_toolbelt) (4.1.1.26)\n",
      "Requirement already satisfied: torch>=1.1 in /home/nelson/miniconda3/lib/python3.7/site-packages (from pytorch_toolbelt) (1.2.0)\n",
      "Requirement already satisfied: torchvision>=0.3 in /home/nelson/miniconda3/lib/python3.7/site-packages (from pytorch_toolbelt) (0.4.0)\n",
      "Requirement already satisfied: numpy>=1.14.5 in /home/nelson/miniconda3/lib/python3.7/site-packages (from opencv-python>=4.0->pytorch_toolbelt) (1.17.2)\n",
      "Requirement already satisfied: pillow>=4.1.1 in /home/nelson/miniconda3/lib/python3.7/site-packages (from torchvision>=0.3->pytorch_toolbelt) (6.2.1)\n",
      "Requirement already satisfied: six in /home/nelson/miniconda3/lib/python3.7/site-packages (from torchvision>=0.3->pytorch_toolbelt) (1.12.0)\n",
      "Requirement already satisfied: torchvision==0.4 in /home/nelson/miniconda3/lib/python3.7/site-packages (0.4.0)\n",
      "Requirement already satisfied: torch==1.2.0 in /home/nelson/miniconda3/lib/python3.7/site-packages (from torchvision==0.4) (1.2.0)\n",
      "Requirement already satisfied: numpy in /home/nelson/miniconda3/lib/python3.7/site-packages (from torchvision==0.4) (1.17.2)\n",
      "Requirement already satisfied: six in /home/nelson/miniconda3/lib/python3.7/site-packages (from torchvision==0.4) (1.12.0)\n",
      "Requirement already satisfied: pillow>=4.1.1 in /home/nelson/miniconda3/lib/python3.7/site-packages (from torchvision==0.4) (6.2.1)\n"
     ]
    }
   ],
   "source": [
    "!pip install catalyst\n",
    "!pip install pretrainedmodels\n",
    "!pip install git+https://github.com/qubvel/segmentation_models.pytorch\n",
    "!pip install pytorch_toolbelt\n",
    "!pip install torchvision==0.4"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Importing libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "_cell_guid": "b1076dfc-b9ad-4769-8c92-a6c4dae69d19",
    "_kg_hide-input": true,
    "_kg_hide-output": true,
    "_uuid": "8f2839f25d086af736a60e9eeb907d3b93b6e0e5"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "pyarrow not available, switching to pickle. To install pyarrow, run `pip install pyarrow`.\n",
      "lz4 not available, disabling compression. To install lz4, run `pip install lz4`.\n",
      "wandb not available, to install wandb, run `pip install wandb`.\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import cv2\n",
    "import collections\n",
    "import time \n",
    "import tqdm\n",
    "from PIL import Image\n",
    "from functools import partial\n",
    "train_on_gpu = True\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "%matplotlib inline\n",
    "\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.metrics import roc_auc_score\n",
    "\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "import torch\n",
    "from torch.utils.data import TensorDataset, DataLoader,Dataset\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "from torch.optim import lr_scheduler\n",
    "from torch.utils.data.sampler import SubsetRandomSampler\n",
    "from torch.optim.lr_scheduler import StepLR, ReduceLROnPlateau, CosineAnnealingLR\n",
    "\n",
    "import albumentations as albu\n",
    "from albumentations import pytorch as AT\n",
    "\n",
    "from catalyst.data import Augmentor\n",
    "from catalyst.dl import utils\n",
    "from catalyst.data.reader import ImageReader, ScalarReader, ReaderCompose, LambdaReader\n",
    "from catalyst.dl.runner import SupervisedRunner\n",
    "from catalyst.contrib.models.segmentation import Unet\n",
    "from catalyst.dl.callbacks import DiceCallback, EarlyStoppingCallback, InferCallback, CheckpointCallback\n",
    "\n",
    "import segmentation_models_pytorch as smp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Helper functions and classes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "_kg_hide-input": true
   },
   "outputs": [],
   "source": [
    "def get_img(x, folder: str='train_images'):\n",
    "    \"\"\"\n",
    "    Return image based on image name and folder.\n",
    "    \"\"\"\n",
    "    data_folder = f\"{path}/{folder}\"\n",
    "    image_path = os.path.join(data_folder, x)\n",
    "    img = cv2.imread(image_path)\n",
    "    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n",
    "    return img\n",
    "\n",
    "\n",
    "def rle_decode(mask_rle: str = '', shape: tuple = (1400, 2100)):\n",
    "    '''\n",
    "    Decode rle encoded mask.\n",
    "    \n",
    "    :param mask_rle: run-length as string formatted (start length)\n",
    "    :param shape: (height, width) of array to return \n",
    "    Returns numpy array, 1 - mask, 0 - background\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, order='F')\n",
    "\n",
    "\n",
    "def make_mask(df: pd.DataFrame, image_name: str='img.jpg', shape: tuple = (1400, 2100)):\n",
    "    \"\"\"\n",
    "    Create mask based on df, image name and shape.\n",
    "    \"\"\"\n",
    "    encoded_masks = df.loc[df['im_id'] == image_name, 'EncodedPixels']\n",
    "    masks = np.zeros((shape[0], shape[1], 4), dtype=np.float32)\n",
    "\n",
    "    for idx, label in enumerate(encoded_masks.values):\n",
    "        if label is not np.nan:\n",
    "            mask = rle_decode(label)\n",
    "            masks[:, :, idx] = mask\n",
    "            \n",
    "    return masks\n",
    "\n",
    "\n",
    "def to_tensor(x, **kwargs):\n",
    "    \"\"\"\n",
    "    Convert image or mask.\n",
    "    \"\"\"\n",
    "    return x.transpose(2, 0, 1).astype('float32')\n",
    "\n",
    "\n",
    "def mask2rle(img):\n",
    "    '''\n",
    "    Convert mask to rle.\n",
    "    img: numpy array, 1 - mask, 0 - background\n",
    "    Returns run length as string formated\n",
    "    '''\n",
    "    pixels= img.T.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",
    "\n",
    "\n",
    "def visualize(image, mask, original_image=None, original_mask=None):\n",
    "    \"\"\"\n",
    "    Plot image and masks.\n",
    "    If two pairs of images and masks are passes, show both.\n",
    "    \"\"\"\n",
    "    fontsize = 14\n",
    "    class_dict = {0: 'Fish', 1: 'Flower', 2: 'Gravel', 3: 'Sugar'}\n",
    "    \n",
    "    if original_image is None and original_mask is None:\n",
    "        f, ax = plt.subplots(1, 5, figsize=(24, 24))\n",
    "\n",
    "        ax[0].imshow(image)\n",
    "        for i in range(4):\n",
    "            ax[i + 1].imshow(mask[:, :, i])\n",
    "            ax[i + 1].set_title(f'Mask {class_dict[i]}', fontsize=fontsize)\n",
    "    else:\n",
    "        f, ax = plt.subplots(2, 5, figsize=(24, 12))\n",
    "\n",
    "        ax[0, 0].imshow(original_image)\n",
    "        ax[0, 0].set_title('Original image', fontsize=fontsize)\n",
    "                \n",
    "        for i in range(4):\n",
    "            ax[0, i + 1].imshow(original_mask[:, :, i])\n",
    "            ax[0, i + 1].set_title(f'Original mask {class_dict[i]}', fontsize=fontsize)\n",
    "        \n",
    "        ax[1, 0].imshow(image)\n",
    "        ax[1, 0].set_title('Transformed image', fontsize=fontsize)\n",
    "        \n",
    "        \n",
    "        for i in range(4):\n",
    "            ax[1, i + 1].imshow(mask[:, :, i])\n",
    "            ax[1, i + 1].set_title(f'Transformed mask {class_dict[i]}', fontsize=fontsize)\n",
    "            \n",
    "            \n",
    "def visualize_with_raw(image, mask, original_image=None, original_mask=None, raw_image=None, raw_mask=None):\n",
    "    \"\"\"\n",
    "    Plot image and masks.\n",
    "    If two pairs of images and masks are passes, show both.\n",
    "    \"\"\"\n",
    "    fontsize = 14\n",
    "    class_dict = {0: 'Fish', 1: 'Flower', 2: 'Gravel', 3: 'Sugar'}\n",
    "\n",
    "    f, ax = plt.subplots(3, 5, figsize=(24, 12))\n",
    "\n",
    "    ax[0, 0].imshow(original_image)\n",
    "    ax[0, 0].set_title('Original image', fontsize=fontsize)\n",
    "\n",
    "    for i in range(4):\n",
    "        ax[0, i + 1].imshow(original_mask[:, :, i])\n",
    "        ax[0, i + 1].set_title(f'Original mask {class_dict[i]}', fontsize=fontsize)\n",
    "\n",
    "\n",
    "    ax[1, 0].imshow(raw_image)\n",
    "    ax[1, 0].set_title('Original image', fontsize=fontsize)\n",
    "\n",
    "    for i in range(4):\n",
    "        ax[1, i + 1].imshow(raw_mask[:, :, i])\n",
    "        ax[1, i + 1].set_title(f'Raw predicted mask {class_dict[i]}', fontsize=fontsize)\n",
    "        \n",
    "    ax[2, 0].imshow(image)\n",
    "    ax[2, 0].set_title('Transformed image', fontsize=fontsize)\n",
    "\n",
    "\n",
    "    for i in range(4):\n",
    "        ax[2, i + 1].imshow(mask[:, :, i])\n",
    "        ax[2, i + 1].set_title(f'Predicted mask with processing {class_dict[i]}', fontsize=fontsize)\n",
    "            \n",
    "            \n",
    "def plot_with_augmentation(image, mask, augment):\n",
    "    \"\"\"\n",
    "    Wrapper for `visualize` function.\n",
    "    \"\"\"\n",
    "    augmented = augment(image=image, mask=mask)\n",
    "    image_flipped = augmented['image']\n",
    "    mask_flipped = augmented['mask']\n",
    "    visualize(image_flipped, mask_flipped, original_image=image, original_mask=mask)\n",
    "    \n",
    "    \n",
    "sigmoid = lambda x: 1 / (1 + np.exp(-x))\n",
    "\n",
    "\n",
    "def post_process(probability, threshold, min_size):\n",
    "    \"\"\"\n",
    "    Post processing of each predicted mask, components with lesser number of pixels\n",
    "    than `min_size` are ignored\n",
    "    \"\"\"\n",
    "    # don't remember where I saw it\n",
    "    mask = cv2.threshold(probability, threshold, 1, cv2.THRESH_BINARY)[1]\n",
    "    num_component, component = cv2.connectedComponents(mask.astype(np.uint8))\n",
    "    predictions = np.zeros((350, 525), np.float32)\n",
    "    num = 0\n",
    "    for c in range(1, num_component):\n",
    "        p = (component == c)\n",
    "        if p.sum() > min_size:\n",
    "            predictions[p] = 1\n",
    "            num += 1\n",
    "    return predictions, num\n",
    "\n",
    "\n",
    "def get_training_augmentation():\n",
    "    train_transform = [\n",
    "\n",
    "        albu.HorizontalFlip(p=0.5),\n",
    "        albu.ShiftScaleRotate(scale_limit=0.5, rotate_limit=0, shift_limit=0.1, p=0.5, border_mode=0),\n",
    "        albu.GridDistortion(p=0.5),\n",
    "        albu.OpticalDistortion(p=0.5, distort_limit=2, shift_limit=0.5),\n",
    "        albu.Resize(320, 640)\n",
    "    ]\n",
    "    return albu.Compose(train_transform)\n",
    "\n",
    "\n",
    "def get_validation_augmentation():\n",
    "    \"\"\"Add paddings to make image shape divisible by 32\"\"\"\n",
    "    test_transform = [\n",
    "        albu.Resize(320, 640)\n",
    "    ]\n",
    "    return albu.Compose(test_transform)\n",
    "\n",
    "\n",
    "def get_preprocessing(preprocessing_fn):\n",
    "    \"\"\"Construct preprocessing transform\n",
    "    \n",
    "    Args:\n",
    "        preprocessing_fn (callbale): data normalization function \n",
    "            (can be specific for each pretrained neural network)\n",
    "    Return:\n",
    "        transform: albumentations.Compose\n",
    "    \n",
    "    \"\"\"\n",
    "    \n",
    "    _transform = [\n",
    "        albu.Lambda(image=preprocessing_fn),\n",
    "        albu.Lambda(image=to_tensor, mask=to_tensor),\n",
    "    ]\n",
    "    return albu.Compose(_transform)\n",
    "\n",
    "\n",
    "def dice(img1, img2):\n",
    "    img1 = np.asarray(img1).astype(np.bool)\n",
    "    img2 = np.asarray(img2).astype(np.bool)\n",
    "\n",
    "    intersection = np.logical_and(img1, img2)\n",
    "\n",
    "    return 2. * intersection.sum() / (img1.sum() + img2.sum())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data overview\n",
    "\n",
    "Let's have a look at the data first."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['ref_b5.py',\n",
       " 'train_fastai_nofunnel.ipynb',\n",
       " 'logs',\n",
       " 'None',\n",
       " '.gitignore',\n",
       " 'LICENSE',\n",
       " 'test_images',\n",
       " 'train_images',\n",
       " 'train_images.zip',\n",
       " 'train_prototype.ipynb',\n",
       " 'unet_fpn_train_fp16_fastai_err.ipynb',\n",
       " 'README.md',\n",
       " 'doc',\n",
       " '.ipynb_checkpoints',\n",
       " 'train.csv',\n",
       " 'ranger.py',\n",
       " 'segmentation-in-pytorch-using-convenient-tools.ipynb',\n",
       " 'unet_fpn_train_fp16_fastai.ipynb',\n",
       " 'test_images.zip',\n",
       " 'ref_b5.ipynb',\n",
       " 'requirements.txt',\n",
       " '.git',\n",
       " 'sample_submission.csv',\n",
       " 'le_jit.ipynb']"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "path = '/home/nelson/ucsi'\n",
    "os.listdir(path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have folders with train and test images, file with train image ids and masks and sample submission."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "train = pd.read_csv(f'{path}/train.csv')\n",
    "sub = pd.read_csv(f'{path}/sample_submission.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "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></th>\n",
       "      <th>Image_Label</th>\n",
       "      <th>EncodedPixels</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0011165.jpg_Fish</td>\n",
       "      <td>264918 937 266318 937 267718 937 269118 937 27...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>0011165.jpg_Flower</td>\n",
       "      <td>1355565 1002 1356965 1002 1358365 1002 1359765...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>0011165.jpg_Gravel</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>0011165.jpg_Sugar</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>002be4f.jpg_Fish</td>\n",
       "      <td>233813 878 235213 878 236613 878 238010 881 23...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "          Image_Label                                      EncodedPixels\n",
       "0    0011165.jpg_Fish  264918 937 266318 937 267718 937 269118 937 27...\n",
       "1  0011165.jpg_Flower  1355565 1002 1356965 1002 1358365 1002 1359765...\n",
       "2  0011165.jpg_Gravel                                                NaN\n",
       "3   0011165.jpg_Sugar                                                NaN\n",
       "4    002be4f.jpg_Fish  233813 878 235213 878 236613 878 238010 881 23..."
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "There are 5546 images in train dataset\n",
      "There are 3698 images in test dataset\n"
     ]
    }
   ],
   "source": [
    "n_train = len(os.listdir(f'{path}/train_images'))\n",
    "n_test = len(os.listdir(f'{path}/test_images'))\n",
    "print(f'There are {n_train} images in train dataset')\n",
    "print(f'There are {n_test} images in test dataset')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Gravel    5546\n",
       "Flower    5546\n",
       "Sugar     5546\n",
       "Fish      5546\n",
       "Name: Image_Label, dtype: int64"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train['Image_Label'].apply(lambda x: x.split('_')[1]).value_counts()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So we have ~5.5k images in train dataset and they can have up to 4 masks: Fish, Flower, Gravel and Sugar."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Sugar     3751\n",
       "Gravel    2939\n",
       "Fish      2781\n",
       "Flower    2365\n",
       "Name: Image_Label, dtype: int64"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train.loc[train['EncodedPixels'].isnull() == False, 'Image_Label'].apply(lambda x: x.split('_')[1]).value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2    2372\n",
       "3    1560\n",
       "1    1348\n",
       "4     266\n",
       "Name: Image_Label, dtype: int64"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train.loc[train['EncodedPixels'].isnull() == False, 'Image_Label'].apply(lambda x: x.split('_')[0]).value_counts().value_counts()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "But there are a lot of empty masks. In fact only 266 images have all four masks. It is important to remember this."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "train['label'] = train['Image_Label'].apply(lambda x: x.split('_')[1])\n",
    "train['im_id'] = train['Image_Label'].apply(lambda x: x.split('_')[0])\n",
    "\n",
    "\n",
    "sub['label'] = sub['Image_Label'].apply(lambda x: x.split('_')[1])\n",
    "sub['im_id'] = sub['Image_Label'].apply(lambda x: x.split('_')[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's have a look at the images and the masks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "#fig = plt.figure(figsize=(25, 16))\n",
    "#for j, im_id in enumerate(np.random.choice(train['im_id'].unique(), 4)):\n",
    "#    for i, (idx, row) in enumerate(train.loc[train['im_id'] == im_id].iterrows()):\n",
    "#        ax = fig.add_subplot(5, 4, j * 4 + i + 1, xticks=[], yticks=[])\n",
    "#        im = Image.open(f\"{path}/train_images/{row['Image_Label'].split('_')[0]}\")\n",
    "#        plt.imshow(im)\n",
    "#       mask_rle = row['EncodedPixels']\n",
    "#        try: # label might not be there!\n",
    "#            mask = rle_decode(mask_rle)\n",
    "#        except:\n",
    "#            mask = np.zeros((1400, 2100))\n",
    "#        plt.imshow(mask, alpha=0.5, cmap='gray')\n",
    "#        ax.set_title(f\"Image: {row['Image_Label'].split('_')[0]}. Label: {row['label']}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see that masks can overlap. Also we can see that clouds are really similar to fish, flower and so on. Another important point: masks are often quite big and can have seemingly empty areas."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Preparing data for modelling\n",
    "\n",
    "At first, let's create a list of unique image ids and the count of masks for images. This will allow us to make a stratified split based on this count."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "id_mask_count = train.loc[train['EncodedPixels'].isnull() == False, 'Image_Label'].apply(lambda x: x.split('_')[0]).value_counts().\\\n",
    "reset_index().rename(columns={'index': 'img_id', 'Image_Label': 'count'})\n",
    "train_ids, valid_ids = train_test_split(id_mask_count['img_id'].values, random_state=42, stratify=id_mask_count['count'], test_size=0.1)\n",
    "test_ids = sub['Image_Label'].apply(lambda x: x.split('_')[0]).drop_duplicates().values"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exploring augmentations with albumentations\n",
    "\n",
    "One of important things while working with images is choosing good augmentations. There are a lot of them, let's have a look at augmentations from albumentations!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "#image_name = '8242ba0.jpg'\n",
    "#image = get_img(image_name)\n",
    "#mask = make_mask(train, image_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "#visualize(image, mask)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is how original image and its masks look like. Let's try adding some augmentations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "#plot_with_augmentation(image, mask, albu.HorizontalFlip(p=1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "#plot_with_augmentation(image, mask, albu.VerticalFlip(p=1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "#plot_with_augmentation(image, mask, albu.RandomRotate90(p=1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "#plot_with_augmentation(image, mask, albu.ElasticTransform(p=1, alpha=120, sigma=120 * 0.05, alpha_affine=120 * 0.03))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "#plot_with_augmentation(image, mask, albu.GridDistortion(p=1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "#plot_with_augmentation(image, mask, albu.OpticalDistortion(p=1, distort_limit=2, shift_limit=0.5))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setting up data for training in Catalyst"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CloudDataset(Dataset):\n",
    "    def __init__(self, df: pd.DataFrame = None, datatype: str = 'train', img_ids: np.array = None,\n",
    "                 transforms = albu.Compose([albu.HorizontalFlip(),AT.ToTensor()]),\n",
    "                preprocessing=None):\n",
    "        self.df = df\n",
    "        if datatype != 'test':\n",
    "            self.data_folder = f\"{path}/train_images\"\n",
    "        else:\n",
    "            self.data_folder = f\"{path}/test_images\"\n",
    "        self.img_ids = img_ids\n",
    "        self.transforms = transforms\n",
    "        self.preprocessing = preprocessing\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        image_name = self.img_ids[idx]\n",
    "        mask = make_mask(self.df, image_name)\n",
    "        image_path = os.path.join(self.data_folder, image_name)\n",
    "        img = cv2.imread(image_path)\n",
    "        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n",
    "        augmented = self.transforms(image=img, mask=mask)\n",
    "        img = augmented['image']\n",
    "        mask = augmented['mask']\n",
    "        if self.preprocessing:\n",
    "            preprocessed = self.preprocessing(image=img, mask=mask)\n",
    "            img = preprocessed['image']\n",
    "            mask = preprocessed['mask']\n",
    "        return img, mask\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.img_ids)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we define model and training parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "ENCODER = 'efficientnet-b5'\n",
    "ENCODER_WEIGHTS = 'imagenet'\n",
    "DEVICE = 'cuda'\n",
    "\n",
    "ACTIVATION = None\n",
    "model = smp.FPN(\n",
    "    encoder_name=ENCODER, \n",
    "    encoder_weights=ENCODER_WEIGHTS, \n",
    "    classes=4, \n",
    "    activation=ACTIVATION,\n",
    ")\n",
    "preprocessing_fn = smp.encoders.get_preprocessing_fn(ENCODER, ENCODER_WEIGHTS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/nelson/miniconda3/lib/python3.7/site-packages/albumentations/augmentations/transforms.py:2879: UserWarning:\n",
      "\n",
      "Using lambda is incompatible with multiprocessing. Consider using regular functions or partial().\n",
      "\n"
     ]
    }
   ],
   "source": [
    "num_workers = 0\n",
    "bs = 4\n",
    "train_dataset = CloudDataset(df=train, datatype='train', img_ids=train_ids, transforms = get_training_augmentation(), preprocessing=get_preprocessing(preprocessing_fn))\n",
    "valid_dataset = CloudDataset(df=train, datatype='valid', img_ids=valid_ids, transforms = get_validation_augmentation(), preprocessing=get_preprocessing(preprocessing_fn))\n",
    "\n",
    "train_loader = DataLoader(train_dataset, batch_size=bs, shuffle=True, num_workers=num_workers)\n",
    "valid_loader = DataLoader(valid_dataset, batch_size=bs, shuffle=False, num_workers=num_workers)\n",
    "\n",
    "loaders = {\n",
    "    \"train\": train_loader,\n",
    "    \"valid\": valid_loader\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_epochs = 40\n",
    "logdir = \"./logs/segmentation\"\n",
    "\n",
    "# model, criterion, optimizer\n",
    "optimizer = torch.optim.Adam([\n",
    "    {'params': model.decoder.parameters(), 'lr': 1e-2}, \n",
    "    {'params': model.encoder.parameters(), 'lr': 1e-3},  \n",
    "])\n",
    "scheduler = ReduceLROnPlateau(optimizer, factor=0.15, patience=2)\n",
    "criterion = smp.utils.losses.BCEDiceLoss(eps=1.)\n",
    "runner = SupervisedRunner()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_kg_hide-output": true
   },
   "outputs": [
    {
     "data": {
      "application/javascript": [
       "IPython.notebook.save_checkpoint();"
      ],
      "text/plain": [
       "<IPython.core.display.Javascript object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "runner.train(\n",
    "    model=model,\n",
    "    criterion=criterion,\n",
    "    optimizer=optimizer,\n",
    "    scheduler=scheduler,\n",
    "    loaders=loaders,\n",
    "    callbacks=[DiceCallback(), EarlyStoppingCallback(patience=5, min_delta=0.001)],\n",
    "    logdir=logdir,\n",
    "    num_epochs=num_epochs,\n",
    "    verbose=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "utils.plot_metrics(\n",
    "    logdir=logdir, \n",
    "    # specify which metrics we want to plot\n",
    "    metrics=[\"loss\", \"dice\", 'lr', '_base/lr']\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exploring predictions\n",
    "Let's make predictions on validation dataset.\n",
    "\n",
    "At first we need to optimize thresholds "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=> loading checkpoint ./logs/segmentation/checkpoints/best.pth\n",
      "loaded checkpoint ./logs/segmentation/checkpoints/best.pth (epoch 15)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      "0it [00:00, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Top best models:\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "555it [00:54, 10.20it/s]\n"
     ]
    }
   ],
   "source": [
    "encoded_pixels = []\n",
    "loaders = {\"infer\": valid_loader}\n",
    "runner.infer(\n",
    "    model=model,\n",
    "    loaders=loaders,\n",
    "    callbacks=[\n",
    "        CheckpointCallback(\n",
    "            resume=f\"{logdir}/checkpoints/best.pth\"),\n",
    "        InferCallback()\n",
    "    ],\n",
    ")\n",
    "valid_masks = []\n",
    "probabilities = np.zeros((2220, 350, 525))\n",
    "for i, (batch, output) in enumerate(tqdm.tqdm(zip(\n",
    "        valid_dataset, runner.callbacks[0].predictions[\"logits\"]))):\n",
    "    image, mask = batch\n",
    "    for m in mask:\n",
    "        if m.shape != (350, 525):\n",
    "            m = cv2.resize(m, dsize=(525, 350), interpolation=cv2.INTER_LINEAR)\n",
    "        valid_masks.append(m)\n",
    "\n",
    "    for j, probability in enumerate(output):\n",
    "        if probability.shape != (350, 525):\n",
    "            probability = cv2.resize(probability, dsize=(525, 350), interpolation=cv2.INTER_LINEAR)\n",
    "        probabilities[i * 4 + j, :, :] = probability"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Find optimal values\n",
    "\n",
    "First of all, my thanks to @samusram for finding a mistake in my validation\n",
    "https://www.kaggle.com/c/understanding_cloud_organization/discussion/107711#622412\n",
    "\n",
    "And now I find optimal values separately for each class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "    threshold   size      dice\n",
      "74       0.70  10000  0.666609\n",
      "79       0.75  10000  0.665326\n",
      "69       0.65  10000  0.665295\n",
      "64       0.60  10000  0.661515\n",
      "78       0.75   5000  0.659725\n",
      "1\n",
      "    threshold   size      dice\n",
      "74       0.70  10000  0.804356\n",
      "69       0.65  10000  0.802267\n",
      "64       0.60  10000  0.800915\n",
      "54       0.50  10000  0.800741\n",
      "59       0.55  10000  0.799427\n",
      "2\n",
      "    threshold   size      dice\n",
      "64       0.60  10000  0.641579\n",
      "59       0.55  10000  0.640747\n",
      "69       0.65  10000  0.637136\n",
      "54       0.50  10000  0.634999\n",
      "68       0.65   5000  0.632911\n",
      "3\n",
      "    threshold   size      dice\n",
      "64       0.60  10000  0.646268\n",
      "74       0.70  10000  0.645251\n",
      "34       0.30  10000  0.641783\n",
      "39       0.35  10000  0.640958\n",
      "69       0.65  10000  0.640638\n"
     ]
    }
   ],
   "source": [
    "class_params = {}\n",
    "for class_id in range(4):\n",
    "    print(class_id)\n",
    "    attempts = []\n",
    "    for t in range(0, 100, 5):\n",
    "        t /= 100\n",
    "        for ms in [0, 100, 1200, 5000, 10000]:\n",
    "            masks = []\n",
    "            for i in range(class_id, len(probabilities), 4):\n",
    "                probability = probabilities[i]\n",
    "                predict, num_predict = post_process(sigmoid(probability), t, ms)\n",
    "                masks.append(predict)\n",
    "\n",
    "            d = []\n",
    "            for i, j in zip(masks, valid_masks[class_id::4]):\n",
    "                if (i.sum() == 0) & (j.sum() == 0):\n",
    "                    d.append(1)\n",
    "                else:\n",
    "                    d.append(dice(i, j))\n",
    "\n",
    "            attempts.append((t, ms, np.mean(d)))\n",
    "\n",
    "    attempts_df = pd.DataFrame(attempts, columns=['threshold', 'size', 'dice'])\n",
    "\n",
    "\n",
    "    attempts_df = attempts_df.sort_values('dice', ascending=False)\n",
    "    print(attempts_df.head())\n",
    "    best_threshold = attempts_df['threshold'].values[0]\n",
    "    best_size = attempts_df['size'].values[0]\n",
    "    \n",
    "    class_params[class_id] = (best_threshold, best_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{0: (0.7, 10000), 1: (0.7, 10000), 2: (0.6, 10000), 3: (0.6, 10000)}\n"
     ]
    }
   ],
   "source": [
    "print(class_params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#sns.lineplot(x='threshold', y='dice', hue='size', data=attempts_df);\n",
    "#plt.title('Threshold and min size vs dice for one of the classes');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's have a look at our masks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#for i, (input, output) in enumerate(zip(\n",
    "#        valid_dataset, runner.callbacks[0].predictions[\"logits\"])):\n",
    "#    image, mask = input\n",
    "#        \n",
    "#    image_vis = image.transpose(1, 2, 0)\n",
    "#    mask = mask.astype('uint8').transpose(1, 2, 0)\n",
    "#    pr_mask = np.zeros((350, 525, 4))\n",
    "#    for j in range(4):\n",
    "#        probability = cv2.resize(output.transpose(1, 2, 0)[:, :, j], dsize=(525, 350), interpolation=cv2.INTER_LINEAR)\n",
    "#        pr_mask[:, :, j], _ = post_process(sigmoid(probability), class_params[j][0], class_params[j][1])\n",
    "    #pr_mask = (sigmoid(output) > best_threshold).astype('uint8').transpose(1, 2, 0)\n",
    "    \n",
    "        \n",
    "#    visualize_with_raw(image=image_vis, mask=pr_mask, original_image=image_vis, original_mask=mask, raw_image=image_vis, raw_mask=output.transpose(1, 2, 0))\n",
    "    \n",
    "#    if i >= 2:\n",
    "#        break"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Predicting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "71"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import gc\n",
    "torch.cuda.empty_cache()\n",
    "gc.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_dataset = CloudDataset(df=sub, datatype='test', img_ids=test_ids, transforms = get_validation_augmentation(), preprocessing=get_preprocessing(preprocessing_fn))\n",
    "test_loader = DataLoader(test_dataset, batch_size=2, shuffle=False, num_workers=0)\n",
    "\n",
    "loaders = {\"test\": test_loader}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1849/1849 [13:03<00:00,  2.36it/s]\n"
     ]
    }
   ],
   "source": [
    "encoded_pixels = []\n",
    "image_id = 0\n",
    "for i, test_batch in enumerate(tqdm.tqdm(loaders['test'])):\n",
    "    runner_out = runner.predict_batch({\"features\": test_batch[0].cuda()})['logits']\n",
    "    for i, batch in enumerate(runner_out):\n",
    "        for probability in batch:\n",
    "            \n",
    "            probability = probability.cpu().detach().numpy()\n",
    "            if probability.shape != (350, 525):\n",
    "                probability = cv2.resize(probability, dsize=(525, 350), interpolation=cv2.INTER_LINEAR)\n",
    "            predict, num_predict = post_process(sigmoid(probability), class_params[image_id % 4][0], class_params[image_id % 4][1])\n",
    "            if num_predict == 0:\n",
    "                encoded_pixels.append('')\n",
    "            else:\n",
    "                r = mask2rle(predict)\n",
    "                encoded_pixels.append(r)\n",
    "            image_id += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "sub['EncodedPixels'] = encoded_pixels\n",
    "sub.to_csv('submission.csv', columns=['Image_Label', 'EncodedPixels'], index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}