mask-rcnn-and-transfer-learning.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "_uuid": "7c1fce19a11f95416168ced03c2c70fa818b21a5",
    "colab_type": "text",
    "id": "KBeAf8WgaeSk"
   },
   "source": [
    "**Mask-RCNN Starter Model for the RSNA Pneumonia Detection Challenge with transfer learning **\n",
    "\n",
    "Using pre-trained COCO weights trained on http://cocodataset.org as in https://github.com/matterport/Mask_RCNN/tree/master/samples/balloon\n",
    "\n",
    "We get the best public kernel performance so far, and also training only within the 6hrs kaggle limit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# install dependencies not included by Colab\n",
    "# use pip3 to ensure compatibility w/ Google Deep Learning Images \n",
    "!pip3 install -q pydicom \n",
    "!pip3 install -q imgaug \n",
    "!pip3 install -q kaggle\n",
    "!pip3 install -q tqdm "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "_uuid": "40c67b3ff0fa04587dec508363308adaa3ceaf34",
    "colab": {},
    "colab_type": "code",
    "id": "4kjcC6QqywWl"
   },
   "outputs": [],
   "source": [
    "import os \n",
    "import sys\n",
    "import random\n",
    "import math\n",
    "import numpy as np\n",
    "import cv2\n",
    "import matplotlib.pyplot as plt\n",
    "import json\n",
    "import pydicom\n",
    "from imgaug import augmenters as iaa\n",
    "from tqdm import tqdm\n",
    "import pandas as pd \n",
    "import glob "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "ename": "ModuleNotFoundError",
     "evalue": "No module named 'google.colab'",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
      "\u001b[1;32m<ipython-input-2-6f075b31acdc>\u001b[0m in \u001b[0;36m<module>\u001b[1;34m()\u001b[0m\n\u001b[0;32m      1\u001b[0m \u001b[1;31m#@title\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m----> 2\u001b[1;33m \u001b[1;32mfrom\u001b[0m \u001b[0mgoogle\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mcolab\u001b[0m \u001b[1;32mimport\u001b[0m \u001b[0mauth\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m      3\u001b[0m \u001b[1;32mfrom\u001b[0m \u001b[0mgoogleapiclient\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mhttp\u001b[0m \u001b[1;32mimport\u001b[0m \u001b[0mMediaFileUpload\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      4\u001b[0m \u001b[1;32mfrom\u001b[0m \u001b[0mgoogleapiclient\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mdiscovery\u001b[0m \u001b[1;32mimport\u001b[0m \u001b[0mbuild\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      5\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;31mModuleNotFoundError\u001b[0m: No module named 'google.colab'"
     ]
    }
   ],
   "source": [
    "#@title\n",
    "from google.colab import auth\n",
    "from googleapiclient.http import MediaFileUpload\n",
    "from googleapiclient.discovery import build\n",
    "\n",
    "auth.authenticate_user()\n",
    "\n",
    "drive_service = build('drive', 'v3')\n",
    "\n",
    "def save_file_to_drive(name, path):\n",
    "    file_metadata = {\n",
    "      'name': name,\n",
    "      'mimeType': 'application/octet-stream'\n",
    "     }\n",
    "  \n",
    "    media = MediaFileUpload(path, \n",
    "                    mimetype='application/octet-stream',\n",
    "                    resumable=True)\n",
    "\n",
    "    created = drive_service.files().create(body=file_metadata,\n",
    "                                   media_body=media,\n",
    "                                   fields='id').execute()\n",
    "\n",
    "    print('File ID: {}'.format(created.get('id')))\n",
    "\n",
    "    return created"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#@title\n",
    "# enter your Kaggle credentionals here\n",
    "os.environ['KAGGLE_USERNAME']=\"skooch\"\n",
    "os.environ['KAGGLE_KEY']=\"42f8a02ee92cc773d1dbe66565673ad3\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Updated property [core/project].\n"
     ]
    }
   ],
   "source": [
    "#@title\n",
    "# upload checkpoint to GCS\n",
    "project_id = 'mammography-198911'\n",
    "bucket_name = 'pneumonia'\n",
    "\n",
    "!gcloud config set project {project_id}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Traceback (most recent call last):\n",
      "  File \"c:\\users\\eric\\anaconda2\\envs\\hse_dl\\lib\\runpy.py\", line 193, in _run_module_as_main\n",
      "    \"__main__\", mod_spec)\n",
      "  File \"c:\\users\\eric\\anaconda2\\envs\\hse_dl\\lib\\runpy.py\", line 85, in _run_code\n",
      "    exec(code, run_globals)\n",
      "  File \"C:\\Users\\eric\\Anaconda2\\envs\\hse_dl\\Scripts\\kaggle.exe\\__main__.py\", line 5, in <module>\n",
      "  File \"c:\\users\\eric\\anaconda2\\envs\\hse_dl\\lib\\site-packages\\kaggle\\__init__.py\", line 23, in <module>\n",
      "    api.authenticate()\n",
      "  File \"c:\\users\\eric\\anaconda2\\envs\\hse_dl\\lib\\site-packages\\kaggle\\api\\kaggle_api_extended.py\", line 110, in authenticate\n",
      "    self._load_config(config_data)\n",
      "  File \"c:\\users\\eric\\anaconda2\\envs\\hse_dl\\lib\\site-packages\\kaggle\\api\\kaggle_api_extended.py\", line 149, in _load_config\n",
      "    raise ValueError('Error: Missing %s in configuration.' % item)\n",
      "ValueError: Error: Missing username in configuration.\n"
     ]
    }
   ],
   "source": [
    "#@title\n",
    "# If you are unable to download the competition dataset, check to see if you have \n",
    "# accepted the user agreement on the competition website. \n",
    "if not os.path.exists(\"./stage_1_detailed_class_info.csv.zip\"):\n",
    "    !kaggle competitions download -c rsna-pneumonia-detection-challenge"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "#@title\n",
    "# unzipping takes a few minutes\n",
    "if not os.path.exists(\"./stage_1_train_labels.csv\"):\n",
    "    print(\"Unzipping...\")\n",
    "    !unzip -q -o stage_1_test_images.zip -d stage_1_test_images\n",
    "    !unzip -q -o stage_1_train_images.zip -d stage_1_train_images\n",
    "    !unzip -q -o stage_1_train_labels.csv.zip"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_uuid": "6e5764759e6a0a9b698b44645658f66873edd807",
    "colab": {},
    "colab_type": "code",
    "id": "yP0XLJx_x_6o"
   },
   "outputs": [],
   "source": [
    "DATA_DIR = './'\n",
    "\n",
    "# Directory to save logs and trained model\n",
    "ROOT_DIR = './working'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "_uuid": "576df4c47a23d08b1bdb384245e09aa69f88bbd3",
    "colab_type": "text",
    "id": "kdYzLq1zfKL4"
   },
   "source": [
    "### Install Matterport's Mask-RCNN model from github.\n",
    "See the [Matterport's implementation of Mask-RCNN](https://github.com/matterport/Mask_RCNN)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_uuid": "b37d22551d332f0f7b722cc7204eb614524b6c21",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 143
    },
    "colab_type": "code",
    "id": "KgllzLnDr7kF",
    "outputId": "6c978df7-2013-437e-acd1-5011048dfb53"
   },
   "outputs": [],
   "source": [
    "!git clone https://www.github.com/matterport/Mask_RCNN.git\n",
    "os.chdir('Mask_RCNN')\n",
    "#!python setup.py -q install"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_uuid": "3acbbbe055b6a409d3c50ae0f893acf51b5ae7ba",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "colab_type": "code",
    "id": "-KZXyWwhzOVU",
    "outputId": "2576cc17-7484-4311-ad72-3c5643dcb5bb"
   },
   "outputs": [],
   "source": [
    "# Import Mask RCNN\n",
    "sys.path.append(os.path.join(ROOT_DIR, 'Mask_RCNN'))  # To find local version of the library\n",
    "from mrcnn.config import Config\n",
    "from mrcnn import utils\n",
    "import mrcnn.model as modellib\n",
    "from mrcnn import visualize\n",
    "from mrcnn.model import log"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_uuid": "50089cc61791871cdf6a5c0037dc4f28b7b7d7cc",
    "colab": {},
    "colab_type": "code",
    "id": "FghMmiMjzOX2"
   },
   "outputs": [],
   "source": [
    "train_dicom_dir = os.path.join(DATA_DIR, 'stage_1_train_images')\n",
    "test_dicom_dir = os.path.join(DATA_DIR, 'stage_1_test_images')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "_uuid": "f108beef7838be8a64dd512d395c5dc0ad952790"
   },
   "source": [
    "### Download COCO pre-trained weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_uuid": "c3ee0cd0ee0b1defdec97b94bc736587c1f7631f"
   },
   "outputs": [],
   "source": [
    "!wget --quiet https://github.com/matterport/Mask_RCNN/releases/download/v2.0/mask_rcnn_coco.h5\n",
    "!ls -lh mask_rcnn_coco.h5\n",
    "\n",
    "COCO_WEIGHTS_PATH = \"mask_rcnn_coco.h5\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "_uuid": "032cc5fe4baa051108106675e6ca4f4fdb2846ed",
    "colab_type": "text",
    "id": "gj-tvDvEaDiC"
   },
   "source": [
    "### Some setup functions and classes for Mask-RCNN\n",
    "\n",
    "- dicom_fps is a list of the dicom image path and filenames \n",
    "- image_annotions is a dictionary of the annotations keyed by the filenames\n",
    "- parsing the dataset returns a list of the image filenames and the annotations dictionary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_uuid": "778cb19865d7cc63440491aef9202b71c61e8bb2",
    "colab": {},
    "colab_type": "code",
    "id": "ivqC4cnszOaM"
   },
   "outputs": [],
   "source": [
    "def get_dicom_fps(dicom_dir):\n",
    "    dicom_fps = glob.glob(dicom_dir+'/'+'*.dcm')\n",
    "    return list(set(dicom_fps))\n",
    "\n",
    "def parse_dataset(dicom_dir, anns): \n",
    "    image_fps = get_dicom_fps(dicom_dir)\n",
    "    image_annotations = {fp: [] for fp in image_fps}\n",
    "    for index, row in anns.iterrows(): \n",
    "        fp = os.path.join(dicom_dir, row['patientId']+'.dcm')\n",
    "        image_annotations[fp].append(row)\n",
    "    return image_fps, image_annotations "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_uuid": "dfcffc4eaa94a41497717851dee9f702d8a2a73b",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 917
    },
    "colab_type": "code",
    "id": "_SfzTa-1zOck",
    "outputId": "91ae8935-bccb-4b8e-9a7e-aa690f95fd9b"
   },
   "outputs": [],
   "source": [
    "# The following parameters have been selected to reduce running time for demonstration purposes \n",
    "# These are not optimal \n",
    "\n",
    "class DetectorConfig(Config):\n",
    "    \"\"\"Configuration for training pneumonia detection on the RSNA pneumonia dataset.\n",
    "    Overrides values in the base Config class.\n",
    "    \"\"\"\n",
    "    \n",
    "    # Give the configuration a recognizable name  \n",
    "    NAME = 'pneumonia'\n",
    "    \n",
    "    # Train on 1 GPU and 8 images per GPU. We can put multiple images on each\n",
    "    # GPU because the images are small. Batch size is 8 (GPUs * images/GPU).\n",
    "    GPU_COUNT = 1\n",
    "    IMAGES_PER_GPU = 12\n",
    "    \n",
    "    BACKBONE = 'resnet50'\n",
    "    \n",
    "    NUM_CLASSES = 2  # background + 1 pneumonia classes\n",
    "    \n",
    "    IMAGE_MIN_DIM = 448\n",
    "    IMAGE_MAX_DIM = 448\n",
    "    RPN_ANCHOR_SCALES = (56, 112, 224, 448)\n",
    "    TRAIN_ROIS_PER_IMAGE = 32\n",
    "    MAX_GT_INSTANCES = 3\n",
    "    DETECTION_MAX_INSTANCES = 3\n",
    "    DETECTION_MIN_CONFIDENCE = 0.7\n",
    "    DETECTION_NMS_THRESHOLD = 0.1\n",
    "\n",
    "    STEPS_PER_EPOCH = 1000\n",
    "    \n",
    "config = DetectorConfig()\n",
    "config.display()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_uuid": "52bd3ffbdde0173a363055482d675da51c2aba99",
    "colab": {},
    "colab_type": "code",
    "id": "8EBVA1M60yAj"
   },
   "outputs": [],
   "source": [
    "class DetectorDataset(utils.Dataset):\n",
    "    \"\"\"Dataset class for training pneumonia detection on the RSNA pneumonia dataset.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, image_fps, image_annotations, orig_height, orig_width):\n",
    "        super().__init__(self)\n",
    "        \n",
    "        # Add classes\n",
    "        self.add_class('pneumonia', 1, 'Lung Opacity')\n",
    "        \n",
    "        # add images \n",
    "        for i, fp in enumerate(image_fps):\n",
    "            annotations = image_annotations[fp]\n",
    "            self.add_image('pneumonia', image_id=i, path=fp, \n",
    "                           annotations=annotations, orig_height=orig_height, orig_width=orig_width)\n",
    "            \n",
    "    def image_reference(self, image_id):\n",
    "        info = self.image_info[image_id]\n",
    "        return info['path']\n",
    "\n",
    "    def load_image(self, image_id):\n",
    "        info = self.image_info[image_id]\n",
    "        fp = info['path']\n",
    "        ds = pydicom.read_file(fp)\n",
    "        image = ds.pixel_array\n",
    "        # If grayscale. Convert to RGB for consistency.\n",
    "        if len(image.shape) != 3 or image.shape[2] != 3:\n",
    "            image = np.stack((image,) * 3, -1)\n",
    "        return image\n",
    "\n",
    "    def load_mask(self, image_id):\n",
    "        info = self.image_info[image_id]\n",
    "        annotations = info['annotations']\n",
    "        count = len(annotations)\n",
    "        if count == 0:\n",
    "            mask = np.zeros((info['orig_height'], info['orig_width'], 1), dtype=np.uint8)\n",
    "            class_ids = np.zeros((1,), dtype=np.int32)\n",
    "        else:\n",
    "            mask = np.zeros((info['orig_height'], info['orig_width'], count), dtype=np.uint8)\n",
    "            class_ids = np.zeros((count,), dtype=np.int32)\n",
    "            for i, a in enumerate(annotations):\n",
    "                if a['Target'] == 1:\n",
    "                    x = int(a['x'])\n",
    "                    y = int(a['y'])\n",
    "                    w = int(a['width'])\n",
    "                    h = int(a['height'])\n",
    "                    mask_instance = mask[:, :, i].copy()\n",
    "                    cv2.rectangle(mask_instance, (x, y), (x+w, y+h), 255, -1)\n",
    "                    mask[:, :, i] = mask_instance\n",
    "                    class_ids[i] = 1\n",
    "        return mask.astype(np.bool), class_ids.astype(np.int32)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "_uuid": "1cb852e262b69d348743767d675573368ab672c9",
    "colab_type": "text",
    "id": "9RlMo04ckd98"
   },
   "source": [
    "### Examine the annotation data, parse the dataset, and view dicom fields"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_uuid": "793b1c6c6ba4e5f0d51e130080aa799f230b5ef6",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 235
    },
    "colab_type": "code",
    "id": "EdhUEFDr0yDA",
    "outputId": "1715a5df-a577-41fd-bf20-f1a27aadb28c"
   },
   "outputs": [],
   "source": [
    "# training dataset\n",
    "anns = pd.read_csv(os.path.join(DATA_DIR, 'stage_1_train_labels.csv'))\n",
    "anns.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_uuid": "7aebc88f910b232e3b8759421914a007c6ffed94",
    "colab": {},
    "colab_type": "code",
    "id": "Mxz-pNbt5txY"
   },
   "outputs": [],
   "source": [
    "image_fps, image_annotations = parse_dataset(train_dicom_dir, anns=anns)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_uuid": "6c386dcef041b972f6209dd19e247d547c3c349f",
    "colab": {},
    "colab_type": "code",
    "id": "YPqjEIXWRhSf"
   },
   "outputs": [],
   "source": [
    "ds = pydicom.read_file(image_fps[0]) # read dicom image from filepath \n",
    "image = ds.pixel_array # get image array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_uuid": "0ef68a41cf1a5e842e86a219b6392e3695004720",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 647
    },
    "colab_type": "code",
    "id": "81lovwF2Ro5R",
    "outputId": "e2263fe2-1a32-432a-ec75-b9220a24e697"
   },
   "outputs": [],
   "source": [
    "# show dicom fields \n",
    "ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_uuid": "74277ae9af4a3b044e62b664d10d76b23848bb43",
    "colab": {},
    "colab_type": "code",
    "id": "gYNSd1AhRqOV"
   },
   "outputs": [],
   "source": [
    "# Original DICOM image size: 1024 x 1024\n",
    "ORIG_SIZE = 1024"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "_uuid": "6563bbca143e4bceb1ea850714d7b43bb1e1178d",
    "colab_type": "text",
    "id": "4FlRu8ML-ceg"
   },
   "source": [
    "### Split the data into training and validation datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_uuid": "6175c72e73639e3190e127f67783988eadced9ba",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "colab_type": "code",
    "id": "7jByVCZt-ZOC",
    "outputId": "f1aa267d-7530-4620-ffc5-2f7aa39083bb"
   },
   "outputs": [],
   "source": [
    "# split dataset into training vs. validation dataset \n",
    "# split ratio is set to 0.9 vs. 0.1 (train vs. validation, respectively)\n",
    "image_fps_list = list(image_fps)\n",
    "random.seed(42)\n",
    "random.shuffle(image_fps_list)\n",
    "\n",
    "val_size = 1500\n",
    "image_fps_val = image_fps_list[:val_size]\n",
    "image_fps_train = image_fps_list[val_size:]\n",
    "\n",
    "print(len(image_fps_train), len(image_fps_val))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "_uuid": "a5143c19dc22bc00d318a3b28cb7e13c7fbacc8a",
    "colab_type": "text",
    "id": "9KUvacUbgiEX"
   },
   "source": [
    "### Create and prepare the training dataset using the DetectorDataset class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_uuid": "86c3333d4dfb8b7d00ce1f401693d0df4e6254e1",
    "colab": {},
    "colab_type": "code",
    "id": "jwMkhotP0yFf"
   },
   "outputs": [],
   "source": [
    "# prepare the training dataset\n",
    "dataset_train = DetectorDataset(image_fps_train, image_annotations, ORIG_SIZE, ORIG_SIZE)\n",
    "dataset_train.prepare()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "_uuid": "4f69286e6b0b640827a3de166326d157a6d86668",
    "colab_type": "text",
    "id": "wPDQ9EVDgxa6"
   },
   "source": [
    "### Let's look at a sample annotation. We see a bounding box with (x, y) of the the top left corner as well as the width and height."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_uuid": "93da5a58731ad483a4bd2b20543f2b1df4b8ad74",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 251
    },
    "colab_type": "code",
    "id": "0xEc47Jz59x5",
    "outputId": "129edfbc-cf9d-46c7-b569-d804a50cd12d"
   },
   "outputs": [],
   "source": [
    "# Show annotation(s) for a DICOM image \n",
    "test_fp = random.choice(image_fps_train)\n",
    "image_annotations[test_fp]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_uuid": "313347d838fa8321a714858c8073f98c50c5be26",
    "colab": {},
    "colab_type": "code",
    "id": "K1TkWuGP0yHl"
   },
   "outputs": [],
   "source": [
    "# prepare the validation dataset\n",
    "dataset_val = DetectorDataset(image_fps_val, image_annotations, ORIG_SIZE, ORIG_SIZE)\n",
    "dataset_val.prepare()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "_uuid": "600a8135d4e382f62797d69e9358f5697873c8f9",
    "colab_type": "text",
    "id": "pEXEt8fygWuC"
   },
   "source": [
    "### Display a random image with bounding boxes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_uuid": "491b78ec96d28fcdbbf8e2d7f9320a05d64c9249",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 359
    },
    "colab_type": "code",
    "id": "4xwsrf9G1lHR",
    "outputId": "a13386d3-a918-41fe-8824-13625c9d7b08"
   },
   "outputs": [],
   "source": [
    "# Load and display random sample and their bounding boxes\n",
    "\n",
    "class_ids = [0]\n",
    "while class_ids[0] == 0:  ## look for a mask\n",
    "    image_id = random.choice(dataset_train.image_ids)\n",
    "    image_fp = dataset_train.image_reference(image_id)\n",
    "    image = dataset_train.load_image(image_id)\n",
    "    mask, class_ids = dataset_train.load_mask(image_id)\n",
    "\n",
    "print(image.shape)\n",
    "\n",
    "plt.figure(figsize=(10, 10))\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.imshow(image)\n",
    "plt.axis('off')\n",
    "\n",
    "plt.subplot(1, 2, 2)\n",
    "masked = np.zeros(image.shape[:2])\n",
    "for i in range(mask.shape[2]):\n",
    "    masked += image[:, :, 0] * mask[:, :, i]\n",
    "plt.imshow(masked, cmap='gray')\n",
    "plt.axis('off')\n",
    "\n",
    "print(image_fp)\n",
    "print(class_ids)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "_uuid": "342b6008873fe7a6a0870a712ee47a87f0d2828d",
    "colab_type": "text",
    "id": "ustAIH78hZI_"
   },
   "source": [
    "### Image Augmentation. Try finetuning some variables to custom values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_uuid": "4ab9d6086ce611a46f189c047956c43b29783e6d",
    "colab": {},
    "colab_type": "code",
    "id": "STZnQTE61lME"
   },
   "outputs": [],
   "source": [
    "# Image augmentation (light but constant)\n",
    "augmentation = iaa.Sequential([\n",
    "    iaa.OneOf([ ## geometric transform\n",
    "        iaa.Affine(\n",
    "            scale={\"x\": (0.98, 1.02), \"y\": (0.98, 1.02)},\n",
    "            translate_percent={\"x\": (-0.02, 0.02), \"y\": (-0.04, 0.04)},\n",
    "            rotate=(-2, 2),\n",
    "            shear=(-1, 1),\n",
    "        ),\n",
    "        iaa.PiecewiseAffine(scale=(0.001, 0.025)),\n",
    "    ]),\n",
    "    iaa.OneOf([ ## brightness or contrast\n",
    "        iaa.Multiply((0.9, 1.1)),\n",
    "        iaa.ContrastNormalization((0.9, 1.1)),\n",
    "    ]),\n",
    "    iaa.OneOf([ ## blur or sharpen\n",
    "        iaa.GaussianBlur(sigma=(0.0, 0.1)),\n",
    "        iaa.Sharpen(alpha=(0.0, 0.1)),\n",
    "    ]),\n",
    "])\n",
    "\n",
    "# test on the same image as above\n",
    "imggrid = augmentation.draw_grid(image[:, :, 0], cols=5, rows=2)\n",
    "plt.figure(figsize=(30, 12))\n",
    "_ = plt.imshow(imggrid[:, :, 0], cmap='gray')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "_uuid": "7e65d2cecb283f446f34cdde19b663a8a8e9590f",
    "colab_type": "text",
    "id": "M4kt7LKuc78e"
   },
   "source": [
    "### Now it's time to train the model. Note that training even a basic model can take a few hours. \n",
    "\n",
    "Note: the following model is for demonstration purpose only. We have limited the training to one epoch, and have set nominal values for the Detector Configuration to reduce run-time. \n",
    "\n",
    "- dataset_train and dataset_val are derived from DetectorDataset \n",
    "- DetectorDataset loads images from image filenames and  masks from the annotation data\n",
    "- model is Mask-RCNN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_uuid": "138d6197fc8dce9f1f8a7b5a6c27aa2069698e03"
   },
   "outputs": [],
   "source": [
    "model = modellib.MaskRCNN(mode='training', config=config, model_dir=ROOT_DIR)\n",
    "\n",
    "# Exclude the last layers because they require a matching\n",
    "# number of classes\n",
    "model.load_weights(COCO_WEIGHTS_PATH, by_name=True, exclude=[\n",
    "    \"mrcnn_class_logits\", \"mrcnn_bbox_fc\",\n",
    "    \"mrcnn_bbox\", \"mrcnn_mask\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_uuid": "64cce2581ffdb8c2b1cb07948ada4a93f64874b0",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 2575
    },
    "colab_type": "code",
    "id": "RVgNhHjl1lOS",
    "outputId": "2cba9efc-eeea-472d-d155-3c3d856585bf"
   },
   "outputs": [],
   "source": [
    "LEARNING_RATE = 0.003\n",
    "\n",
    "# Train Mask-RCNN Model \n",
    "import warnings \n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_uuid": "cf339a499519d174bcdf2311a1802f0e3acb1758"
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "## train heads with higher lr to speedup the learning\n",
    "model.train(dataset_train, dataset_val,\n",
    "            learning_rate=LEARNING_RATE*2,\n",
    "            epochs=2,\n",
    "            layers='heads',\n",
    "            augmentation=None)  ## no need to augment yet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_uuid": "8004790d27f041793562e994bbe95edf67f8978b"
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "model.train(dataset_train, dataset_val,\n",
    "            learning_rate=LEARNING_RATE,\n",
    "            epochs=6,\n",
    "            layers='all',\n",
    "            augmentation=augmentation)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_uuid": "ccea214a520c686735e138f64977dcd7f3e3330a"
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "model.train(dataset_train, dataset_val,\n",
    "            learning_rate=LEARNING_RATE/2,\n",
    "            epochs=16,\n",
    "            layers='all',\n",
    "            augmentation=augmentation)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "model.train(dataset_train, dataset_val,\n",
    "            learning_rate=LEARNING_RATE/4,\n",
    "            epochs=16,\n",
    "            layers='all',\n",
    "            augmentation=augmentation)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_uuid": "fb3b69242b91dcc49697ff076ceeb957347372e1"
   },
   "outputs": [],
   "source": [
    "history = model.history\n",
    "plt.figure(figsize=(16,5))\n",
    "\n",
    "plt.subplot(1)\n",
    "plt.plot(history.epoch, history.history[\"loss\"], label=\"Train loss\")\n",
    "plt.plot(history.epoch, history.history[\"val_loss\"], label=\"Valid loss\")\n",
    "plt.legend()\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_uuid": "db5c10d3f7da099e5751a04a6e6d49819882ecd4",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "colab_type": "code",
    "id": "eraRlzgPmmIZ",
    "outputId": "de9e688c-ba4f-4b62-f842-dbcf00ce397c"
   },
   "outputs": [],
   "source": [
    "# select trained model \n",
    "dir_names = next(os.walk(model.model_dir))[1]\n",
    "key = config.NAME.lower()\n",
    "dir_names = filter(lambda f: f.startswith(key), dir_names)\n",
    "dir_names = sorted(dir_names)\n",
    "\n",
    "if not dir_names:\n",
    "    import errno\n",
    "    raise FileNotFoundError(\n",
    "        errno.ENOENT,\n",
    "        \"Could not find model directory under {}\".format(self.model_dir))\n",
    "    \n",
    "fps = []\n",
    "# Pick last directory\n",
    "for d in dir_names: \n",
    "    dir_name = os.path.join(model.model_dir, d)\n",
    "    # Find the last checkpoint\n",
    "    checkpoints = next(os.walk(dir_name))[2]\n",
    "    checkpoints = filter(lambda f: f.startswith(\"mask_rcnn\"), checkpoints)\n",
    "    checkpoints = sorted(checkpoints)\n",
    "    if not checkpoints:\n",
    "        print('No weight files in {}'.format(dir_name))\n",
    "    else:\n",
    "        checkpoint = os.path.join(dir_name, checkpoints[-1])\n",
    "        fps.append(checkpoint)\n",
    "\n",
    "model_path = sorted(fps)[-1]\n",
    "print('Found model {}'.format(model_path))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_uuid": "52138636b2ae5bf444bba808518cd8313bde65cd",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 53
    },
    "colab_type": "code",
    "id": "TgpT9AzC2Bgz",
    "outputId": "60f5a175-4666-497d-b4e8-0bdab39a92d0"
   },
   "outputs": [],
   "source": [
    "class InferenceConfig(DetectorConfig):\n",
    "    GPU_COUNT = 1\n",
    "    IMAGES_PER_GPU = 1\n",
    "\n",
    "inference_config = InferenceConfig()\n",
    "\n",
    "# Recreate the model in inference mode\n",
    "model = modellib.MaskRCNN(mode='inference', \n",
    "                          config=inference_config,\n",
    "                          model_dir=ROOT_DIR)\n",
    "\n",
    "# Load trained weights (fill in path to trained weights here)\n",
    "assert model_path != \"\", \"Provide path to trained weights\"\n",
    "print(\"Loading weights from \", model_path)\n",
    "model.load_weights(model_path, by_name=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_uuid": "e13c61bee23b791c61ecf1256f7512295cd4d9ab",
    "colab": {},
    "colab_type": "code",
    "id": "9mTBig7D2BjU"
   },
   "outputs": [],
   "source": [
    "# set color for class\n",
    "def get_colors_for_class_ids(class_ids):\n",
    "    colors = []\n",
    "    for class_id in class_ids:\n",
    "        if class_id == 1:\n",
    "            colors.append((.941, .204, .204))\n",
    "    return colors"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "_uuid": "f99fbd3f31ff1a2bd66764835c9b646375364598",
    "colab_type": "text",
    "id": "A8EiL2LOiCr_"
   },
   "source": [
    "### How does the predicted box compared to the expected value? Let's use the validation dataset to check. \n",
    "\n",
    "Note that we trained only one epoch for **demonstration purposes ONLY**. You might be able to improve performance running more epochs. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_uuid": "186412199e25b98719f71cfe5e8869abcce516c4",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1394
    },
    "colab_type": "code",
    "id": "irheTbrW2Bl0",
    "outputId": "56041ad4-173d-45ab-af67-f54e8333511e"
   },
   "outputs": [],
   "source": [
    "# Show few example of ground truth vs. predictions on the validation dataset \n",
    "dataset = dataset_val\n",
    "fig = plt.figure(figsize=(10, 30))\n",
    "\n",
    "for i in range(6):\n",
    "\n",
    "    image_id = random.choice(dataset.image_ids)\n",
    "    \n",
    "    original_image, image_meta, gt_class_id, gt_bbox, gt_mask =\\\n",
    "        modellib.load_image_gt(dataset_val, inference_config, \n",
    "                               image_id, use_mini_mask=False)\n",
    "    \n",
    "    print(original_image.shape)\n",
    "    plt.subplot(6, 2, 2*i + 1)\n",
    "    visualize.display_instances(original_image, gt_bbox, gt_mask, gt_class_id, \n",
    "                                dataset.class_names,\n",
    "                                colors=get_colors_for_class_ids(gt_class_id), ax=fig.axes[-1])\n",
    "    \n",
    "    plt.subplot(6, 2, 2*i + 2)\n",
    "    results = model.detect([original_image]) #, verbose=1)\n",
    "    r = results[0]\n",
    "    visualize.display_instances(original_image, r['rois'], r['masks'], r['class_ids'], \n",
    "                                dataset.class_names, r['scores'], \n",
    "                                colors=get_colors_for_class_ids(r['class_ids']), ax=fig.axes[-1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_uuid": "fd9f53fa319a425693e07fe4898ddeeaa5d07f99",
    "colab": {},
    "colab_type": "code",
    "id": "qRWBVJKYNdWM"
   },
   "outputs": [],
   "source": [
    "# Get filenames of test dataset DICOM images\n",
    "test_image_fps = get_dicom_fps(test_dicom_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "_uuid": "164e18701a830bc6c42a791feea13549de37289b",
    "colab_type": "text",
    "id": "WcV1cL_aiSc4"
   },
   "source": [
    "### Final steps - Create the submission file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_uuid": "4a5c0c6134408ddbf5a34496d7e9d7be5692e9a1",
    "colab": {},
    "colab_type": "code",
    "id": "C6UWVrbM2Bob"
   },
   "outputs": [],
   "source": [
    "# Make predictions on test images, write out sample submission\n",
    "def predict(image_fps, filepath='submission.csv', min_conf=0.95):\n",
    "    # assume square image\n",
    "    resize_factor = ORIG_SIZE / config.IMAGE_SHAPE[0]\n",
    "    #resize_factor = ORIG_SIZE\n",
    "    with open(filepath, 'w') as file:\n",
    "        file.write(\"patientId,PredictionString\\n\")\n",
    "\n",
    "        for image_id in tqdm(image_fps):\n",
    "            ds = pydicom.read_file(image_id)\n",
    "            image = ds.pixel_array\n",
    "            # If grayscale. Convert to RGB for consistency.\n",
    "            if len(image.shape) != 3 or image.shape[2] != 3:\n",
    "                image = np.stack((image,) * 3, -1)\n",
    "            image, window, scale, padding, crop = utils.resize_image(\n",
    "                image,\n",
    "                min_dim=config.IMAGE_MIN_DIM,\n",
    "                min_scale=config.IMAGE_MIN_SCALE,\n",
    "                max_dim=config.IMAGE_MAX_DIM,\n",
    "                mode=config.IMAGE_RESIZE_MODE)\n",
    "\n",
    "            patient_id = os.path.splitext(os.path.basename(image_id))[0]\n",
    "\n",
    "            results = model.detect([image])\n",
    "            r = results[0]\n",
    "\n",
    "            out_str = \"\"\n",
    "            out_str += patient_id\n",
    "            out_str += \",\"\n",
    "            assert( len(r['rois']) == len(r['class_ids']) == len(r['scores']) )\n",
    "            if len(r['rois']) == 0:\n",
    "                pass\n",
    "            else:\n",
    "                num_instances = len(r['rois'])\n",
    "\n",
    "                for i in range(num_instances):\n",
    "                    if r['scores'][i] > min_conf:\n",
    "                        out_str += ' '\n",
    "                        out_str += str(round(r['scores'][i], 2))\n",
    "                        out_str += ' '\n",
    "\n",
    "                        # x1, y1, width, height\n",
    "                        x1 = r['rois'][i][1]\n",
    "                        y1 = r['rois'][i][0]\n",
    "                        width = r['rois'][i][3] - x1\n",
    "                        height = r['rois'][i][2] - y1\n",
    "                        bboxes_str = \"{} {} {} {}\".format(x1*resize_factor, y1*resize_factor, \\\n",
    "                                                           width*resize_factor, height*resize_factor)\n",
    "                        out_str += bboxes_str\n",
    "\n",
    "            file.write(out_str+\"\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_uuid": "0406e7f5aaa4867782c4f9c064f90bba386128e7",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "colab_type": "code",
    "id": "C5cBpNka2Bsv",
    "outputId": "a2af9176-d9d6-49f6-f22a-5a1c455d144f"
   },
   "outputs": [],
   "source": [
    "submission_fp = os.path.join(ROOT_DIR, 'submission.csv')\n",
    "predict(test_image_fps, filepath=submission_fp)\n",
    "print(submission_fp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_uuid": "3fd8d178fc51ef0bca94fbb3f423160f08a77edc",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1599
    },
    "colab_type": "code",
    "id": "_BjPE_Ee9rbA",
    "outputId": "67b5f053-112b-494a-9ab3-d017bfb440c2"
   },
   "outputs": [],
   "source": [
    "output = pd.read_csv(submission_fp, names=['patientId', 'PredictionString'])\n",
    "output.head(60)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_uuid": "ea110f197abc2acb1c3435383f7259079dc0eb0e"
   },
   "outputs": [],
   "source": [
    "# show a few test image detection example\n",
    "def visualize(): \n",
    "    image_id = random.choice(test_image_fps)\n",
    "    ds = pydicom.read_file(image_id)\n",
    "    \n",
    "    # original image \n",
    "    image = ds.pixel_array\n",
    "    \n",
    "    # assume square image \n",
    "    resize_factor = ORIG_SIZE / config.IMAGE_SHAPE[0]\n",
    "    \n",
    "    # If grayscale. Convert to RGB for consistency.\n",
    "    if len(image.shape) != 3 or image.shape[2] != 3:\n",
    "        image = np.stack((image,) * 3, -1) \n",
    "    resized_image, window, scale, padding, crop = utils.resize_image(\n",
    "        image,\n",
    "        min_dim=config.IMAGE_MIN_DIM,\n",
    "        min_scale=config.IMAGE_MIN_SCALE,\n",
    "        max_dim=config.IMAGE_MAX_DIM,\n",
    "        mode=config.IMAGE_RESIZE_MODE)\n",
    "\n",
    "    patient_id = os.path.splitext(os.path.basename(image_id))[0]\n",
    "    print(patient_id)\n",
    "\n",
    "    results = model.detect([resized_image])\n",
    "    r = results[0]\n",
    "    for bbox in r['rois']: \n",
    "        print(bbox)\n",
    "        x1 = int(bbox[1] * resize_factor)\n",
    "        y1 = int(bbox[0] * resize_factor)\n",
    "        x2 = int(bbox[3] * resize_factor)\n",
    "        y2 = int(bbox[2]  * resize_factor)\n",
    "        cv2.rectangle(image, (x1,y1), (x2,y2), (77, 255, 9), 3, 1)\n",
    "        width = x2 - x1 \n",
    "        height = y2 - y1 \n",
    "        print(\"x {} y {} h {} w {}\".format(x1, y1, width, height))\n",
    "    plt.figure() \n",
    "    plt.imshow(image, cmap=plt.cm.gist_gray)\n",
    "\n",
    "visualize()\n",
    "visualize()\n",
    "visualize()\n",
    "visualize()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_uuid": "835a15c9d018acd5deb16e9e02f9b765f68d0e78"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "lesson-3-rsna-pneumonia-detection-challenge-kaggle",
   "provenance": [],
   "version": "0.3.2"
  },
  "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}