yolo_480x480_16.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "igEd7obtcxR2"
   },
   "source": [
    "## Approach\n",
    "\n",
    "* This is a simplied version of YOLO. The input image is 512x512. YOLO was intended to train on images which all contained at least one object, most multiple objects. The architecture and loss functions have been adapted for this case, where not every image contains an object and there are at most 4 ROIs per image. \n",
    "* Since the maximum number of ROIs in each image is 4 we decided to use a 2x2 grid to simplify the network. In the case where multiple ROIs occur in the same cell we combine the ROIs into a single ROI by expanding the box to include both. Each cell outputs five values:\n",
    "    * The confidence that there is pneumonia present\n",
    "    * The x, y, w, and h of the bounding box\n",
    "    * We have removed the entire classification section, using the confidence to indicate whether there is an ROI in the cell instead.\n",
    "* A sigmoid is applied to the output of the network to result in values between 0 and 1    \n",
    "* The x and y coordinates are offsets from the upper left corner of each cell, the w and h are percentage of the total width.\n",
    "* The loss function is based on YOLO with some differences:\n",
    "    * The weights of the components have been updated\n",
    "    * The classification loss has been removed.\n",
    "    * The \"objectness\" loss of YOLO tries to make the confidence match the actual IOU. Since our model only outputs one box per cell, this doesn't really make any sense and ends up driving the confidence down to the IOU. We replace this by an IOU loss which subtracts the IOU for that cell from 1.\n",
    "* Some of the hyperparameters are tweaked during training. We train in sets of 10 epochs and evaluate which hyperparameters may need to be adjusted in between.\n",
    "    * The weight for the positive class in the confidence often needs to be tweaked to focus the model on predicting more or less positive cells. \n",
    "    * The weights of the various components of the loss function are also tweaked as needed to focus the model on minimizing certain losses over others.\n",
    "\n",
    "## Network\n",
    "\n",
    "* The network consists of a number of residual blocks with convolutions and downsampling blocks with max pooling.\n",
    "* There are two dilated convolutions at the end of the network to provide context.\n",
    "* The dilated convolutions are followed by two strided convolutions which downsize to 2x2.\n",
    "* Finally there are a series of 1x1 convolutions which output a 2x2x5 tensor.\n",
    "* We are using an Adam optimizer with gradient clipping to avoid exploding gradients, which had been a problem.\n",
    "\n",
    "## Predictions\n",
    "* To generate our predictions we loop through each cell of the output\n",
    "* We unnormalize the output to get the actual values\n",
    "* If the confidence is greater than 0.5 and the box has width and height we append it to the list of candidates.\n",
    "* We apply non-max suppression to remove duplicates.\n",
    "* The remaining boxes are concatenated onto the output string.\n",
    "\n",
    "**Change Log:**\n",
    "* v3 - changing output to 8x8 grid from 16x16; changed model to downsample one more time; adjusted network accordingly. \n",
    "* v4 - changed output to 4x4 grid, no image has more than 3 ROIs so this may work better? \n",
    "    * Using center point of ROI to predict instead of upper left corner.\n",
    "* v5 - We only calculate MSE loss for boxes with a confidence over 0.5 or actual truth since we don't care about predictions for boxes that are not ROIs. This will prevent the network from being constrained by outputting 0s for boxes that don't exist.\n",
    "* v6 - centering input data so maybe bboxes can be output more accurately? Also centering the image\n",
    "* v6.2 - labels have ROI centered in center of cell by default instead of mean location.\n",
    "* v8 - using custom loss function based on YOLO loss. Set default height and width to 1 px because 0 sent the gradients to -inf which screwed everything up.\n",
    "* v10 - using 2x2 grid as output\n",
    "* v12 - changing layout of model slightly - replaced final pools with convs with strides; reduced number of params; ensure that there is at least one positive image per batch.\n",
    "* v13 - tweaking architecture slightly.\n",
    "* v14 - going to 5x5 grid\n",
    "* v15 - downsizing with average and max pool followed by 1x1 convolutions\n",
    "* v16 - minor tweaks to architecture, trying 3x3 output grid. If the center of an ROI is near the cell border also assign it to the neighboring cell. We also expand ROIs into cells which they cover more than 50% of, adjusting the height and width appropriately."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "LHncOor-cxSS"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import csv\n",
    "import random\n",
    "import pydicom\n",
    "import numpy as np\n",
    "import pickle\n",
    "import pandas as pd\n",
    "from skimage import measure\n",
    "from skimage.transform import resize\n",
    "import datetime\n",
    "import math\n",
    "\n",
    "import tensorflow as tf\n",
    "from tensorflow import keras\n",
    "\n",
    "from matplotlib import pyplot as plt\n",
    "import matplotlib.patches as patches"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Malisiewicz et al.\n",
    "def non_max_suppression_fast(boxes, overlapThresh=0.3):\n",
    "    # if there are no boxes, return an empty list\n",
    "    if len(boxes) == 0:\n",
    "        return []\n",
    " \n",
    "    # if the bounding boxes integers, convert them to floats --\n",
    "    # this is important since we'll be doing a bunch of divisions\n",
    "    if boxes.dtype.kind == \"i\":\n",
    "        boxes = boxes.astype(\"float\")\n",
    " \n",
    "    # initialize the list of picked indexes\t\n",
    "    pick = []\n",
    " \n",
    "    # grab the coordinates of the bounding boxes\n",
    "    x1 = boxes[:,0]\n",
    "    y1 = boxes[:,1]\n",
    "    w = boxes[:,2]\n",
    "    h = boxes[:,3]\n",
    " \n",
    "    x2 = x1 + w\n",
    "    y2 = y1 + h\n",
    "\n",
    "    # compute the area of the bounding boxes and sort the bounding\n",
    "    # boxes by the bottom-right y-coordinate of the bounding box\n",
    "    area = (w + 1) * (h + 1)\n",
    "    idxs = np.argsort(y2)\n",
    " \n",
    "    # keep looping while some indexes still remain in the indexes\n",
    "    # list\n",
    "    while len(idxs) > 0:\n",
    "        # grab the last index in the indexes list and add the\n",
    "        # index value to the list of picked indexes\n",
    "        last = len(idxs) - 1\n",
    "        i = idxs[last]\n",
    "        pick.append(i)\n",
    " \n",
    "        # find the largest (x, y) coordinates for the start of\n",
    "        # the bounding box and the smallest (x, y) coordinates\n",
    "        # for the end of the bounding box\n",
    "        xx1 = np.maximum(x1[i], x1[idxs[:last]])\n",
    "        yy1 = np.maximum(y1[i], y1[idxs[:last]])\n",
    "        xx2 = np.minimum(x2[i], x2[idxs[:last]])\n",
    "        yy2 = np.minimum(y2[i], y2[idxs[:last]])\n",
    " \n",
    "        # compute the width and height of the bounding box\n",
    "        w = np.maximum(0, xx2 - xx1 + 1)\n",
    "        h = np.maximum(0, yy2 - yy1 + 1)\n",
    " \n",
    "        # compute the ratio of overlap\n",
    "        overlap = (w * h) / area[idxs[:last]]\n",
    " \n",
    "        # delete all indexes from the index list that have\n",
    "        idxs = np.delete(idxs, np.concatenate(([last],\n",
    "            np.where(overlap > overlapThresh)[0])))\n",
    " \n",
    "    # return only the bounding boxes that were picked using the\n",
    "    # integer data type\n",
    "    return boxes[pick].astype(\"int\")\n",
    "\n",
    "# get actual x and y values from sigmoid output and what cell they are in\n",
    "def unnorm(val, idx, cell_size=256):\n",
    "    x = (val * cell_size) + (cell_size * idx)\n",
    "    return x\n",
    "\n",
    "# sigmoid in numpy, with limit to avoid nans                             \n",
    "def sigmoid(x):\n",
    "    # to avoid NaNs set a lower floor on x values\n",
    "    y = np.maximum(x, -700)\n",
    "    y = np.minimum(y, 700)\n",
    "    return 1 / (1 + np.exp(-y))    \n",
    "\n",
    "# adjust contrast of image\n",
    "def change_contrast(img, contrast_factor):\n",
    "    mean = np.mean(img)\n",
    "    img = (img - mean) * contrast_factor + mean\n",
    "    return img"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "hE3jDmgudDDP"
   },
   "outputs": [],
   "source": [
    "# enter your Kaggle credentionals here\n",
    "os.environ['KAGGLE_USERNAME']=\"skooch\"\n",
    "os.environ['KAGGLE_KEY']=\"42f8a02ee92cc773d1dbe66565673ad3\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "sSWY-U2TcxSn"
   },
   "source": [
    "# Load pneumonia locations\n",
    "\n",
    "Table contains [filename : pneumonia location] pairs per row. \n",
    "* If a filename contains multiple pneumonia, the table contains multiple rows with the same filename but different pneumonia locations. \n",
    "* If a filename contains no pneumonia it contains a single row with an empty pneumonia location.\n",
    "\n",
    "The code below loads the table and transforms it into a dictionary. \n",
    "* The dictionary uses the filename as key and a list of pneumonia locations in that filename as value. \n",
    "* If a filename is not present in the dictionary it means that it contains no pneumonia."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "m_0ZGQyLdQS8",
    "outputId": "4848c0ce-4d75-4781-d01f-91cbd7626d77"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train dir: ./stage_1_train_images\n"
     ]
    }
   ],
   "source": [
    "ROOT_DIR = \"./\"\n",
    "\n",
    "train_dicom_dir = os.path.join(ROOT_DIR, 'stage_1_train_images')\n",
    "test_dicom_dir = os.path.join(ROOT_DIR, 'stage_1_test_images')\n",
    "print(\"Train dir:\", train_dicom_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "hopsRmICcxSs"
   },
   "outputs": [],
   "source": [
    "with open('yolo_labels_centered_5x5_10a.p', 'rb') as handle:\n",
    "    pneumonia_locations = pickle.load(handle)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "UoB6GBhIcxS6"
   },
   "source": [
    "# Load filenames"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 51
    },
    "colab_type": "code",
    "id": "b_iB-IftcxTE",
    "outputId": "89a05e05-f2d9-45fc-d1ef-1a0cc289f98d"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "n train samples 23116\n",
      "n valid samples 2568\n"
     ]
    }
   ],
   "source": [
    "random.seed(17)\n",
    "\n",
    "# load and shuffle filenames\n",
    "folder = './stage_1_train_images'\n",
    "filenames = os.listdir(folder)\n",
    "random.shuffle(filenames)\n",
    "# split into train and validation filenames\n",
    "n_valid_samples = int(len(filenames) * 0.1)\n",
    "train_filenames = filenames[n_valid_samples:]\n",
    "valid_filenames = filenames[:n_valid_samples]\n",
    "print('n train samples', len(train_filenames))\n",
    "print('n valid samples', len(valid_filenames))\n",
    "n_train_samples = len(filenames) - n_valid_samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "positive_images = []\n",
    "\n",
    "for filename in pneumonia_locations:\n",
    "    label = pneumonia_locations[filename][...,0]\n",
    "    if np.max(label) > 1e-6:\n",
    "        if filename + \".dcm\" in train_filenames:\n",
    "            positive_images.append(filename + \".dcm\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "lorV7ZmDcxTa"
   },
   "source": [
    " # Data generator\n",
    "\n",
    "The dataset is too large to fit into memory, so we need to create a generator that loads data on the fly.\n",
    "\n",
    "* The generator takes in some filenames, batch_size and other parameters.\n",
    "\n",
    "* The generator outputs a random batch of numpy images and numpy masks.\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "dFbLdFxVehhB"
   },
   "outputs": [],
   "source": [
    "BATCH_SIZE = 16\n",
    "IMAGE_SIZE = 480\n",
    "CHECKPOINT_PATH = \"yolo16_2_480.h5\"\n",
    "GRID_SIZE = 5\n",
    "CELL_SIZE = 1024 / GRID_SIZE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "TueY1bVlcxTg"
   },
   "outputs": [],
   "source": [
    "class generator(keras.utils.Sequence):\n",
    "    \n",
    "    def __init__(self, folder, filenames, pneumonia_locations=None, batch_size=BATCH_SIZE, image_size=IMAGE_SIZE, shuffle=True, augment=False, predict=False):\n",
    "        self.folder = folder\n",
    "        self.filenames = filenames\n",
    "        self.pneumonia_locations = pneumonia_locations\n",
    "        self.batch_size = batch_size\n",
    "        self.image_size = image_size\n",
    "        self.shuffle = shuffle\n",
    "        self.augment = augment\n",
    "        self.predict = predict\n",
    "        self.on_epoch_end()\n",
    "    \n",
    "    def augment_imgs(self, img, confs, boxes):\n",
    "        # flip the image half the time\n",
    "        if random.random() > 0.5:\n",
    "            img = np.fliplr(img)\n",
    "            \n",
    "            # update our x coords\n",
    "            mask = (confs != 0)\n",
    "            \n",
    "            # flip\n",
    "            boxes[mask, 0] = (1 - boxes[mask,0])\n",
    "            \n",
    "            # flip our boxes lr on axis 0\n",
    "            boxes = np.flip(boxes, axis=1)   \n",
    "            \n",
    "            # flip the confidences lr as well\n",
    "            confs = np.flip(confs, axis=1)            \n",
    "            \n",
    "        ## small random shifts\n",
    "        h_offset = np.random.randint(low=0, high=5)\n",
    "        v_offset = np.random.randint(low=0, high=5)\n",
    "\n",
    "        # crop the images\n",
    "        img = img[v_offset:,h_offset:]\n",
    "           \n",
    "        ## adjust contrast half the time\n",
    "        if random.random() > 0.5:  \n",
    "            # generate a random contrast adjustment\n",
    "            contrast_factor = np.random.normal(loc=1.0, scale=0.10)\n",
    "            \n",
    "            # put some limits on the contrast\n",
    "            contrast_factor = np.minimum(contrast_factor, 1.25)\n",
    "            contrast_factor = np.maximum(contrast_factor, 0.75)\n",
    "            \n",
    "            # adjust the image\n",
    "            img = change_contrast(img, contrast_factor)\n",
    "            \n",
    "        return img, confs, boxes\n",
    "            \n",
    "            \n",
    "    def __load__(self, filename):\n",
    "        # load dicom file as numpy array\n",
    "        img = pydicom.dcmread(os.path.join(self.folder, filename)).pixel_array\n",
    "        # get filename without extension\n",
    "        filename = filename.split('.')[0]\n",
    "        label = pneumonia_locations[filename].copy()\n",
    "        \n",
    "        # remove the confidence and bboxes because they will be flipped separately\n",
    "        # round the confidences since some are 0.99 or maybe even 1e-6\n",
    "        confs = np.round(label[:,:,0])\n",
    "        boxes = label[:,:,1:]\n",
    "        \n",
    "        ## augment the data with flips, small shifts and contrast adjustment\n",
    "        if self.augment:\n",
    "            img, confs, boxes = self.augment_imgs(img, confs, boxes)\n",
    "            \n",
    "        # resize both image and mask\n",
    "        img = resize(img, (self.image_size, self.image_size), mode='reflect')\n",
    "        \n",
    "        # scale and center the image\n",
    "        img = (img - np.mean(img)) / np.max(img)\n",
    "        \n",
    "        boxes = np.concatenate([confs.reshape((GRID_SIZE,GRID_SIZE,1)), boxes], axis=2)\n",
    "        \n",
    "        # add trailing channel dimension\n",
    "        img = np.expand_dims(img, -1)\n",
    "        return img, boxes\n",
    "    \n",
    "    def __loadpredict__(self, filename):\n",
    "        # load dicom file as numpy array\n",
    "        img = pydicom.dcmread(os.path.join(self.folder, filename)).pixel_array\n",
    "        # resize image\n",
    "        img = resize(img, (self.image_size, self.image_size), mode='reflect')\n",
    "        # add trailing channel dimension\n",
    "        img = np.expand_dims(img, -1)\n",
    "        return img\n",
    "        \n",
    "    def __getitem__(self, index):\n",
    "        # select batch\n",
    "        filenames = self.filenames[index*self.batch_size:(index+1)*self.batch_size]\n",
    "        # predict mode: return images and filenames\n",
    "        if self.predict:\n",
    "            # load files\n",
    "            imgs = [self.__loadpredict__(filename) for filename in filenames]\n",
    "            # create numpy batch\n",
    "            imgs = np.array(imgs)\n",
    "            return imgs, filenames\n",
    "        # train mode: return images and masks\n",
    "        else:\n",
    "            # load files\n",
    "            items = [self.__load__(filename) for filename in filenames]\n",
    "            \n",
    "            # unzip images and masks\n",
    "            imgs, bboxes = zip(*items)\n",
    "            \n",
    "            # create numpy batch\n",
    "            imgs = np.array(imgs)\n",
    "            bboxes = np.array(bboxes)\n",
    "            \n",
    "            # make sure there is at least one positive image in the batch\n",
    "            pos = np.max(bboxes[:,:,:,0])\n",
    "            if pos < 1:\n",
    "                # pick a random positive image\n",
    "                filename = np.random.choice(positive_images)\n",
    "                img, label = self.__load__(filename)\n",
    "                \n",
    "                # add the positive image to our batch\n",
    "                imgs[-1] = img\n",
    "                bboxes[-1] = label\n",
    "                \n",
    "            labels = bboxes\n",
    "            return imgs, labels\n",
    "        \n",
    "    def on_epoch_end(self):\n",
    "        if self.shuffle:\n",
    "            random.shuffle(self.filenames)\n",
    "        \n",
    "    def __len__(self):\n",
    "        if self.predict:\n",
    "            # return everything\n",
    "            return int(np.ceil(len(self.filenames) / self.batch_size))\n",
    "        else:\n",
    "            # return full batches only\n",
    "            return int(len(self.filenames) / self.batch_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# train_gen = generator(folder, train_filenames, pneumonia_locations, batch_size=BATCH_SIZE, image_size=IMAGE_SIZE, shuffle=True, augment=True, predict=False)\n",
    "\n",
    "# counter = 0\n",
    "# for imgs, labels in train_gen:\n",
    "#     for label in labels:\n",
    "#         for i in range(4):\n",
    "#             for j in range(4):\n",
    "#                 if label[i,j,0] == 1:\n",
    "#                     if label[i,j,4] ==  0.001:\n",
    "#                         print(\"Error!\")\n",
    "#     counter += 1\n",
    "#     if counter > 15:\n",
    "#         break"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "3_GgDL7ncxT3"
   },
   "source": [
    "# Network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "bqobZQZQcxUE"
   },
   "outputs": [],
   "source": [
    "def create_downsample(channels, inputs):\n",
    "    x = keras.layers.BatchNormalization(momentum=0.9)(inputs)\n",
    "    x = keras.layers.LeakyReLU(0.1)(x)\n",
    "    x_1 = keras.layers.MaxPool2D(2)(x)\n",
    "    x_2 = keras.layers.AveragePooling2D(2)(x)\n",
    "    x = keras.layers.concatenate([x_1, x_2])\n",
    "    x = keras.layers.Conv2D(channels, 1, padding=\"same\", use_bias=False)(x)\n",
    "    return x\n",
    "\n",
    "def create_resblock(channels, inputs):\n",
    "    x = keras.layers.BatchNormalization(momentum=0.9)(inputs)\n",
    "    x = keras.layers.LeakyReLU(0.1)(x)\n",
    "    x_1 = keras.layers.Conv2D(channels, 3, padding='same', use_bias=False)(x)\n",
    "    x = keras.layers.BatchNormalization(momentum=0.9)(x_1)\n",
    "    x = keras.layers.LeakyReLU(0.1)(x)\n",
    "    x = keras.layers.Conv2D(channels, 3, padding='same', use_bias=False)(x)\n",
    "#     x = keras.layers.add([x, inputs])\n",
    "    x = keras.layers.BatchNormalization(momentum=0.9)(x)\n",
    "    x = keras.layers.LeakyReLU(0.1)(x)\n",
    "    x = keras.layers.Conv2D(channels, 3, padding='same', use_bias=False)(x)\n",
    "    x = keras.layers.BatchNormalization(momentum=0.9)(x)\n",
    "    x = keras.layers.LeakyReLU(0.1)(x)\n",
    "    x = keras.layers.Conv2D(channels, 3, padding='same', use_bias=False)(x)\n",
    "    x = keras.layers.add([x, x_1])\n",
    "    return x\n",
    "\n",
    "def create_network(input_size, channels, n_blocks=2, depth=4):\n",
    "    # input - 512x512x3\n",
    "    inputs = keras.Input(shape=(input_size, input_size, 1))\n",
    "    \n",
    "    # downsample to 256x256x24\n",
    "    x = keras.layers.Conv2D(channels, 3, strides=(2,2), padding='same', use_bias=False)(inputs)\n",
    "    x = keras.layers.BatchNormalization(momentum=0.9)(x)\n",
    "    x = keras.layers.LeakyReLU(0.1)(x)\n",
    "    \n",
    "#     x = keras.layers.Conv2D(channels, 3, padding='same', use_bias=False)(inputs)\n",
    "#     x = keras.layers.BatchNormalization(momentum=0.9)(x)\n",
    "#     x = keras.layers.LeakyReLU(0.1)(x)\n",
    "    \n",
    "    x = keras.layers.Conv2D(channels, 3, padding='same', use_bias=False)(x)\n",
    "    x = keras.layers.BatchNormalization(momentum=0.9)(x)\n",
    "    x = keras.layers.LeakyReLU(0.1)(x)\n",
    "    \n",
    "    x = keras.layers.Conv2D(channels, 3, padding='same', use_bias=False)(x)\n",
    "    \n",
    "    # residual blocks\n",
    "    for d in range(depth):\n",
    "        \n",
    "        x = create_downsample(channels, x)\n",
    "        channels = channels * 2\n",
    "        for b in range(n_blocks):\n",
    "            x = create_resblock(channels, x)\n",
    "        \n",
    "    x_2 = x\n",
    "    x = keras.layers.BatchNormalization(momentum=0.9)(x)\n",
    "    x = keras.layers.LeakyReLU(0)(x)\n",
    "    \n",
    "    # dilated convolutions for context - 15x15x512\n",
    "    x = keras.layers.Conv2D(channels, (3,3), padding='same', dilation_rate=(2,2), activation=None, name=\"dilated_conv_1\")(x)\n",
    "    x = keras.layers.BatchNormalization(momentum=0.9)(x)\n",
    "    x = keras.layers.LeakyReLU(0.1)(x)\n",
    "    \n",
    "    x = keras.layers.Conv2D(channels, (3,3), padding='same', strides=(1,1), activation=None, name=\"last_conv\")(x)\n",
    "    x = keras.layers.BatchNormalization(momentum=0.9)(x)\n",
    "    x = keras.layers.LeakyReLU(0.1)(x)\n",
    "    \n",
    "    # downsample to 8x8 with stride 2\n",
    "    x = keras.layers.Conv2D(512, (5,5), padding='same', strides=(3,3), activation=None, name=\"downsample_1\", kernel_regularizer=keras.regularizers.l2(l=0.002))(x)\n",
    "    x = keras.layers.BatchNormalization(momentum=0.9)(x)\n",
    "    x = keras.layers.LeakyReLU(0.1)(x)\n",
    "    x = keras.layers.Dropout(0.10)(x)\n",
    "    \n",
    "    # bounding box branch\n",
    "    b = keras.layers.Conv2D(768, (1,1), padding='same', activation=None, name=\"fc_1_b\", kernel_regularizer=keras.regularizers.l2(l=0.01))(x)\n",
    "    b = keras.layers.BatchNormalization(momentum=0.9)(b)\n",
    "    b = keras.layers.LeakyReLU(0.01)(b)\n",
    "    b = keras.layers.Dropout(0.20)(b)\n",
    "    \n",
    "    b = keras.layers.Conv2D(1024, (1,1), padding='same', activation=None, name=\"fc_2_b\", kernel_regularizer=keras.regularizers.l2(l=0.005))(b)\n",
    "    b = keras.layers.BatchNormalization(momentum=0.9)(b)\n",
    "    b = keras.layers.LeakyReLU(0.01)(b)\n",
    "    \n",
    "    boxes = keras.layers.Conv2D(5, (1,1), padding='same', activation=\"linear\")(b)\n",
    "#     boxes = keras.layers.concatenate([confidence, boxes], name=\"bboxes_output\")\n",
    "    \n",
    "    # return both outputs\n",
    "    model = keras.Model(inputs=inputs, outputs=boxes)\n",
    "    \n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "AvdSzGI4cxUL"
   },
   "source": [
    "# Train network\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "# after each epoch run our function to calculate the actual IOU for half of the validation images. \n",
    "# I was not able to figure out how to implement this in tensorflow since I don't think we can iterate\n",
    "# through tensors. So we do it as a callback instead. Hopefully this will approximate the actual IOU score.\n",
    "class Calc_IOU_CB(keras.callbacks.Callback):\n",
    "    \n",
    "    def on_epoch_end(self, epoch, logs={}):\n",
    "        valid_gen2 = generator(folder, valid_filenames, pneumonia_locations, batch_size=32, image_size=IMAGE_SIZE, shuffle=True, predict=False)\n",
    "        ious = []\n",
    "        counter = 0\n",
    "        for imgs, labels in valid_gen2:\n",
    "            preds = self.model.predict(imgs)\n",
    "            iou = overlap_iou2(labels, preds)\n",
    "            ious.append(iou)\n",
    "            counter += BATCH_SIZE\n",
    "            \n",
    "            if counter > 900:\n",
    "                break\n",
    "            \n",
    "        print(\"Epoch\", epoch, \": Mean IOU:\", np.mean(ious))\n",
    "        \n",
    "        return"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 3451
    },
    "colab_type": "code",
    "id": "DCgIz8CwcxUO",
    "outputId": "ce38cdfb-1806-488d-838f-d180e198a099"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "__________________________________________________________________________________________________\n",
      "Layer (type)                    Output Shape         Param #     Connected to                     \n",
      "==================================================================================================\n",
      "input_2 (InputLayer)            (None, 480, 480, 1)  0                                            \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_25 (Conv2D)              (None, 240, 240, 24) 216         input_2[0][0]                    \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_29 (BatchNo (None, 240, 240, 24) 96          conv2d_25[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_29 (LeakyReLU)      (None, 240, 240, 24) 0           batch_normalization_29[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_26 (Conv2D)              (None, 240, 240, 24) 5184        leaky_re_lu_29[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_30 (BatchNo (None, 240, 240, 24) 96          conv2d_26[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_30 (LeakyReLU)      (None, 240, 240, 24) 0           batch_normalization_30[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_27 (Conv2D)              (None, 240, 240, 24) 5184        leaky_re_lu_30[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_31 (BatchNo (None, 240, 240, 24) 96          conv2d_27[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_31 (LeakyReLU)      (None, 240, 240, 24) 0           batch_normalization_31[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "max_pooling2d_5 (MaxPooling2D)  (None, 120, 120, 24) 0           leaky_re_lu_31[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "average_pooling2d_5 (AveragePoo (None, 120, 120, 24) 0           leaky_re_lu_31[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "concatenate_5 (Concatenate)     (None, 120, 120, 48) 0           max_pooling2d_5[0][0]            \n",
      "                                                                 average_pooling2d_5[0][0]        \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_28 (Conv2D)              (None, 120, 120, 24) 1152        concatenate_5[0][0]              \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_32 (BatchNo (None, 120, 120, 24) 96          conv2d_28[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_32 (LeakyReLU)      (None, 120, 120, 24) 0           batch_normalization_32[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_29 (Conv2D)              (None, 120, 120, 48) 10368       leaky_re_lu_32[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_33 (BatchNo (None, 120, 120, 48) 192         conv2d_29[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_33 (LeakyReLU)      (None, 120, 120, 48) 0           batch_normalization_33[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_30 (Conv2D)              (None, 120, 120, 48) 20736       leaky_re_lu_33[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_34 (BatchNo (None, 120, 120, 48) 192         conv2d_30[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_34 (LeakyReLU)      (None, 120, 120, 48) 0           batch_normalization_34[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_31 (Conv2D)              (None, 120, 120, 48) 20736       leaky_re_lu_34[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_35 (BatchNo (None, 120, 120, 48) 192         conv2d_31[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_35 (LeakyReLU)      (None, 120, 120, 48) 0           batch_normalization_35[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_32 (Conv2D)              (None, 120, 120, 48) 20736       leaky_re_lu_35[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "add_5 (Add)                     (None, 120, 120, 48) 0           conv2d_32[0][0]                  \n",
      "                                                                 conv2d_29[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_36 (BatchNo (None, 120, 120, 48) 192         add_5[0][0]                      \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_36 (LeakyReLU)      (None, 120, 120, 48) 0           batch_normalization_36[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "max_pooling2d_6 (MaxPooling2D)  (None, 60, 60, 48)   0           leaky_re_lu_36[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "average_pooling2d_6 (AveragePoo (None, 60, 60, 48)   0           leaky_re_lu_36[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "concatenate_6 (Concatenate)     (None, 60, 60, 96)   0           max_pooling2d_6[0][0]            \n",
      "                                                                 average_pooling2d_6[0][0]        \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_33 (Conv2D)              (None, 60, 60, 48)   4608        concatenate_6[0][0]              \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_37 (BatchNo (None, 60, 60, 48)   192         conv2d_33[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_37 (LeakyReLU)      (None, 60, 60, 48)   0           batch_normalization_37[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_34 (Conv2D)              (None, 60, 60, 96)   41472       leaky_re_lu_37[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_38 (BatchNo (None, 60, 60, 96)   384         conv2d_34[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_38 (LeakyReLU)      (None, 60, 60, 96)   0           batch_normalization_38[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_35 (Conv2D)              (None, 60, 60, 96)   82944       leaky_re_lu_38[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_39 (BatchNo (None, 60, 60, 96)   384         conv2d_35[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_39 (LeakyReLU)      (None, 60, 60, 96)   0           batch_normalization_39[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_36 (Conv2D)              (None, 60, 60, 96)   82944       leaky_re_lu_39[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_40 (BatchNo (None, 60, 60, 96)   384         conv2d_36[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_40 (LeakyReLU)      (None, 60, 60, 96)   0           batch_normalization_40[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_37 (Conv2D)              (None, 60, 60, 96)   82944       leaky_re_lu_40[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "add_6 (Add)                     (None, 60, 60, 96)   0           conv2d_37[0][0]                  \n",
      "                                                                 conv2d_34[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_41 (BatchNo (None, 60, 60, 96)   384         add_6[0][0]                      \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_41 (LeakyReLU)      (None, 60, 60, 96)   0           batch_normalization_41[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "max_pooling2d_7 (MaxPooling2D)  (None, 30, 30, 96)   0           leaky_re_lu_41[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "average_pooling2d_7 (AveragePoo (None, 30, 30, 96)   0           leaky_re_lu_41[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "concatenate_7 (Concatenate)     (None, 30, 30, 192)  0           max_pooling2d_7[0][0]            \n",
      "                                                                 average_pooling2d_7[0][0]        \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_38 (Conv2D)              (None, 30, 30, 96)   18432       concatenate_7[0][0]              \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_42 (BatchNo (None, 30, 30, 96)   384         conv2d_38[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_42 (LeakyReLU)      (None, 30, 30, 96)   0           batch_normalization_42[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_39 (Conv2D)              (None, 30, 30, 192)  165888      leaky_re_lu_42[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_43 (BatchNo (None, 30, 30, 192)  768         conv2d_39[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_43 (LeakyReLU)      (None, 30, 30, 192)  0           batch_normalization_43[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_40 (Conv2D)              (None, 30, 30, 192)  331776      leaky_re_lu_43[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_44 (BatchNo (None, 30, 30, 192)  768         conv2d_40[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_44 (LeakyReLU)      (None, 30, 30, 192)  0           batch_normalization_44[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_41 (Conv2D)              (None, 30, 30, 192)  331776      leaky_re_lu_44[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_45 (BatchNo (None, 30, 30, 192)  768         conv2d_41[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_45 (LeakyReLU)      (None, 30, 30, 192)  0           batch_normalization_45[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_42 (Conv2D)              (None, 30, 30, 192)  331776      leaky_re_lu_45[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "add_7 (Add)                     (None, 30, 30, 192)  0           conv2d_42[0][0]                  \n",
      "                                                                 conv2d_39[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_46 (BatchNo (None, 30, 30, 192)  768         add_7[0][0]                      \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_46 (LeakyReLU)      (None, 30, 30, 192)  0           batch_normalization_46[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "max_pooling2d_8 (MaxPooling2D)  (None, 15, 15, 192)  0           leaky_re_lu_46[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "average_pooling2d_8 (AveragePoo (None, 15, 15, 192)  0           leaky_re_lu_46[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "concatenate_8 (Concatenate)     (None, 15, 15, 384)  0           max_pooling2d_8[0][0]            \n",
      "                                                                 average_pooling2d_8[0][0]        \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_43 (Conv2D)              (None, 15, 15, 192)  73728       concatenate_8[0][0]              \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_47 (BatchNo (None, 15, 15, 192)  768         conv2d_43[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_47 (LeakyReLU)      (None, 15, 15, 192)  0           batch_normalization_47[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_44 (Conv2D)              (None, 15, 15, 384)  663552      leaky_re_lu_47[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_48 (BatchNo (None, 15, 15, 384)  1536        conv2d_44[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_48 (LeakyReLU)      (None, 15, 15, 384)  0           batch_normalization_48[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_45 (Conv2D)              (None, 15, 15, 384)  1327104     leaky_re_lu_48[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_49 (BatchNo (None, 15, 15, 384)  1536        conv2d_45[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_49 (LeakyReLU)      (None, 15, 15, 384)  0           batch_normalization_49[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_46 (Conv2D)              (None, 15, 15, 384)  1327104     leaky_re_lu_49[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_50 (BatchNo (None, 15, 15, 384)  1536        conv2d_46[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_50 (LeakyReLU)      (None, 15, 15, 384)  0           batch_normalization_50[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_47 (Conv2D)              (None, 15, 15, 384)  1327104     leaky_re_lu_50[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "add_8 (Add)                     (None, 15, 15, 384)  0           conv2d_47[0][0]                  \n",
      "                                                                 conv2d_44[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_51 (BatchNo (None, 15, 15, 384)  1536        add_8[0][0]                      \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_51 (LeakyReLU)      (None, 15, 15, 384)  0           batch_normalization_51[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "dilated_conv_1 (Conv2D)         (None, 15, 15, 384)  1327488     leaky_re_lu_51[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_52 (BatchNo (None, 15, 15, 384)  1536        dilated_conv_1[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_52 (LeakyReLU)      (None, 15, 15, 384)  0           batch_normalization_52[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "last_conv (Conv2D)              (None, 15, 15, 384)  1327488     leaky_re_lu_52[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_53 (BatchNo (None, 15, 15, 384)  1536        last_conv[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_53 (LeakyReLU)      (None, 15, 15, 384)  0           batch_normalization_53[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "downsample_1 (Conv2D)           (None, 5, 5, 512)    4915712     leaky_re_lu_53[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_54 (BatchNo (None, 5, 5, 512)    2048        downsample_1[0][0]               \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_54 (LeakyReLU)      (None, 5, 5, 512)    0           batch_normalization_54[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "dropout_3 (Dropout)             (None, 5, 5, 512)    0           leaky_re_lu_54[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "fc_1_b (Conv2D)                 (None, 5, 5, 768)    393984      dropout_3[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_55 (BatchNo (None, 5, 5, 768)    3072        fc_1_b[0][0]                     \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_55 (LeakyReLU)      (None, 5, 5, 768)    0           batch_normalization_55[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "dropout_4 (Dropout)             (None, 5, 5, 768)    0           leaky_re_lu_55[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "fc_2_b (Conv2D)                 (None, 5, 5, 1024)   787456      dropout_4[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_56 (BatchNo (None, 5, 5, 1024)   4096        fc_2_b[0][0]                     \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_56 (LeakyReLU)      (None, 5, 5, 1024)   0           batch_normalization_56[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_48 (Conv2D)              (None, 5, 5, 5)      5125        leaky_re_lu_56[0][0]             \n",
      "==================================================================================================\n",
      "Total params: 15,060,253\n",
      "Trainable params: 15,047,485\n",
      "Non-trainable params: 12,768\n",
      "__________________________________________________________________________________________________\n",
      "None\n"
     ]
    }
   ],
   "source": [
    "def binary_accuracy(y_true, y_pred):\n",
    "    # apply sigmoid to our predictions\n",
    "    y_pred = tf.sigmoid(y_pred)\n",
    "    # round both since our negative truths are 1e-16 instead of 0\n",
    "    y_true = tf.round(tf.reshape(y_true, [-1, 5]))\n",
    "    y_pred = tf.round(tf.reshape(y_pred, [-1, 5]))\n",
    "    \n",
    "    acc = tf.reduce_mean(tf.cast(tf.equal(y_true[:,0], y_pred[:,0]), dtype=tf.float32))\n",
    "    return acc\n",
    "\n",
    "def overlap_iou(y_true, y_pred):\n",
    "    \"\"\"\n",
    "    Args:\n",
    "        bboxes1: shape (batch_size, 16, 16, 4)\n",
    "            with x1, y1, x2, y2 point order.\n",
    "        bboxes2: shape (batch_size, 16, 16, 4)\n",
    "            with x1, y1, x2, y2 point order.\n",
    "        p1 *-----\n",
    "           |     |\n",
    "           |_____* p2\n",
    "    Returns:\n",
    "        Tensor with shape (total_bboxes1, total_bboxes2)\n",
    "        with the IoU (intersection over union) of bboxes1[i] and bboxes2[j]\n",
    "        in [i, j].\n",
    "    \"\"\"\n",
    "    y_pred = tf.sigmoid(y_pred)\n",
    "    \n",
    "    # flatten the data because it's easier that way\n",
    "    bboxes1 = tf.reshape(y_true, (-1, 5))\n",
    "    bboxes2 = tf.reshape(y_pred, (-1, 5))\n",
    "    \n",
    "    # split the components out\n",
    "    true_boxes, x11, y11, w1, h1 = tf.split(bboxes1, 5, axis=1)\n",
    "    pred_conf, x21, y21, w2, h2 = tf.split(bboxes2, 5, axis=1)\n",
    "    \n",
    "    # uncenter the data - make sure the numbers are positive\n",
    "    x11 = x11 * CELL_SIZE\n",
    "    x21 = x21 * CELL_SIZE\n",
    "    y11 = y11 * CELL_SIZE\n",
    "    y21 = y21 * CELL_SIZE\n",
    "    \n",
    "    w1 = w1 * 1024\n",
    "    w2 = w2 * 1024\n",
    "    h1 = h1 * 1024\n",
    "    h2 = h2 * 1024\n",
    "    \n",
    "    # is there either a box predicted here or a box actually here?\n",
    "    mask = (pred_conf >= 0.5) | (true_boxes > 0.5)\n",
    "    \n",
    "    # get the far corners of the boxes\n",
    "    x12 = x11 + (w1 / 2)\n",
    "    y12 = y11 + (h1 / 2)\n",
    "    x22 = x21 + (w2 / 2)\n",
    "    y22 = y21 + (h2 / 2)\n",
    "    \n",
    "    x11 = x11 - (w1 / 2)\n",
    "    y11 = y11 - (h1 / 2)\n",
    "    x21 = x21 - (w2 / 2)\n",
    "    y21 = y21 - (h2 / 2)\n",
    "\n",
    "    # find the corners of the intersection area\n",
    "    xI1 = tf.maximum(x11, x21)\n",
    "    yI1 = tf.maximum(y11, y21)\n",
    "\n",
    "    xI2 = tf.minimum(x12, x22)\n",
    "    yI2 = tf.minimum(y12, y22)\n",
    "    \n",
    "    # get the intersection area, if the truth has no boxes it is 0\n",
    "    inter_area = true_boxes * (xI2 - xI1 + 1) * (yI2 - yI1 + 1)\n",
    "\n",
    "    # get the area of each box\n",
    "    bboxes1_area = (w1 + 1) * (h1 + 1)\n",
    "    bboxes2_area = (w2 + 1) * (h2 + 1)\n",
    "    \n",
    "    # union is area of both boxes - intersection\n",
    "    union = (bboxes1_area + bboxes2_area) - inter_area + 1\n",
    "    \n",
    "    iou = tf.maximum(inter_area/(union + 1e-6), 0)\n",
    "    \n",
    "    # apply the mask\n",
    "    iou = tf.boolean_mask(iou, mask)\n",
    "    \n",
    "    # reduce the mean so we have mean iou for our inputs\n",
    "    return tf.reduce_mean(iou)\n",
    "\n",
    "def overlap_iou2(y_true, y_pred):\n",
    "    \"\"\"\n",
    "    Args:\n",
    "        y_true and y_pred - arrays of boxes containing center points, h and w of boxes\n",
    "        p1 *-----\n",
    "           |     |\n",
    "           |_____* p2\n",
    "    Returns:\n",
    "        Average IOU over boxes\n",
    "    \"\"\"\n",
    "    OVERLAP = 0.3\n",
    "    \n",
    "    # apply the sigmoid\n",
    "    y_pred = sigmoid(y_pred)\n",
    "    \n",
    "    ious = []\n",
    "    \n",
    "    # loop over both sets of boxes\n",
    "    for truth, pred in zip(y_true, y_pred):\n",
    "        # get true area\n",
    "        true_area = 0\n",
    "        true_boxes = []\n",
    "        for i in range(GRID_SIZE):\n",
    "            for j in range(GRID_SIZE):\n",
    "                c, x, y, w, h = truth[i,j,:]\n",
    "                # if we have an ROI\n",
    "                if c > 0.5:\n",
    "                    # unnormalize the data\n",
    "                    w, h = w*1024, h*1024\n",
    "                    true_area += w * h\n",
    "                    \n",
    "                    # unnormalize the x and y\n",
    "                    x = unnorm(x, j, CELL_SIZE)\n",
    "                    y = unnorm(y, i, CELL_SIZE)\n",
    "                    \n",
    "                    true_boxes.append([x, y, x+w, y+h])\n",
    "                    \n",
    "        pred_area = 0\n",
    "        pred_boxes = []\n",
    "        # get the predicted area\n",
    "        for i in range(GRID_SIZE):\n",
    "            for j in range(GRID_SIZE):\n",
    "                c, x, y, w, h = pred[i,j,:]\n",
    "                # if we have an ROI\n",
    "                if c > 0.5:\n",
    "                    # unnormalize the data\n",
    "                    w, h = w*1024, h*1024\n",
    "                    pred_area += w * h\n",
    "                    \n",
    "                    # unnormalize the x and y\n",
    "                    x = unnorm(x, j, CELL_SIZE)\n",
    "                    y = unnorm(y, i, CELL_SIZE)\n",
    "                    \n",
    "                    pred_boxes.append([x, y, x+w, y+h])\n",
    "            \n",
    "        intersect_area = 0\n",
    "        \n",
    "        # non-max suppression?\n",
    "        pred_boxes = non_max_suppression_fast(np.array(pred_boxes), OVERLAP)\n",
    "        \n",
    "        # get the intersection\n",
    "        for pred_box in pred_boxes:\n",
    "            x1_p, y1_p, x2_p, y2_p = pred_box\n",
    "            \n",
    "            for true_box in true_boxes:\n",
    "                x1_t, y1_t, x2_t, y2_t = true_box\n",
    "                \n",
    "                # if the boxes overlap at all\n",
    "                if (x1_p >= x1_t and y1_p >= y1_t) or (x1_p <= x1_t and y1_p <= y1_t):\n",
    "                    # get the intersection corners\n",
    "                    x1_i = np.maximum(x1_p, x1_t)\n",
    "                    y1_i = np.maximum(y1_p, y1_t)\n",
    "                    x2_i = np.minimum(x2_p, x2_t)\n",
    "                    y2_i = np.minimum(y2_p, y2_t)\n",
    "                    \n",
    "                    # get area of intersect\n",
    "                    i_w, i_h = x2_i - x1_i, y2_i - y1_i\n",
    "\n",
    "                    # trap for negative numbers\n",
    "                    i_w = np.maximum(i_w, 0)\n",
    "                    i_h = np.maximum(i_h, 0)\n",
    "\n",
    "                    intersection = i_w * i_h\n",
    "                    intersect_area += intersection\n",
    "\n",
    "        union = true_area + pred_area - intersect_area\n",
    "        \n",
    "        iou = intersect_area / (union + 1e-16)\n",
    "        \n",
    "        # only count the IOU if there are truths or predictions\n",
    "        if len(true_boxes) or len(pred_boxes):\n",
    "            ious.append(iou)\n",
    "    \n",
    "    iou_area = np.mean(ious)\n",
    "    return iou_area\n",
    "                    \n",
    "def loss_fn(y_true, y_pred):\n",
    "    # get the iou loss\n",
    "    iou_loss = iou_loss_fn(y_true, y_pred)\n",
    "    \n",
    "    # get the xe loss\n",
    "    xe_loss = binary_cross_entropy(y_true, y_pred)\n",
    "    \n",
    "    # get the box loss\n",
    "    box_loss = adj_mse(y_true, y_pred)\n",
    "    \n",
    "    # add the losses and return them\n",
    "    return (1.5 * xe_loss) + (box_loss * 3.0) + (iou_loss * 3.0)\n",
    "\n",
    "# only apply mse to layers with high confidence that there is an ROI or if there actually is an ROI\n",
    "def adj_mse(y_true, y_pred):\n",
    "    # coefficients\n",
    "    lam_coord = 5\n",
    "    lam_noobj = 0.5\n",
    "    \n",
    "    y_pred = tf.sigmoid(y_pred)\n",
    "        \n",
    "    # flatten the inputs\n",
    "    y_true = tf.reshape(y_true, (-1, 5))\n",
    "    y_pred = tf.reshape(y_pred, (-1, 5))\n",
    "\n",
    "    # separate the confidence from the boxes\n",
    "    conf_true, x_true, y_true, w_true, h_true = tf.split(y_true, 5, axis=1)\n",
    "    conf_pred, x_pred, y_pred, w_pred, h_pred = tf.split(y_pred, 5, axis=1)\n",
    "\n",
    "    # center squared error\n",
    "    center_se = tf.square(x_true - y_pred) + tf.square(y_true - y_pred)\n",
    "    size_se =  tf.square(tf.sqrt(w_true) - tf.sqrt(w_pred)) + tf.square(tf.sqrt(h_true) - tf.sqrt(h_pred))\n",
    "    \n",
    "    box_se = center_se + size_se\n",
    "    \n",
    "    # only get loss for boxes which are actually positive\n",
    "    mask = tf.greater(conf_true, 0.5)\n",
    "    box_se = tf.boolean_mask(box_se, mask)\n",
    "    \n",
    "    # weight the loss higher\n",
    "    box_se = tf.multiply(box_se, 5.0)\n",
    "    \n",
    "    loss = tf.reduce_sum(box_se)\n",
    "    \n",
    "    return loss\n",
    "    \n",
    "# use weight of 0.5 for negative cells, 19 for positive ones\n",
    "def binary_cross_entropy(y_true, y_pred):\n",
    "    conf_true = y_true[...,0]\n",
    "    conf_pred = tf.sigmoid(y_pred[...,0])\n",
    "    \n",
    "    weight = 12.0\n",
    "    weights = tf.multiply(y_true[...,0], weight) + 0.5\n",
    "    \n",
    "    xe = tf.multiply(tf.square(conf_true - conf_pred), weights)\n",
    "    \n",
    "    return tf.reduce_sum(xe)\n",
    "\n",
    "def objectness_loss_fn(y_true, y_pred):\n",
    "    y_pred = tf.sigmoid(y_pred)\n",
    "    \n",
    "    # pred box conf\n",
    "    pred_box_conf = y_pred[...,0]\n",
    "    \n",
    "    # separate the x, y from the w, h\n",
    "    true_box_xy = y_true[...,1:3] * CELL_SIZE\n",
    "    true_box_wh = y_true[...,3:] * 1024\n",
    "    \n",
    "    pred_box_xy = y_pred[...,1:3] * CELL_SIZE\n",
    "    pred_box_wh = y_pred[...,3:] * 1024\n",
    "    \n",
    "    # get the corners of the boxes by subtracting or adding half of h, w\n",
    "    true_wh_half = true_box_wh / 2.\n",
    "    true_mins    = true_box_xy - true_wh_half\n",
    "    true_maxes   = true_box_xy + true_wh_half\n",
    "    \n",
    "    pred_wh_half = pred_box_wh / 2.\n",
    "    pred_mins    = pred_box_xy - pred_wh_half\n",
    "    pred_maxes   = pred_box_xy + pred_wh_half  \n",
    "    \n",
    "    # get the corners of the intersect area\n",
    "    intersect_mins  = tf.maximum(pred_mins,  true_mins)\n",
    "    intersect_maxes = tf.minimum(pred_maxes, true_maxes)\n",
    "    intersect_wh    = tf.maximum(intersect_maxes - intersect_mins, 0.)\n",
    "    intersect_areas = intersect_wh[..., 0] * intersect_wh[..., 1]\n",
    "    \n",
    "    true_areas = true_box_wh[..., 0] * true_box_wh[..., 1]\n",
    "    pred_areas = pred_box_wh[..., 0] * pred_box_wh[..., 1]\n",
    "\n",
    "    union_areas = pred_areas + true_areas - intersect_areas\n",
    "    iou_scores  = tf.truediv(intersect_areas, union_areas + 1e-16)\n",
    "    \n",
    "    true_box_conf = iou_scores * y_true[..., 0]\n",
    "    \n",
    "    conf_mask  = tf.zeros_like(iou_scores)\n",
    "    conf_mask = conf_mask + tf.cast((iou_scores < 0.6), dtype=tf.float32) * (1 - y_true[...,0])\n",
    "    conf_mask = conf_mask + y_true[..., 0] * 5.0\n",
    "    \n",
    "    nb_conf_box  = tf.reduce_sum(tf.cast((conf_mask > 0.0), dtype=tf.float32)) \n",
    "    \n",
    "    loss_conf  =  tf.reduce_sum(tf.square(true_box_conf-pred_box_conf) * conf_mask)  / (nb_conf_box + 1e-6) / 2.\n",
    "    \n",
    "    return tf.reduce_sum(loss_conf)\n",
    "\n",
    "def iou_loss_fn(y_true, y_pred):\n",
    "    # apply sigmoid to predictions\n",
    "    y_pred = tf.sigmoid(y_pred)\n",
    "    \n",
    "    # separate the x, y from the w, h and unnormalize them\n",
    "    true_box_xy = y_true[...,1:3] * CELL_SIZE\n",
    "    true_box_wh = y_true[...,3:] * 1024\n",
    "    \n",
    "    pred_box_xy = y_pred[...,1:3] * CELL_SIZE\n",
    "    pred_box_wh = y_pred[...,3:] * 1024\n",
    "    \n",
    "    # get the corners of the boxes by subtracting or adding half of h, w\n",
    "    true_wh_half = true_box_wh / 2.\n",
    "    true_mins    = true_box_xy - true_wh_half\n",
    "    true_maxes   = true_box_xy + true_wh_half\n",
    "    \n",
    "    pred_wh_half = pred_box_wh / 2.\n",
    "    pred_mins    = pred_box_xy - pred_wh_half\n",
    "    pred_maxes   = pred_box_xy + pred_wh_half  \n",
    "    \n",
    "    # get the corners of the intersect area\n",
    "    intersect_mins  = tf.maximum(pred_mins,  true_mins)\n",
    "    intersect_maxes = tf.minimum(pred_maxes, true_maxes)\n",
    "    intersect_wh    = tf.maximum(intersect_maxes - intersect_mins, 0.)\n",
    "    intersect_areas = intersect_wh[..., 0] * intersect_wh[..., 1]\n",
    "    \n",
    "    # get the area of the boxes\n",
    "    true_areas = true_box_wh[..., 0] * true_box_wh[..., 1]\n",
    "    pred_areas = pred_box_wh[..., 0] * pred_box_wh[..., 1]\n",
    "\n",
    "    # calculate the IOU\n",
    "    union_areas = pred_areas + true_areas - intersect_areas\n",
    "    iou_scores  = tf.truediv(intersect_areas, union_areas + 1e-16)\n",
    "    \n",
    "    # only use the IOU from boxes which actually have ROIs\n",
    "    mask = tf.greater(y_true[...,0], 0.5)\n",
    "    use_iou = tf.boolean_mask(1 - iou_scores, mask)\n",
    "    \n",
    "    return tf.reduce_sum(use_iou)\n",
    "    \n",
    "# create network and compiler\n",
    "model = create_network(input_size=IMAGE_SIZE, channels=24, n_blocks=1, depth=4)\n",
    "\n",
    "optimizer = keras.optimizers.Adam(lr=0.001, beta_1=0.9, beta_2=0.999, decay=1e-7, clipnorm=2.0, clipvalue=0.75)\n",
    "\n",
    "model.compile(optimizer=optimizer,\n",
    "              loss=loss_fn,\n",
    "              metrics=[binary_accuracy, overlap_iou, adj_mse, binary_cross_entropy])\n",
    "\n",
    "# cosine learning rate annealing\n",
    "def exp_decay(x):\n",
    "    lr0 = 0.004\n",
    "    epochs_drop = 10\n",
    "    drop = 0.85\n",
    "    lrate = lr0 * math.pow(drop, math.floor((1+x)/epochs_drop))\n",
    "    return lrate\n",
    "\n",
    "learning_rate = tf.keras.callbacks.LearningRateScheduler(exp_decay)\n",
    "checkpoint = keras.callbacks.ModelCheckpoint(CHECKPOINT_PATH, monitor='val_loss', verbose=1, save_best_only=False, save_weights_only=True, mode='auto', period=1)\n",
    "calc_iou = Calc_IOU_CB()\n",
    "\n",
    "# create train and validation generators\n",
    "folder = './stage_1_train_images'\n",
    "train_gen = generator(folder, train_filenames, pneumonia_locations, batch_size=BATCH_SIZE, image_size=IMAGE_SIZE, shuffle=True, augment=True, predict=False)\n",
    "valid_gen = generator(folder, valid_filenames, pneumonia_locations, batch_size=BATCH_SIZE, image_size=IMAGE_SIZE, shuffle=False, predict=False)\n",
    "\n",
    "print(model.summary())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "jSSSpmNE-L1L"
   },
   "outputs": [],
   "source": [
    "model.load_weights(CHECKPOINT_PATH)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 527
    },
    "colab_type": "code",
    "id": "fM5VeiuDcxUY",
    "outputId": "c3d66d6d-3c06-4aa6-f58d-46cfa42cdd9f"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.5/dist-packages/tensorflow/python/ops/gradients_impl.py:100: UserWarning: Converting sparse IndexedSlices to a dense Tensor of unknown shape. This may consume a large amount of memory.\n",
      "  \"Converting sparse IndexedSlices to a dense Tensor of unknown shape. \"\n",
      "/home/eric/.local/lib/python3.5/site-packages/skimage/transform/_warps.py:110: UserWarning: Anti-aliasing will be enabled by default in skimage 0.15 to avoid aliasing artifacts when down-sampling images.\n",
      "  warn(\"Anti-aliasing will be enabled by default in skimage 0.15 to \"\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/10\n",
      "1443/1444 [============================>.] - ETA: 0s - loss: 191.9186 - binary_accuracy: 0.7871 - overlap_iou: 0.0595 - adj_mse: 17.3663 - binary_cross_entropy: 66.2203\n",
      "Epoch 00001: saving model to yolo16_2_480.h5\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/eric/.local/lib/python3.5/site-packages/ipykernel_launcher.py:71: RuntimeWarning: overflow encountered in exp\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0 : Mean IOU: 0.011441588870911526\n",
      "1444/1444 [==============================] - 1061s 735ms/step - loss: 191.9178 - binary_accuracy: 0.7871 - overlap_iou: 0.0595 - adj_mse: 17.3667 - binary_cross_entropy: 66.2218 - val_loss: 151.2438 - val_binary_accuracy: 0.8307 - val_overlap_iou: 0.0524 - val_adj_mse: 14.1366 - val_binary_cross_entropy: 48.3534\n",
      "Epoch 2/10\n",
      "1443/1444 [============================>.] - ETA: 0s - loss: 143.0745 - binary_accuracy: 0.8434 - overlap_iou: 0.0552 - adj_mse: 13.2516 - binary_cross_entropy: 47.7908\n",
      "Epoch 00002: saving model to yolo16_2_480.h5\n",
      "Epoch 1 : Mean IOU: 0.012228542550636962\n",
      "1444/1444 [==============================] - 1044s 723ms/step - loss: 143.1154 - binary_accuracy: 0.8434 - overlap_iou: 0.0552 - adj_mse: 13.2567 - binary_cross_entropy: 47.8068 - val_loss: 137.6820 - val_binary_accuracy: 0.8920 - val_overlap_iou: 0.0795 - val_adj_mse: 13.6959 - val_binary_cross_entropy: 45.9999\n",
      "Epoch 3/10\n",
      " 307/1444 [=====>........................] - ETA: 11:46 - loss: 137.5453 - binary_accuracy: 0.8447 - overlap_iou: 0.0578 - adj_mse: 13.6169 - binary_cross_entropy: 46.4905"
     ]
    }
   ],
   "source": [
    "history = model.fit_generator(train_gen, validation_data=valid_gen, callbacks=[learning_rate, checkpoint, calc_iou], epochs=10, shuffle=True, verbose=1, initial_epoch=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "7tyG7P79Y42U"
   },
   "outputs": [],
   "source": [
    "history = model.fit_generator(train_gen, validation_data=valid_gen, callbacks=[learning_rate, checkpoint, calc_iou], epochs=25, shuffle=True, verbose=1, initial_epoch=18)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "history = model.fit_generator(train_gen, validation_data=valid_gen, callbacks=[learning_rate, checkpoint, calc_iou], epochs=35, shuffle=True, verbose=1, initial_epoch=31)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "history = model.fit_generator(train_gen, validation_data=valid_gen, callbacks=[learning_rate, checkpoint, calc_iou], epochs=50, shuffle=True, verbose=1, initial_epoch=37)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(20,5))\n",
    "plt.subplot(141)\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",
    "plt.subplot(142)\n",
    "plt.plot(history.epoch, history.history[\"adj_mse\"], label=\"Train MSE loss\")\n",
    "plt.plot(history.epoch, history.history[\"val_adj_mse\"], label=\"Valid MSE loss\")\n",
    "plt.legend()\n",
    "plt.subplot(143)\n",
    "plt.plot(history.epoch, history.history[\"binary_accuracy\"], label=\"Train Confidence accuracy\")\n",
    "plt.plot(history.epoch, history.history[\"val_binary_accuracy\"], label=\"Valid Confidence accuracy\")\n",
    "plt.legend()\n",
    "plt.subplot(144)\n",
    "plt.plot(history.epoch, history.history[\"overlap_iou\"], label=\"Train iou\")\n",
    "plt.plot(history.epoch, history.history[\"val_overlap_iou\"], label=\"Valid iou\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 269
    },
    "colab_type": "code",
    "id": "y_bLpJPBcxUh",
    "outputId": "49c438f6-aabc-41b9-f98e-1bf8519da3cc"
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(20,5))\n",
    "plt.subplot(141)\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",
    "plt.subplot(142)\n",
    "plt.plot(history.epoch, history.history[\"adj_mse\"], label=\"Train MSE loss\")\n",
    "plt.plot(history.epoch, history.history[\"val_adj_mse\"], label=\"Valid MSE loss\")\n",
    "plt.legend()\n",
    "plt.subplot(143)\n",
    "plt.plot(history.epoch, history.history[\"binary_accuracy\"], label=\"Train Confidence accuracy\")\n",
    "plt.plot(history.epoch, history.history[\"val_binary_accuracy\"], label=\"Valid Confidence accuracy\")\n",
    "plt.legend()\n",
    "plt.subplot(144)\n",
    "plt.plot(history.epoch, history.history[\"overlap_iou\"], label=\"Train iou\")\n",
    "plt.plot(history.epoch, history.history[\"val_overlap_iou\"], label=\"Valid iou\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(20,5))\n",
    "plt.subplot(141)\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",
    "plt.subplot(142)\n",
    "plt.plot(history.epoch, history.history[\"adj_mse\"], label=\"Train MSE loss\")\n",
    "plt.plot(history.epoch, history.history[\"val_adj_mse\"], label=\"Valid MSE loss\")\n",
    "plt.legend()\n",
    "plt.subplot(143)\n",
    "plt.plot(history.epoch, history.history[\"binary_accuracy\"], label=\"Train Confidence accuracy\")\n",
    "plt.plot(history.epoch, history.history[\"val_binary_accuracy\"], label=\"Valid Confidence accuracy\")\n",
    "plt.legend()\n",
    "plt.subplot(144)\n",
    "plt.plot(history.epoch, history.history[\"overlap_iou\"], label=\"Train iou\")\n",
    "plt.plot(history.epoch, history.history[\"val_overlap_iou\"], label=\"Valid iou\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "7dpxfbugzcLT"
   },
   "source": [
    "## Sample Predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.metrics_names"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mses = []\n",
    "problem_names = []\n",
    "\n",
    "# find which images are causing the high mse loss\n",
    "for filename in valid_filenames[:20]:    \n",
    "    # load and resize the image\n",
    "    img = pydicom.dcmread(os.path.join(train_dicom_dir, filename)).pixel_array\n",
    "    img = resize(img, (IMAGE_SIZE, IMAGE_SIZE), mode='reflect')\n",
    "    \n",
    "    # get the label\n",
    "    filename = filename.split('.')[0]\n",
    "    label = pneumonia_locations[filename]\n",
    "    # predict the image\n",
    "    loss = model.evaluate(img.reshape(1, IMAGE_SIZE, IMAGE_SIZE, 1), label.reshape(1, GRID_SIZE, GRID_SIZE, 5), verbose=0)\n",
    "    mses.append(loss[3])\n",
    "    if loss[3] > 10:\n",
    "        print(filename, \"Score:\", loss)\n",
    "        problem_names.append(filename + \".dcm\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "THRESHOLD = 0.75\n",
    "OVERLAP = 0.4\n",
    "NMS = True\n",
    "# look at some sample predictions\n",
    "samples = np.random.choice(valid_filenames, size=15, replace=False)\n",
    "# samples = problem_names\n",
    "\n",
    "coords = np.arange(0, 1024, CELL_SIZE)\n",
    "overall_ious = []\n",
    "\n",
    "for filename in samples:\n",
    "    # load the image\n",
    "    img = pydicom.dcmread(os.path.join(train_dicom_dir, filename)).pixel_array\n",
    "    \n",
    "    filename = filename.split('.')[0]\n",
    "    \n",
    "    f, ax = plt.subplots(figsize=(6, 6))\n",
    "    plt.imshow(img)\n",
    "    \n",
    "    # initialize our lists\n",
    "    ious = []\n",
    "    truths = []\n",
    "    boxes = []\n",
    "    \n",
    "    # draw the truth boxes\n",
    "    if filename in pneumonia_locations:\n",
    "        locs = pneumonia_locations[filename].copy()\n",
    "        for i in range(GRID_SIZE):\n",
    "            for j in range(GRID_SIZE):\n",
    "                pixel_data = locs[i,j,:]\n",
    "                if pixel_data[0] > 0.5:\n",
    "                    x, y, w, h = pixel_data[1:]\n",
    "                    \n",
    "                    # unnormalize the data\n",
    "                    w = w * 1024\n",
    "                    h = h * 1024\n",
    "                    \n",
    "                    x = unnorm(x, j, CELL_SIZE)\n",
    "                    y = unnorm(y, i, CELL_SIZE)\n",
    "                    \n",
    "                    # get the corners\n",
    "                    x = x - (w // 2)\n",
    "                    y = y - (h // 2)\n",
    "                    \n",
    "                    x = int(x)\n",
    "                    y = int(y)\n",
    "                    w = int(w)\n",
    "                    h = int(h)\n",
    "                    locs[i,j,:] = [1, x, y, w, h]\n",
    "                    print(\"Truth:\", i, j, x, y, w, h)\n",
    "                    truths.append([x, y, w, h])\n",
    "                    \n",
    "                    rect = patches.Rectangle((x,y),w,h,linewidth=1,edgecolor='b',facecolor='none')\n",
    "                    ax.add_patch(rect)\n",
    "                    \n",
    "    # predict the image\n",
    "    img = resize(img, (IMAGE_SIZE, IMAGE_SIZE), mode='reflect')\n",
    "    yhat = model.predict(img.reshape(1, IMAGE_SIZE, IMAGE_SIZE, 1))\n",
    "    yhat = sigmoid(yhat)\n",
    "#     break\n",
    "    conf = np.squeeze(yhat)[:,:,0]\n",
    "    bboxes = np.squeeze(yhat)[:,:,1:]    \n",
    "    pred_boxes = np.zeros_like(bboxes)\n",
    "#     print(conf)\n",
    "\n",
    "    # loop through our predictions\n",
    "    for i in range(GRID_SIZE):\n",
    "        for j in range(GRID_SIZE):\n",
    "            conf_ = conf[i,j]\n",
    "            # if we have a prediction\n",
    "            if conf_ > THRESHOLD:\n",
    "                x,y,w,h = bboxes[i,j,:]\n",
    "                \n",
    "                # unnormalize the data\n",
    "                w = w * 1024\n",
    "                h = h * 1024\n",
    "\n",
    "                x = unnorm(x, j, CELL_SIZE)\n",
    "                y = unnorm(y, i, CELL_SIZE)\n",
    "                \n",
    "                # convert to upper left corner from center\n",
    "                x = np.maximum(x - (w // 2), 0)\n",
    "                y = np.maximum(y - (h // 2), 0)\n",
    "                \n",
    "                x = int(x)\n",
    "                y = int(y)\n",
    "                w = int(w)\n",
    "                h = int(h)\n",
    "                \n",
    "                pred_boxes[i,j,:] = [x, y, w, h]\n",
    "                \n",
    "                print(\"Pred:\", i, j, \"conf:\", conf_, x, y, w, h)\n",
    "                \n",
    "                # if the boxes have width and height add them to our list\n",
    "                if w > 30 and h > 30:\n",
    "                    boxes.append([x,y,w,h])\n",
    "                \n",
    "    # do non-max suppression of our boxes\n",
    "    if NMS:\n",
    "        nms_boxes = non_max_suppression_fast(np.array(boxes), OVERLAP)\n",
    "    else:\n",
    "        nms_boxes = boxes\n",
    "    \n",
    "    # plot our boxes\n",
    "    for box in nms_boxes:\n",
    "        x,y,w,h = box\n",
    "        rect = patches.Rectangle((x,y), w, h, linewidth=1, edgecolor='r', facecolor='none')\n",
    "        ax.add_patch(rect)\n",
    "    \n",
    "    ## calculate the IOU\n",
    "    true_area = 0\n",
    "    # get the area of the true boxes\n",
    "    for true_box in truths:\n",
    "        x, y, w, h = true_box\n",
    "        area = w * h\n",
    "        true_area += area\n",
    "        \n",
    "    pred_area = 0\n",
    "    # get the area of the predictions\n",
    "    for pred_box in nms_boxes:\n",
    "        x, y, w, h = pred_box\n",
    "        area = w * h\n",
    "        pred_area += area\n",
    "        \n",
    "    overall_intersect_area = 0\n",
    "    # get the IOU by checking all combinations of boxes\n",
    "    for true_box in truths:\n",
    "        x1_t, y1_t, w_t, h_t = true_box\n",
    "        for pred_box in nms_boxes:\n",
    "            x1_p, y1_p, w_p, h_p = pred_box\n",
    "            \n",
    "            # get the far corners\n",
    "            x2_p, y2_p = x1_p + w_p, y1_p + h_p\n",
    "            x2_t, y2_t = x1_t + w_t, y1_t + h_t\n",
    "        \n",
    "            # get corners of intersection\n",
    "            x1_i = np.maximum(x1_p, x1_t)\n",
    "            y1_i = np.maximum(y1_p, y1_t)\n",
    "            x2_i = np.minimum(x2_p, x2_t)\n",
    "            y2_i = np.minimum(y2_p, y2_t)\n",
    "\n",
    "            # get area of intersect\n",
    "            i_w, i_h = x2_i - x1_i, y2_i - y1_i\n",
    "            \n",
    "            # trap for negative numbers\n",
    "            i_w = np.maximum(i_w, 0)\n",
    "            i_h = np.maximum(i_h, 0)\n",
    "            \n",
    "            intersect_area = i_w * i_h\n",
    "            overall_intersect_area += intersect_area\n",
    "            \n",
    "    for item in coords:\n",
    "        plt.axvline(item, linewidth=0.5)\n",
    "        plt.axhline(item, linewidth=0.5)\n",
    "    \n",
    "    print(\"True:\", true_area, \"Pred:\", pred_area, \"Intersect:\", overall_intersect_area)\n",
    "    union_area = true_area + pred_area - overall_intersect_area\n",
    "    \n",
    "    iou = overall_intersect_area / (union_area + 1e-6)\n",
    "    print(\"IOU:\", iou)\n",
    "    \n",
    "    overall_ious.append(iou)\n",
    "    plt.title(filename)\n",
    "    plt.show()\n",
    "    \n",
    "print(\"Overall Mean IOU:\", np.mean(overall_ious))    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "3tAntxescxUq"
   },
   "source": [
    "# Predict test images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "fnu7ebfScxUu",
    "outputId": "2e91d334-3204-4213-b74c-4c112b1c3512"
   },
   "outputs": [],
   "source": [
    "THRESHOLD = 0.75\n",
    "\n",
    "# load and shuffle filenames\n",
    "folder = './stage_1_test_images'\n",
    "test_filenames = os.listdir(folder)\n",
    "print('n test samples:', len(test_filenames))\n",
    "\n",
    "# create test generator with predict flag set to True\n",
    "test_gen = generator(folder, test_filenames, None, batch_size=24, image_size=IMAGE_SIZE, shuffle=False, predict=True)\n",
    "\n",
    "# create submission dictionary\n",
    "submission_dict = {}\n",
    "# loop through testset\n",
    "for imgs, filenames in test_gen:\n",
    "    \n",
    "    # predict batch of images\n",
    "    yhats = model.predict(imgs)\n",
    "    \n",
    "    # apply sigmoid\n",
    "    yhats = sigmoid(yhats)\n",
    "    \n",
    "    # loop through batch\n",
    "    for yhat, filename in zip(yhats, filenames):\n",
    "        predictionString = \"\"\n",
    "        boxes = []\n",
    "        for i in range(GRID_SIZE):\n",
    "            for j in range(GRID_SIZE):\n",
    "                conf = yhat[i, j, 0]\n",
    "                if conf > THRESHOLD:\n",
    "                    x, y, w, h = yhat[i,j, 1:]\n",
    "                    \n",
    "                    # possible thresholds to keep our boxes within reasonable sizes?\n",
    "                    if True: #w < 600 and h < 1000:\n",
    "                        w = w * 1024\n",
    "                        h = h * 1024\n",
    "\n",
    "                        x = unnorm(x, j, CELL_SIZE)\n",
    "                        y = unnorm(y, i, CELL_SIZE)\n",
    "\n",
    "                        # convert to upper left corner from center\n",
    "                        x = x - (w // 2)\n",
    "                        y = y - (h // 2)\n",
    "                        \n",
    "                        if w > 20 and h > 20:\n",
    "                            # make sure our boxes don't run off the edges of the images\n",
    "                            w = np.minimum(w, 1024 - x)\n",
    "                            h = np.minimum(h, 1024 - y)\n",
    "                            boxes.append([x,y,w,h])\n",
    "\n",
    "        # do our non-max suppression here\n",
    "        boxes = non_max_suppression_fast(np.array(boxes), 0.3)\n",
    "        \n",
    "        # loop through our suppressed boxes and creat the prediction string\n",
    "        for box in boxes:\n",
    "            x,y,w,h = box\n",
    "            \n",
    "            x = int(x)\n",
    "            y = int(y)\n",
    "            w = int(w)\n",
    "            h = int(h)\n",
    "        \n",
    "            # create the prediction string\n",
    "            predictionString += str(0.9) + ' ' + str(x) + ' ' + str(y) + ' ' + str(w) + ' ' + str(h) + ' '\n",
    "            \n",
    "        # add filename and predictionString to dictionary\n",
    "        filename = filename.split('.')[0]\n",
    "        submission_dict[filename] = predictionString\n",
    "\n",
    "    # stop if we've got them all\n",
    "    if len(submission_dict) >= len(test_filenames):\n",
    "        break\n",
    "    \n",
    "# save dictionary as csv file\n",
    "sub = pd.DataFrame.from_dict(submission_dict,orient='index')\n",
    "sub.index.names = ['patientId']\n",
    "sub.columns = ['PredictionString']\n",
    "\n",
    "now = datetime.datetime.now()\n",
    "today = str(now)[:10]\n",
    "submission_file = today + \"_yolo_submission.csv\" \n",
    "sub.to_csv(submission_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "U1NEg9QCcxU4",
    "outputId": "c7e7ccf3-3e94-4996-a07d-a1098b69f3ce"
   },
   "outputs": [],
   "source": [
    "!kaggle competitions submit -c rsna-pneumonia-detection-challenge -f {submission_file} -m \"YOLOv15 480x480 20 epochs\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 85
    },
    "colab_type": "code",
    "id": "yf06y04TVRi2",
    "outputId": "1ca39a9a-1aa6-4e2a-ec86-39a826c676dd"
   },
   "outputs": [],
   "source": [
    "# upload checkpoint to GCS\n",
    "project_id = 'mammography-198911'\n",
    "bucket_name = 'pneumonia'\n",
    "\n",
    "!gcloud config set project {project_id}\n",
    "!gsutil cp ./{CHECKPOINT_PATH} gs://{bucket_name}/"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "best_cnn_segmentation_connected_components_384x384_2.ipynb",
   "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.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}