notebooks/segmentation_pytorch.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 19988,
     "status": "ok",
     "timestamp": 1642896029652,
     "user": {
      "displayName": "Sergio Aizcorbe",
      "photoUrl": "https://lh3.googleusercontent.com/a/default-user=s64",
      "userId": "06957247140334446809"
     },
     "user_tz": -60
    },
    "id": "JnIWCiPGfviG",
    "outputId": "610aaa4c-af55-44f3-d740-a9b60ca6dcf7",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Mounted at /content/gdrive\n"
     ]
    }
   ],
   "source": [
    "from google.colab import drive\n",
    "drive.mount('/content/gdrive')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "pepXmJfkSQR2",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Reqirements\n",
    "- keras >= 2.2.0 or tensorflow >= 1.13\n",
    "- segmenation-models==1.0.*\n",
    "- albumentations==0.3.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 28059,
     "status": "ok",
     "timestamp": 1642896057708,
     "user": {
      "displayName": "Sergio Aizcorbe",
      "photoUrl": "https://lh3.googleusercontent.com/a/default-user=s64",
      "userId": "06957247140334446809"
     },
     "user_tz": -60
    },
    "id": "EyOohIbnSQR5",
    "outputId": "4143072b-99e2-401f-9c77-5668b570e071",
    "pycharm": {
     "name": "#%%\n",
     "is_executing": true
    }
   },
   "outputs": [],
   "source": [
    "# Install required libs\n",
    "\n",
    "### please update Albumentations to version>=0.3.0 for `Lambda` transform support\n",
    "!pip install -U git+https://github.com/albu/albumentations --no-cache-dir\n",
    "\n",
    "!pip uninstall -y opencv-python\n",
    "!pip install opencv-python"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "2Gf65DtIj_uS",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# !git clone https://github.com/qubvel/segmentation_models.pytorch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 5910,
     "status": "ok",
     "timestamp": 1642896063606,
     "user": {
      "displayName": "Sergio Aizcorbe",
      "photoUrl": "https://lh3.googleusercontent.com/a/default-user=s64",
      "userId": "06957247140334446809"
     },
     "user_tz": -60
    },
    "id": "qEl3C_OGSnir",
    "outputId": "49b9cc23-e695-4fbb-abf0-701a2b7528be",
    "pycharm": {
     "name": "#%%\n",
     "is_executing": true
    }
   },
   "outputs": [],
   "source": [
    "%cd ./gdrive/MyDrive/Colab Notebooks/Solar Panels\n",
    "\n",
    "!pip install -r requirements.txt\n",
    "!ls"
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Loading dataset"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 56251,
     "status": "ok",
     "timestamp": 1642896119846,
     "user": {
      "displayName": "Sergio Aizcorbe",
      "photoUrl": "https://lh3.googleusercontent.com/a/default-user=s64",
      "userId": "06957247140334446809"
     },
     "user_tz": -60
    },
    "id": "_ca6no_ISQR_",
    "outputId": "160a14b8-8767-4e73-d8d2-fb0eff5902f9",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "0.2.1\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import cv2\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import albumentations as A\n",
    "\n",
    "import segmentation_models_pytorch as smp\n",
    "\n",
    "from pathlib import Path\n",
    "\n",
    "from torch.utils.data import DataLoader\n",
    "from torch.utils.data import Dataset\n",
    "\n",
    "print(smp.__version__)\n",
    "# os.environ['CUDA_VISIBLE_DEVICES'] = '0'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "EvwrRpc1SQSC",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Dataloader and utility functions "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "C8HpZzGlSQSC",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "from src.models.segmentation.datasets import *\n",
    "\n",
    "\n",
    "# # classes for data loading and preprocessing\n",
    "# class SolarPanelsDataset(Dataset):\n",
    "#     \"\"\"CamVid Dataset. Read images, apply augmentation and preprocessing transformations.\n",
    "#\n",
    "#     Args:\n",
    "#         images_dir (str): path to images folder\n",
    "#         masks_dir (str): path to segmentation masks folder\n",
    "#         class_values (list): values of classes to extract from segmentation mask\n",
    "#         augmentation (albumentations.Compose): data transformation pipeline\n",
    "#             (e.g. flip, scale, etc.)\n",
    "#         preprocessing (albumentations.Compose): data preprocessing\n",
    "#             (e.g. normalization, shape manipulation, etc.)\n",
    "#\n",
    "#     \"\"\"\n",
    "#\n",
    "#     CLASSES = ['solar_panel']\n",
    "#\n",
    "#     def __init__(\n",
    "#             self,\n",
    "#             images_dir,\n",
    "#             masks_dir,\n",
    "#             classes=None,\n",
    "#             augmentation=None,\n",
    "#             preprocessing=None,\n",
    "#     ):\n",
    "#         self.ids = os.listdir(images_dir)\n",
    "#         self.images_fps = [os.path.join(images_dir, image_id) for image_id in self.ids]\n",
    "#         self.masks_fps = [os.path.join(masks_dir, image_id.split('.')[0]+'_label.png') for image_id in self.ids]\n",
    "#\n",
    "#         # convert str names to class values on masks\n",
    "#         self.class_values = [self.CLASSES.index(cls.lower()) for cls in classes]\n",
    "#\n",
    "#         self.augmentation = augmentation\n",
    "#         self.preprocessing = preprocessing\n",
    "#\n",
    "#\n",
    "#     def __getitem__(self, i):\n",
    "#         # read data\n",
    "#         image = cv2.imread(self.images_fps[i])\n",
    "#         image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
    "#         mask = cv2.imread(self.masks_fps[i],0)\n",
    "#         mask = cv2.threshold(mask, 0, 255, cv2.THRESH_BINARY)[1]\n",
    "#\n",
    "#         # extract certain classes from mask (e.g. cars)\n",
    "#         masks = [(mask != v) for v in self.class_values]\n",
    "#         # masks = [(mask == v) for v in self.class_values]\n",
    "#         mask = np.stack(masks, axis=-1).astype('float')\n",
    "#\n",
    "#         # add background if mask is not binary\n",
    "#         if mask.shape[-1] != 1:\n",
    "#             background = 1 - mask.sum(axis=-1, keepdims=True)\n",
    "#             mask = np.concatenate((mask, background), axis=-1)\n",
    "#\n",
    "#         # apply augmentations\n",
    "#         if self.augmentation:\n",
    "#             sample = self.augmentation(image=image, mask=mask)\n",
    "#             image, mask = sample['image'], sample['mask']\n",
    "#\n",
    "#         # apply preprocessing\n",
    "#         if self.preprocessing:\n",
    "#             sample = self.preprocessing(image=image, mask=mask)\n",
    "#             image, mask = sample['image'], sample['mask']\n",
    "#\n",
    "#         return image, mask\n",
    "#\n",
    "#     def __len__(self):\n",
    "#         return len(self.ids)\n",
    "#\n",
    "# class GoogleMapsDataset(Dataset):\n",
    "#     \"\"\"CamVid Dataset. Read images, apply augmentation and preprocessing transformations.\n",
    "#\n",
    "#     Args:\n",
    "#         images_dir (str): path to images folder\n",
    "#         augmentation (albumentations.Compose): data transfromation pipeline\n",
    "#             (e.g. flip, scale, etc.)\n",
    "#         preprocessing (albumentations.Compose): data preprocessing\n",
    "#             (e.g. noralization, shape manipulation, etc.)\n",
    "#     \"\"\"\n",
    "#\n",
    "#     def __init__(\n",
    "#             self,\n",
    "#             images_dir,\n",
    "#             augmentation=None,\n",
    "#             preprocessing=None,\n",
    "#     ):\n",
    "#         self.ids = os.listdir(images_dir)\n",
    "#         self.images_fps = [os.path.join(images_dir, image_id) for image_id in self.ids]\n",
    "#\n",
    "#         self.augmentation = augmentation\n",
    "#         self.preprocessing = preprocessing\n",
    "#\n",
    "#     def __getitem__(self, i):\n",
    "#\n",
    "#         # read data\n",
    "#         image = cv2.imread(self.images_fps[i])\n",
    "#         image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
    "#\n",
    "#         # apply augmentations\n",
    "#         if self.augmentation:\n",
    "#             sample = self.augmentation(image=image)\n",
    "#             image = sample['image']\n",
    "#\n",
    "#         # apply preprocessing\n",
    "#         if self.preprocessing:\n",
    "#             sample = self.preprocessing(image=image)\n",
    "#             image = sample['image']\n",
    "#\n",
    "#         return image\n",
    "#\n",
    "#     def __len__(self):\n",
    "#         return len(self.ids)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "UTwvnsLLSQSH",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Augmentations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7EDufX9vSQSI",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Data augmentation is a powerful technique to increase the amount of your data and prevent model overfitting.  \n",
    "If you not familiar with such trick read some of these articles:\n",
    " - [The Effectiveness of Data Augmentation in Image Classification using Deep\n",
    "Learning](http://cs231n.stanford.edu/reports/2017/pdfs/300.pdf)\n",
    " - [Data Augmentation | How to use Deep Learning when you have Limited Data](https://medium.com/nanonets/how-to-use-deep-learning-when-you-have-limited-data-part-2-data-augmentation-c26971dc8ced)\n",
    " - [Data Augmentation Experimentation](https://towardsdatascience.com/data-augmentation-experimentation-3e274504f04b)\n",
    "\n",
    "Since our dataset is very small we will apply a large number of different augmentations:\n",
    " - horizontal flip\n",
    " - affine transforms\n",
    " - perspective transforms\n",
    " - brightness/contrast/colors manipulations\n",
    " - image bluring and sharpening\n",
    " - gaussian noise\n",
    " - random crops\n",
    "\n",
    "All this transforms can be easily applied with [**Albumentations**](https://github.com/albu/albumentations/) - fast augmentation library.\n",
    "For detailed explanation of image transformations you can look at [kaggle salt segmentation exmaple](https://github.com/albu/albumentations/blob/master/notebooks/example_kaggle_salt.ipynb) provided by [**Albumentations**](https://github.com/albu/albumentations/) authors.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "nnzfqi86SQSJ",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# define heavy augmentations\n",
    "def get_training_augmentation():\n",
    "    train_transform = [\n",
    "\n",
    "        A.HorizontalFlip(p=0.5),\n",
    "\n",
    "        A.ShiftScaleRotate(scale_limit=0.5, rotate_limit=0, shift_limit=0.1, p=1, border_mode=0),\n",
    "\n",
    "        A.PadIfNeeded(min_height=256, min_width=256, always_apply=True, border_mode=0),\n",
    "        A.RandomCrop(height=256, width=256, always_apply=True),\n",
    "\n",
    "        A.GaussNoise(p=0.2),\n",
    "        #A.Perspective(p=0.5),\n",
    "\n",
    "        A.OneOf(\n",
    "            [\n",
    "                A.CLAHE(p=1),\n",
    "                A.RandomBrightnessContrast(p=1),\n",
    "                A.RandomGamma(p=1),\n",
    "            ],\n",
    "            p=0.9,\n",
    "        ),\n",
    "\n",
    "        A.OneOf(\n",
    "            [\n",
    "                A.Sharpen(p=1),\n",
    "                A.Blur(blur_limit=3, p=1),\n",
    "                A.MotionBlur(blur_limit=3, p=1),\n",
    "            ],\n",
    "            p=0.9,\n",
    "        ),\n",
    "\n",
    "        A.OneOf(\n",
    "            [\n",
    "                A.RandomBrightnessContrast(p=1),\n",
    "                A.HueSaturationValue(p=1),\n",
    "            ],\n",
    "            p=0.9,\n",
    "        ),\n",
    "    ]\n",
    "    return A.Compose(train_transform)\n",
    "\n",
    "\n",
    "def get_validation_augmentation():\n",
    "    \"\"\"Add paddings to make image shape divisible by 32\"\"\"\n",
    "    test_transform = [\n",
    "        A.PadIfNeeded(256, 256)\n",
    "    ]\n",
    "    return A.Compose(test_transform)\n",
    "\n",
    "\n",
    "def to_tensor(x, **kwargs):\n",
    "    return x.transpose(2, 0, 1).astype('float32')\n",
    "\n",
    "\n",
    "def get_preprocessing(preprocessing_fn):\n",
    "    \"\"\"Construct preprocessing transform\n",
    "    \n",
    "    Args:\n",
    "        preprocessing_fn (callable): data normalization function\n",
    "            (can be specific for each pretrained neural network)\n",
    "    Return:\n",
    "        transform: albumentations.Compose\n",
    "    \n",
    "    \"\"\"\n",
    "    \n",
    "    _transform = [\n",
    "        A.Lambda(image=preprocessing_fn),\n",
    "        A.Lambda(image=to_tensor, mask=to_tensor),\n",
    "    ]\n",
    "    return A.Compose(_transform)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "SNH8F04tSQSK",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Segmentation model training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "_akS4ikJwR25",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def get_sp_dataset(folder, augmentation, params):\n",
    "    x_dir = os.path.join(params['data_dir'], f'{folder}/images')\n",
    "    y_dir = os.path.join(params['data_dir'], f'{folder}/masks')\n",
    "\n",
    "    return SolarPanelsDataset(\n",
    "        x_dir, y_dir,\n",
    "        classes=params['classes'],\n",
    "        augmentation=augmentation(),\n",
    "        preprocessing=get_preprocessing(smp.encoders.get_preprocessing_fn(params['encoder'])),\n",
    "    )\n",
    "\n",
    "def get_gm_dataset(folder, augmentation, params):\n",
    "    x_dir = os.path.join(params['data_dir'], folder)\n",
    "    return GoogleMapsDataset(\n",
    "        x_dir,\n",
    "        augmentation=augmentation(),\n",
    "        preprocessing=get_preprocessing(smp.encoders.get_preprocessing_fn(params['encoder'])),\n",
    "    )\n",
    "\n",
    "def get_model(model, encoder, n_classes, activation):\n",
    "    return model(encoder, classes=n_classes, activation=activation)\n",
    "\n",
    "def get_model_info(model_name):\n",
    "    info, ext = model_name.split('.')\n",
    "    arch, *enc, epochs = info.split('_')\n",
    "    \n",
    "    enc = '_'.join(enc[:-1])\n",
    "    raw_name = arch + '_' + enc\n",
    "    return raw_name, enc, int(epochs)\n",
    "\n",
    "def model_exists(model_name):\n",
    "    parent = Path(model_name).parent\n",
    "    name, _, _ = get_model_info(model_name)\n",
    "    for model in os.listdir(parent):\n",
    "        if model.startswith(name):\n",
    "            return os.path.join(parent, model)\n",
    "\n",
    "def get_optimizer(model, optimizer, lr):\n",
    "    return optimizer(params=model.parameters(), lr=lr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "hZv7hpCsN7Wb",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def train(train_params, device, verbose=True):\n",
    "\n",
    "    model_name = model_exists(train_params['model_name'])\n",
    "    n_classes = 1 if len(train_params['classes']) == 1 else (len(train_params['classes']) + 1)  # case for binary and multiclass segmentation\n",
    "\n",
    "    if model_name is not None:\n",
    "\n",
    "        model = torch.load(model_name)\n",
    "        raw_name, _, prev_epochs = get_model_info(model_name)\n",
    "\n",
    "        if prev_epochs == 0:\n",
    "            print(f'There already exists a model: {model_name}')\n",
    "            return\n",
    "\n",
    "        train_params['epochs'] -= prev_epochs\n",
    "\n",
    "    else:\n",
    "        model = get_model(\n",
    "            model=train_params['architecture'],\n",
    "            encoder=train_params['encoder'],\n",
    "            activation='sigmoid' if n_classes == 1 else 'softmax',\n",
    "            n_classes=n_classes,\n",
    "        )\n",
    "        \n",
    "    train_dataset = get_sp_dataset('train', get_training_augmentation, train_params)  # Dataset for training images\n",
    "    valid_dataset = get_sp_dataset('val', get_validation_augmentation, train_params)  # Dataset for validation images\n",
    "\n",
    "    train_loader = DataLoader(\n",
    "        train_dataset,\n",
    "        batch_size=train_params['batch_size'],\n",
    "        shuffle=True,\n",
    "        num_workers=2\n",
    "    )\n",
    "    valid_loader = DataLoader(\n",
    "        valid_dataset,\n",
    "        batch_size=1,\n",
    "        shuffle=False,\n",
    "        num_workers=2\n",
    "    )\n",
    "    return trainloop(model, train_loader, valid_loader, train_params, device, verbose)\n",
    "\n",
    "\n",
    "def trainloop(model, train_loader, valid_loader, train_params, device, verbose):\n",
    "    optimizer = get_optimizer(model, train_params['optimizer'], train_params['lr'])\n",
    "\n",
    "    train_epoch = smp.train.TrainEpoch(\n",
    "        model, \n",
    "        loss=train_params['loss'], \n",
    "        metrics=train_params['metrics'], \n",
    "        optimizer=optimizer,\n",
    "        device=device,\n",
    "        verbose=verbose,\n",
    "    )\n",
    "    valid_epoch = smp.train.ValidEpoch(\n",
    "        model, \n",
    "        loss=train_params['loss'], \n",
    "        metrics=train_params['metrics'],\n",
    "        device=device,\n",
    "        verbose=verbose,\n",
    "    )\n",
    "\n",
    "    max_score = 0\n",
    "    print(train_params['model_name'])\n",
    "    for epoch in range(train_params['epochs']):\n",
    "        print(f'\\nEpoch: {epoch + 1}')\n",
    "        train_logs = train_epoch.run(train_loader)\n",
    "        valid_logs = valid_epoch.run(valid_loader)\n",
    "\n",
    "    if not os.path.exists(train_params['model_name']):\n",
    "        torch.save(model, train_params['model_name'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "bdTMBTUXSQSM",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "CLASSES = ['solar_panel']\n",
    "BATCH_SIZE = 16\n",
    "LR = 0.0001\n",
    "LOSS = smp.losses.DiceLoss(smp.losses.BINARY_MODE, from_logits=True)\n",
    "\n",
    "DEVICE = 'cuda'\n",
    "DATA_DIR = 'data'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "q1IEoN_LMCQ3",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def gen_params(arch, encoder, epochs):\n",
    "    return {\n",
    "        'architecture': arch,\n",
    "        'encoder': encoder,\n",
    "        'model_name': f'models_pytorch/{arch.__name__.lower()}_{encoder}_model_{epochs}.pth',\n",
    "        'data_dir': DATA_DIR,\n",
    "\n",
    "        'classes': CLASSES,\n",
    "        'lr': LR,\n",
    "        'epochs': epochs,\n",
    "        'batch_size': BATCH_SIZE,\n",
    "\n",
    "        'loss': LOSS,\n",
    "\n",
    "        'metrics': [smp.metrics.IoU(threshold=0.5),\n",
    "                    smp.metrics.Fscore(threshold=0.5)],\n",
    "        'optimizer': torch.optim.Adam\n",
    "    }\n",
    "\n",
    "def gen_test_params(model_name):\n",
    "    _, encoder, _ = get_model_info(model_name)\n",
    "    return {\n",
    "        'encoder': encoder,\n",
    "        'data_dir': DATA_DIR,\n",
    "        'classes': CLASSES,\n",
    "\n",
    "        'loss': LOSS,\n",
    "\n",
    "        'metrics': [smp.metrics.IoU(threshold=0.5),\n",
    "                    smp.metrics.Fscore(threshold=0.5)],\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "wvR1n76VQJsv",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def test(model, test_params, device):\n",
    "\n",
    "    test_dataset = get_sp_dataset('test', get_validation_augmentation, test_params)  # Dataset for validation images\n",
    "    test_dataloader = DataLoader(test_dataset, batch_size=1, shuffle=False, num_workers=2)\n",
    "\n",
    "    test_epoch = smp.train.ValidEpoch(\n",
    "        model=model,\n",
    "        loss=test_params['loss'],\n",
    "        metrics=test_params['metrics'],\n",
    "        device=device,\n",
    "    )\n",
    "\n",
    "    logs = test_epoch.run(test_dataloader)\n",
    "\n",
    "    return test_dataset, logs\n",
    "\n",
    "def inference(model, params, device, img_folder='gmaps'):\n",
    "    dataset = get_gm_dataset(img_folder, get_validation_augmentation, params)\n",
    "    dataloader = DataLoader(dataset, batch_size=1, shuffle=False, num_workers=2)\n",
    "\n",
    "    masks = []\n",
    "    model.eval()\n",
    "    with torch.no_grad():\n",
    "        for image in dataloader:\n",
    "            image = image.to(device)\n",
    "            pr_mask = model.predict(image).cpu()\n",
    "            pr_mask = (pr_mask.squeeze().numpy().round())\n",
    "            prob = model(image)\n",
    "            masks.append(pr_mask)\n",
    "\n",
    "    return masks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "CFcDi9F9h4m2",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# from itertools import product\n",
    "\n",
    "# epochs = 25\n",
    "# architectures = [smp.Unet, smp.UnetPlusPlus, smp.MAnet, smp.Linknet, smp.FPN, smp.PSPNet, smp.PAN, smp.DeepLabV3, smp.DeepLabV3Plus]\n",
    "\n",
    "# encoders = [\n",
    "#     # 'resnet50',\n",
    "#     # 'resnext50_32x4d',\n",
    "#     # 'timm-resnest50d_4s2x40d',\n",
    "#     # 'timm-res2next50', 'timm-regnetx_064', 'timm-gernet_m',\n",
    "#     # 'se_resnext101_32x4d',\n",
    "#     # 'densenet201',\n",
    "#     # 'xception',\n",
    "#     'efficientnet-b2',\n",
    "#     # 'timm-efficientnet-b3',\n",
    "#     'timm-mobilenetv3_large_100',\n",
    "#     'vgg16_bn', 'vgg19_bn'\n",
    "# ]\n",
    "# count = 0\n",
    "# for arch, encoder in product(architectures, encoders):\n",
    "#     train_params = gen_params(arch, encoder, epochs)\n",
    "#     print('ARCH:', train_params['architecture'].__name__)\n",
    "#     print('ENCODER:', train_params['encoder'])\n",
    "\n",
    "#     train(train_params, DEVICE, verbose=False)\n",
    "\n",
    "#     best_model = torch.load(train_params['model_name'])\n",
    "#     _, logs = test(best_model, train_params, DEVICE)\n",
    "\n",
    "#     if logs.get('fscore') > 0.92:\n",
    "#         break\n",
    "\n",
    "#     print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "V1ZdJcRI23dl",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "ARCHITECTURE = smp.UnetPlusPlus\n",
    "ENCODER = 'se_resnext101_32x4d'\n",
    "EPOCHS = 50\n",
    "\n",
    "train_params = gen_params(ARCHITECTURE, ENCODER, EPOCHS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "UVOKWcKSXNds",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "train(train_params, DEVICE)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TU1g4n1nSQSP",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Model Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "0qWxzcN3HNhp",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import matplotlib.patches as patches\n",
    "from PIL import Image, ImageDraw\n",
    "from skimage.filters import threshold_otsu\n",
    "import skimage.measure as km\n",
    "from scipy import ndimage as nd\n",
    "import skimage.morphology as morph\n",
    "\n",
    "\n",
    "def post_process(image):\n",
    "    image = nd.binary_closing(image)\n",
    "    image = nd.binary_fill_holes(image)\n",
    "    image = morph.erosion(image, selem=morph.disk(5))\n",
    "    image = morph.dilation(image, selem=morph.disk(5))\n",
    "    return image\n",
    "\n",
    "def overlap(image, mask):\n",
    "    color = np.array([255, 0, 0], dtype='uint8')  # color to fill\n",
    "\n",
    "    # equal color where mask, else image\n",
    "    # this would paint your object silhouette entirely with `color`\n",
    "    masked_img = np.where(mask[...,None], color, image)\n",
    "\n",
    "    # use `addWeighted` to blend the two images\n",
    "    # the object will be tinted toward `color`\n",
    "    out = cv2.addWeighted(image, 0.7, masked_img, 0.2,0)\n",
    "    return out\n",
    "\n",
    "def visualize(**images):\n",
    "    \"\"\"\n",
    "    Helper function for data visualization\n",
    "    Plot images in one row.\"\"\"\n",
    "    n = len(images)\n",
    "    plt.figure(figsize=(16, 5))\n",
    "    for i, (name, image) in enumerate(images.items()):\n",
    "        plt.subplot(1, n, i + 1)\n",
    "        plt.xticks([])\n",
    "        plt.yticks([])\n",
    "        plt.title(' '.join(name.split('_')).title())\n",
    "        plt.imshow(image)\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "7jw9Czi1zJSA",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "status": "ok",
     "timestamp": 1642874737725,
     "user_tz": -60,
     "elapsed": 9393,
     "user": {
      "displayName": "Sergio Aizcorbe",
      "photoUrl": "https://lh3.googleusercontent.com/a/default-user=s64",
      "userId": "06957247140334446809"
     }
    },
    "outputId": "9c29ca3f-6741-4db5-b6d0-20bddc326cab",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "data/plots\n"
     ]
    }
   ],
   "source": [
    "# model_name = 'deeplabv3plus_timm-efficientnet-b3_model_25.pth' # !\n",
    "model_name = 'unetplusplus_timm-resnest50d_4s2x40d_model_50.pth'\n",
    "best_model = torch.load(os.path.join('models_pytorch', model_name))\n",
    "path = 'gmaps/images'\n",
    "path = 'plots'\n",
    "\n",
    "params = gen_test_params(model_name)\n",
    "masks = inference(best_model, params, DEVICE, img_folder=path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "TL2Zz76cWjO1",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "dataset = GoogleMapsDataset('data/' + path)\n",
    "for idx, (img, mask) in enumerate(zip(dataset, masks)):\n",
    "    mask = post_process(mask)\n",
    "    if 1 in mask:\n",
    "        mask = overlap(img, mask)\n",
    "        mask_img = Image.fromarray(mask)\n",
    "        mask_img.save(f'data/plots{idx}.png')\n",
    "        # visualize(image=img, mask=mask)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000,
     "output_embedded_package_id": "1LNbif3fKRjNw-3ruh38nujGruvx3hgul"
    },
    "executionInfo": {
     "elapsed": 30345,
     "status": "ok",
     "timestamp": 1642816992245,
     "user": {
      "displayName": "Sergio Aizcorbe",
      "photoUrl": "https://lh3.googleusercontent.com/a/default-user=s64",
      "userId": "06957247140334446809"
     },
     "user_tz": -60
    },
    "id": "cLNFYNYy9il6",
    "outputId": "f20a6647-548a-4d5d-83b3-9103445df0c6",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "Output hidden; open in https://colab.research.google.com to view."
     },
     "metadata": {}
    }
   ],
   "source": [
    "path = 'plots'\n",
    "dataset = GoogleMapsDataset(os.path.join(DATA_DIR, path))\n",
    "\n",
    "for model in os.listdir('models_pytorch'):\n",
    "    print(model)\n",
    "\n",
    "    best_model = torch.load(os.path.join('models_pytorch', model))\n",
    "    masks = inference(best_model, gen_test_params(model), DEVICE, img_folder=path)\n",
    "\n",
    "    for img, mask in zip(dataset, masks):\n",
    "        mask = post_process(mask)\n",
    "        if 1 in mask:\n",
    "            mask_img = Image.fromarray(overlap(img, mask))\n",
    "            visualize(image=img, mask=mask_img)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "7618tQ0ZyBLn",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def test_all(models_path, arch=None, enc=None, epochs=None, inference=False):\n",
    "    results = {}\n",
    "    arch = str(arch) if arch is not None else ''\n",
    "    enc = str(enc) if enc is not None else ''\n",
    "    epochs = f'{epochs}.pth' if epochs is not None else ''    \n",
    "\n",
    "    for model in os.listdir(models_path):\n",
    "        \n",
    "        if model.startswith(arch) and enc in model and model.endswith(epochs):\n",
    "            print(model)\n",
    "            best_model = torch.load(f'{models_path}/{model}')\n",
    "\n",
    "            params = gen_test_params(model)\n",
    "\n",
    "            if inference:\n",
    "                result = inference(best_model, params, DEVICE)\n",
    "            else:\n",
    "                _, result = test(best_model, params, DEVICE)\n",
    "\n",
    "            # dataset = get_gm_dataset('test_gm', get_validation_augmentation, params)\n",
    "            results[model.split('.')[0]] = result\n",
    "\n",
    "    \n",
    "    return results"
   ]
  },
  {
   "cell_type": "code",
   "source": [
    "from pprint import pprint\n",
    "\n",
    "logs = {}\n",
    "for arch in ['unetplusplus', 'fpn', 'pspnet', 'deeplabv3plus']:\n",
    "    logs[arch] = test_all('models_pytorch', arch=arch)\n",
    "\n",
    "# pprint(logs['unetplusplus'])"
   ],
   "metadata": {
    "id": "H_FEivRvR3jX",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [
    "def compute_scores(results):\n",
    "    scores = {}\n",
    "    for arch in results:\n",
    "        iou = np.mean([x.get('iou_score') for x in results[arch].values()])\n",
    "        fsc = np.mean([x.get('fscore') for x in results[arch].values()])\n",
    "        scores[arch] = {'iou': iou, 'fscore': fsc}\n",
    "    return scores\n",
    "\n",
    "scores = compute_scores(logs)\n",
    "pprint(scores)"
   ],
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "VaiLWRYccAL3",
    "executionInfo": {
     "status": "ok",
     "timestamp": 1642901160651,
     "user_tz": -60,
     "elapsed": 367,
     "user": {
      "displayName": "Sergio Aizcorbe",
      "photoUrl": "https://lh3.googleusercontent.com/a/default-user=s64",
      "userId": "06957247140334446809"
     }
    },
    "outputId": "804a0a6b-831c-44e5-8b05-0c27dad0c583",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": null,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "{'deeplabv3plus': {'fscore': 0.8635911337846935, 'iou': 0.8025295631626088},\n",
      " 'fpn': {'fscore': 0.8770014081403256, 'iou': 0.8223624959882697},\n",
      " 'pspnet': {'fscore': 0.8174557797752686, 'iou': 0.7396663873501653},\n",
      " 'unetplusplus': {'fscore': 0.8874491387298703, 'iou': 0.8433931577938459}}\n"
     ]
    }
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "pytorch_sp_segmentation.ipynb",
   "provenance": []
  },
  "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.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}