yolo_512x512_10.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "igEd7obtcxR2"
   },
   "source": [
    "## Approach\n",
    "\n",
    "* This is a simplied version of YOLO. We take our 512x512 input, divide it into a 4x4 grid, 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 \"objectness\" loss of YOLO tries to make the confidence match the actual IOU. Since our model only outputs one box per cell, this effectively drives the confidence down towards the IOU rather than driving the IOU up.\n",
    "    * We replaced this component of the loss by calculating the IOU for cells with ground truth, subtracting that from 1 and using that as the loss. Since we want to maximize the IOU, why not just do it directly?\n",
    "\n",
    "## Network\n",
    "\n",
    "* The network consists of a number of residual blocks with convolutions and downsampling blocks with max pooling.\n",
    "* The network outputs a 4x4x5 array as described above.\n",
    "* We are using an Adam optimizer with gradient clipping to avoid exploding gradients.\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"
   ]
  },
  {
   "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",
    "    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": 50,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "hopsRmICcxSs"
   },
   "outputs": [],
   "source": [
    "with open('yolo_labels_centered_2x2_7.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": 51,
   "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(72)\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": 52,
   "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": 53,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "dFbLdFxVehhB"
   },
   "outputs": [],
   "source": [
    "BATCH_SIZE = 16\n",
    "IMAGE_SIZE = 512\n",
    "CHECKPOINT_PATH = \"yolo10_2_512.h5\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [],
   "source": [
    "# means to center data\n",
    "mu_x = 391.456158\n",
    "mu_y = 363.1358768\n",
    "mu_w = 220.8453815\n",
    "mu_h = 334.1743641\n",
    "mu_center_x = 501.8788487\n",
    "mu_center_y = 530.2230589"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "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 __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",
    "        confs = label[:,:,0]\n",
    "        boxes = label[:,:,1:]\n",
    "        \n",
    "        ## data augmentation may be complicated, let's do that later\n",
    "        # if augment then horizontal flip half the time\n",
    "        if self.augment and random.random() > 0.5:\n",
    "            # flip the image\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=0)         \n",
    "            \n",
    "            # flip the confidences lr as well\n",
    "            confs = np.flip(confs, axis=0)\n",
    "        \n",
    "        # very small shifts to image - if these are bigger they will mess with the label alignment\n",
    "        if self.augment:\n",
    "            h_offset = np.random.randint(low=0, high=5)\n",
    "            v_offset = np.random.randint(low=0, high=5)\n",
    "            # crop the images\n",
    "            img = img[v_offset:,h_offset:]\n",
    "        \n",
    "        # random contrast adjustment\n",
    "        if self.augment and 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",
    "        # 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((2,2,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": 56,
   "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": 57,
   "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 = keras.layers.MaxPool2D(2)(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.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",
    "    # 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)(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",
    "    x = keras.layers.BatchNormalization(momentum=0.9)(x)\n",
    "    x = keras.layers.LeakyReLU(0.1)(x)\n",
    "    \n",
    "    # 256x256x24 for residual connection\n",
    "    x = keras.layers.Conv2D(channels, 1, padding='same', use_bias=False)(x)\n",
    "    \n",
    "    # residual blocks\n",
    "    for d in range(depth):\n",
    "        channels = channels * 2\n",
    "        \n",
    "        x = create_downsample(channels, x)\n",
    "            \n",
    "        for b in range(n_blocks):\n",
    "            x = create_resblock(channels, x)\n",
    "    \n",
    "    x = keras.layers.BatchNormalization(momentum=0.9)(x)\n",
    "    x = keras.layers.LeakyReLU(0)(x)\n",
    "    \n",
    "    # dilated convolutions for context - 16x16x512\n",
    "    x_2 = keras.layers.Conv2D(512, (3,3), padding='same', dilation_rate=(2,2), activation=None, name=\"dilated_conv_1\")(x)\n",
    "    x = keras.layers.BatchNormalization(momentum=0.9)(x_2)\n",
    "    x = keras.layers.LeakyReLU(0.1)(x)\n",
    "    \n",
    "    x = keras.layers.Conv2D(512, (3,3), padding='same', dilation_rate=(2,2), activation=None, name=\"dilated_conv_2\")(x)\n",
    "    x = keras.layers.BatchNormalization(momentum=0.9)(x)\n",
    "    x = keras.layers.LeakyReLU(0.1)(x)\n",
    "    \n",
    "    x = keras.layers.Conv2D(512, (3,3), padding='same', activation=None, name=\"last_conv\")(x)\n",
    "    x = keras.layers.add([x, x_2])\n",
    "    x = keras.layers.BatchNormalization(momentum=0.9)(x)\n",
    "    x = keras.layers.LeakyReLU(0.1)(x)\n",
    "    \n",
    "    # downsample to 4x4x512 with stride 4\n",
    "    x = keras.layers.Conv2D(512, (4,4), padding='same', strides=(4,4), activation=None, name=\"downsample_1\")(x)\n",
    "    x = keras.layers.BatchNormalization(momentum=0.9)(x)\n",
    "    x = keras.layers.LeakyReLU(0.1)(x)\n",
    "    \n",
    "    # downsample to 2x2 with average pool\n",
    "    x = keras.layers.AveragePooling2D(2, 2, padding=\"same\", name=\"average_pool\")(x)\n",
    "    \n",
    "#     # confidence branch\n",
    "#     c = keras.layers.Conv2D(384, (1,1), padding='same', activation=None, name=\"fc_1_c\", kernel_regularizer=keras.regularizers.l2(l=0.01))(x)\n",
    "#     c = keras.layers.BatchNormalization(momentum=0.9)(c)\n",
    "#     c = keras.layers.LeakyReLU(0.1)(c)\n",
    "#     c = keras.layers.Dropout(0.25)(c)\n",
    "    \n",
    "#     c = keras.layers.Conv2D(256, (1,1), padding='same', activation=None, name=\"fc_2_c\", kernel_regularizer=keras.regularizers.l2(l=0.01))(c)\n",
    "#     c = keras.layers.BatchNormalization(momentum=0.9)(c)\n",
    "#     c = keras.layers.LeakyReLU(0.1)(c)\n",
    "#     c = keras.layers.Dropout(0.10)(c)\n",
    "    \n",
    "#     confidence = keras.layers.Conv2D(1, (1,1), padding='same', activation=None, name=\"confidence_output\")(c)\n",
    "    \n",
    "    # bounding box branch\n",
    "    b = keras.layers.Conv2D(1024, (1,1), padding='same', activation=None, name=\"fc_1_b\", kernel_regularizer=keras.regularizers.l2(l=0.005))(x)\n",
    "    b = keras.layers.BatchNormalization(momentum=0.9)(b)\n",
    "    b = keras.layers.LeakyReLU(0.01)(b)\n",
    "    b = keras.layers.Dropout(0.15)(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.001))(b)\n",
    "    b = keras.layers.BatchNormalization(momentum=0.9)(b)\n",
    "    b = keras.layers.LeakyReLU(0.01)(b)\n",
    "#     b = keras.layers.Dropout(0.05)(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": 58,
   "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_9 (InputLayer)            (None, 512, 512, 1)  0                                            \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_169 (Conv2D)             (None, 256, 256, 24) 216         input_9[0][0]                    \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_241 (BatchN (None, 256, 256, 24) 96          conv2d_169[0][0]                 \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_241 (LeakyReLU)     (None, 256, 256, 24) 0           batch_normalization_241[0][0]    \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_170 (Conv2D)             (None, 256, 256, 24) 5184        leaky_re_lu_241[0][0]            \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_242 (BatchN (None, 256, 256, 24) 96          conv2d_170[0][0]                 \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_242 (LeakyReLU)     (None, 256, 256, 24) 0           batch_normalization_242[0][0]    \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_171 (Conv2D)             (None, 256, 256, 24) 5184        leaky_re_lu_242[0][0]            \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_243 (BatchN (None, 256, 256, 24) 96          conv2d_171[0][0]                 \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_243 (LeakyReLU)     (None, 256, 256, 24) 0           batch_normalization_243[0][0]    \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_172 (Conv2D)             (None, 256, 256, 24) 576         leaky_re_lu_243[0][0]            \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_244 (BatchN (None, 256, 256, 24) 96          conv2d_172[0][0]                 \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_244 (LeakyReLU)     (None, 256, 256, 24) 0           batch_normalization_244[0][0]    \n",
      "__________________________________________________________________________________________________\n",
      "max_pooling2d_33 (MaxPooling2D) (None, 128, 128, 24) 0           leaky_re_lu_244[0][0]            \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_245 (BatchN (None, 128, 128, 24) 96          max_pooling2d_33[0][0]           \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_245 (LeakyReLU)     (None, 128, 128, 24) 0           batch_normalization_245[0][0]    \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_173 (Conv2D)             (None, 128, 128, 48) 10368       leaky_re_lu_245[0][0]            \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_246 (BatchN (None, 128, 128, 48) 192         conv2d_173[0][0]                 \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_246 (LeakyReLU)     (None, 128, 128, 48) 0           batch_normalization_246[0][0]    \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_174 (Conv2D)             (None, 128, 128, 48) 20736       leaky_re_lu_246[0][0]            \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_247 (BatchN (None, 128, 128, 48) 192         conv2d_174[0][0]                 \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_247 (LeakyReLU)     (None, 128, 128, 48) 0           batch_normalization_247[0][0]    \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_175 (Conv2D)             (None, 128, 128, 48) 20736       leaky_re_lu_247[0][0]            \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_248 (BatchN (None, 128, 128, 48) 192         conv2d_175[0][0]                 \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_248 (LeakyReLU)     (None, 128, 128, 48) 0           batch_normalization_248[0][0]    \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_176 (Conv2D)             (None, 128, 128, 48) 20736       leaky_re_lu_248[0][0]            \n",
      "__________________________________________________________________________________________________\n",
      "add_41 (Add)                    (None, 128, 128, 48) 0           conv2d_176[0][0]                 \n",
      "                                                                 conv2d_173[0][0]                 \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_249 (BatchN (None, 128, 128, 48) 192         add_41[0][0]                     \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_249 (LeakyReLU)     (None, 128, 128, 48) 0           batch_normalization_249[0][0]    \n",
      "__________________________________________________________________________________________________\n",
      "max_pooling2d_34 (MaxPooling2D) (None, 64, 64, 48)   0           leaky_re_lu_249[0][0]            \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_250 (BatchN (None, 64, 64, 48)   192         max_pooling2d_34[0][0]           \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_250 (LeakyReLU)     (None, 64, 64, 48)   0           batch_normalization_250[0][0]    \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_177 (Conv2D)             (None, 64, 64, 96)   41472       leaky_re_lu_250[0][0]            \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_251 (BatchN (None, 64, 64, 96)   384         conv2d_177[0][0]                 \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_251 (LeakyReLU)     (None, 64, 64, 96)   0           batch_normalization_251[0][0]    \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_178 (Conv2D)             (None, 64, 64, 96)   82944       leaky_re_lu_251[0][0]            \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_252 (BatchN (None, 64, 64, 96)   384         conv2d_178[0][0]                 \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_252 (LeakyReLU)     (None, 64, 64, 96)   0           batch_normalization_252[0][0]    \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_179 (Conv2D)             (None, 64, 64, 96)   82944       leaky_re_lu_252[0][0]            \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_253 (BatchN (None, 64, 64, 96)   384         conv2d_179[0][0]                 \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_253 (LeakyReLU)     (None, 64, 64, 96)   0           batch_normalization_253[0][0]    \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_180 (Conv2D)             (None, 64, 64, 96)   82944       leaky_re_lu_253[0][0]            \n",
      "__________________________________________________________________________________________________\n",
      "add_42 (Add)                    (None, 64, 64, 96)   0           conv2d_180[0][0]                 \n",
      "                                                                 conv2d_177[0][0]                 \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_254 (BatchN (None, 64, 64, 96)   384         add_42[0][0]                     \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_254 (LeakyReLU)     (None, 64, 64, 96)   0           batch_normalization_254[0][0]    \n",
      "__________________________________________________________________________________________________\n",
      "max_pooling2d_35 (MaxPooling2D) (None, 32, 32, 96)   0           leaky_re_lu_254[0][0]            \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_255 (BatchN (None, 32, 32, 96)   384         max_pooling2d_35[0][0]           \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_255 (LeakyReLU)     (None, 32, 32, 96)   0           batch_normalization_255[0][0]    \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_181 (Conv2D)             (None, 32, 32, 192)  165888      leaky_re_lu_255[0][0]            \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_256 (BatchN (None, 32, 32, 192)  768         conv2d_181[0][0]                 \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_256 (LeakyReLU)     (None, 32, 32, 192)  0           batch_normalization_256[0][0]    \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_182 (Conv2D)             (None, 32, 32, 192)  331776      leaky_re_lu_256[0][0]            \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_257 (BatchN (None, 32, 32, 192)  768         conv2d_182[0][0]                 \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_257 (LeakyReLU)     (None, 32, 32, 192)  0           batch_normalization_257[0][0]    \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_183 (Conv2D)             (None, 32, 32, 192)  331776      leaky_re_lu_257[0][0]            \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_258 (BatchN (None, 32, 32, 192)  768         conv2d_183[0][0]                 \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_258 (LeakyReLU)     (None, 32, 32, 192)  0           batch_normalization_258[0][0]    \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_184 (Conv2D)             (None, 32, 32, 192)  331776      leaky_re_lu_258[0][0]            \n",
      "__________________________________________________________________________________________________\n",
      "add_43 (Add)                    (None, 32, 32, 192)  0           conv2d_184[0][0]                 \n",
      "                                                                 conv2d_181[0][0]                 \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_259 (BatchN (None, 32, 32, 192)  768         add_43[0][0]                     \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_259 (LeakyReLU)     (None, 32, 32, 192)  0           batch_normalization_259[0][0]    \n",
      "__________________________________________________________________________________________________\n",
      "max_pooling2d_36 (MaxPooling2D) (None, 16, 16, 192)  0           leaky_re_lu_259[0][0]            \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_260 (BatchN (None, 16, 16, 192)  768         max_pooling2d_36[0][0]           \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_260 (LeakyReLU)     (None, 16, 16, 192)  0           batch_normalization_260[0][0]    \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_185 (Conv2D)             (None, 16, 16, 384)  663552      leaky_re_lu_260[0][0]            \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_261 (BatchN (None, 16, 16, 384)  1536        conv2d_185[0][0]                 \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_261 (LeakyReLU)     (None, 16, 16, 384)  0           batch_normalization_261[0][0]    \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_186 (Conv2D)             (None, 16, 16, 384)  1327104     leaky_re_lu_261[0][0]            \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_262 (BatchN (None, 16, 16, 384)  1536        conv2d_186[0][0]                 \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_262 (LeakyReLU)     (None, 16, 16, 384)  0           batch_normalization_262[0][0]    \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_187 (Conv2D)             (None, 16, 16, 384)  1327104     leaky_re_lu_262[0][0]            \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_263 (BatchN (None, 16, 16, 384)  1536        conv2d_187[0][0]                 \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_263 (LeakyReLU)     (None, 16, 16, 384)  0           batch_normalization_263[0][0]    \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_188 (Conv2D)             (None, 16, 16, 384)  1327104     leaky_re_lu_263[0][0]            \n",
      "__________________________________________________________________________________________________\n",
      "add_44 (Add)                    (None, 16, 16, 384)  0           conv2d_188[0][0]                 \n",
      "                                                                 conv2d_185[0][0]                 \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_264 (BatchN (None, 16, 16, 384)  1536        add_44[0][0]                     \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_264 (LeakyReLU)     (None, 16, 16, 384)  0           batch_normalization_264[0][0]    \n",
      "__________________________________________________________________________________________________\n",
      "dilated_conv_1 (Conv2D)         (None, 16, 16, 512)  1769984     leaky_re_lu_264[0][0]            \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_265 (BatchN (None, 16, 16, 512)  2048        dilated_conv_1[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_265 (LeakyReLU)     (None, 16, 16, 512)  0           batch_normalization_265[0][0]    \n",
      "__________________________________________________________________________________________________\n",
      "dilated_conv_2 (Conv2D)         (None, 16, 16, 512)  2359808     leaky_re_lu_265[0][0]            \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_266 (BatchN (None, 16, 16, 512)  2048        dilated_conv_2[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_266 (LeakyReLU)     (None, 16, 16, 512)  0           batch_normalization_266[0][0]    \n",
      "__________________________________________________________________________________________________\n",
      "last_conv (Conv2D)              (None, 16, 16, 512)  2359808     leaky_re_lu_266[0][0]            \n",
      "__________________________________________________________________________________________________\n",
      "add_45 (Add)                    (None, 16, 16, 512)  0           last_conv[0][0]                  \n",
      "                                                                 dilated_conv_1[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_267 (BatchN (None, 16, 16, 512)  2048        add_45[0][0]                     \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_267 (LeakyReLU)     (None, 16, 16, 512)  0           batch_normalization_267[0][0]    \n",
      "__________________________________________________________________________________________________\n",
      "downsample_1 (Conv2D)           (None, 4, 4, 512)    4194816     leaky_re_lu_267[0][0]            \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_268 (BatchN (None, 4, 4, 512)    2048        downsample_1[0][0]               \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_268 (LeakyReLU)     (None, 4, 4, 512)    0           batch_normalization_268[0][0]    \n",
      "__________________________________________________________________________________________________\n",
      "average_pool (AveragePooling2D) (None, 2, 2, 512)    0           leaky_re_lu_268[0][0]            \n",
      "__________________________________________________________________________________________________\n",
      "fc_1_b (Conv2D)                 (None, 2, 2, 1024)   525312      average_pool[0][0]               \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_269 (BatchN (None, 2, 2, 1024)   4096        fc_1_b[0][0]                     \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_269 (LeakyReLU)     (None, 2, 2, 1024)   0           batch_normalization_269[0][0]    \n",
      "__________________________________________________________________________________________________\n",
      "dropout_9 (Dropout)             (None, 2, 2, 1024)   0           leaky_re_lu_269[0][0]            \n",
      "__________________________________________________________________________________________________\n",
      "fc_2_b (Conv2D)                 (None, 2, 2, 1024)   1049600     dropout_9[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_270 (BatchN (None, 2, 2, 1024)   4096        fc_2_b[0][0]                     \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_270 (LeakyReLU)     (None, 2, 2, 1024)   0           batch_normalization_270[0][0]    \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_189 (Conv2D)             (None, 2, 2, 5)      5125        leaky_re_lu_270[0][0]            \n",
      "==================================================================================================\n",
      "Total params: 18,474,301\n",
      "Trainable params: 18,459,437\n",
      "Non-trainable params: 14,864\n",
      "__________________________________________________________________________________________________\n",
      "None\n"
     ]
    }
   ],
   "source": [
    "GRID_SIZE = 2\n",
    "CELL_SIZE = 1024 // GRID_SIZE\n",
    "\n",
    "def binary_accuracy(y_true, y_pred):\n",
    "    y_pred = tf.sigmoid(y_pred)\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 * 256\n",
    "    x21 = x21 * 256\n",
    "    y11 = y11 * 256\n",
    "    y21 = y21 * 256\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 == 1)\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, 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",
    "    # apply the sigmoid\n",
    "    y_pred = tf.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 == 1:\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",
    "        # 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",
    "        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 (xe_loss + (box_loss * 3.0) + (iou_loss * 2.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.equal(conf_true, 1)\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 = 3.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 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] * 256\n",
    "    true_box_wh = y_true[...,3:] * 1024\n",
    "    \n",
    "    pred_box_xy = y_pred[...,1:3] * 256\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.equal(y_true[...,0], 1)\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=1.5, 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.002\n",
    "    epochs_drop = 10\n",
    "    drop = 0.75\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",
    "\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": 59,
   "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 6/10\n",
      "1443/1444 [============================>.] - ETA: 0s - loss: 24.9172 - binary_accuracy: 0.8178 - overlap_iou: 0.1322 - adj_mse: 3.4210 - binary_cross_entropy: 8.3833\n",
      "Epoch 00006: saving model to yolo10_2_512.h5\n",
      "1444/1444 [==============================] - 1170s 810ms/step - loss: 24.9120 - binary_accuracy: 0.8177 - overlap_iou: 0.1322 - adj_mse: 3.4199 - binary_cross_entropy: 8.3820 - val_loss: 22.9564 - val_binary_accuracy: 0.8459 - val_overlap_iou: 0.1526 - val_adj_mse: 3.1555 - val_binary_cross_entropy: 7.5580\n",
      "Epoch 7/10\n",
      "1443/1444 [============================>.] - ETA: 0s - loss: 24.2088 - binary_accuracy: 0.8217 - overlap_iou: 0.1397 - adj_mse: 3.3093 - binary_cross_entropy: 8.1685\n",
      "Epoch 00007: saving model to yolo10_2_512.h5\n",
      "1444/1444 [==============================] - 1131s 783ms/step - loss: 24.2100 - binary_accuracy: 0.8217 - overlap_iou: 0.1397 - adj_mse: 3.3096 - binary_cross_entropy: 8.1683 - val_loss: 21.1850 - val_binary_accuracy: 0.7773 - val_overlap_iou: 0.1303 - val_adj_mse: 2.9122 - val_binary_cross_entropy: 7.0486\n",
      "Epoch 8/10\n",
      "1373/1444 [===========================>..] - ETA: 52s - loss: 23.3759 - binary_accuracy: 0.8295 - overlap_iou: 0.1482 - adj_mse: 3.1802 - binary_cross_entropy: 7.8542"
     ]
    }
   ],
   "source": [
    "history = model.fit_generator(train_gen, validation_data=valid_gen, callbacks=[learning_rate, checkpoint], epochs=10, shuffle=True, verbose=1, initial_epoch=5)"
   ]
  },
  {
   "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], epochs=15, shuffle=True, verbose=1, initial_epoch=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "history = model.fit_generator(train_gen, validation_data=valid_gen, callbacks=[learning_rate, checkpoint], epochs=25, shuffle=True, verbose=1, initial_epoch=15)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "history = model.fit_generator(train_gen, validation_data=valid_gen, callbacks=[learning_rate, checkpoint], epochs=30, shuffle=True, verbose=1, initial_epoch=25)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "history = model.fit_generator(train_gen, validation_data=valid_gen, callbacks=[learning_rate, checkpoint], epochs=40, shuffle=True, verbose=1, initial_epoch=30)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "history = model.fit_generator(train_gen, validation_data=valid_gen, callbacks=[learning_rate, checkpoint], epochs=45, shuffle=True, verbose=1, initial_epoch=40)"
   ]
  },
  {
   "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(151)\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(152)\n",
    "plt.plot(history.epoch, history.history[\"iou_loss_fn\"], label=\"Train IOU loss\")\n",
    "plt.plot(history.epoch, history.history[\"val_iou_loss_fn\"], label=\"Valid IOU loss\")\n",
    "plt.legend()\n",
    "plt.subplot(153)\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(154)\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(155)\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": [
    "THRESHOLD = 0.5\n",
    "OVERLAP = 0.3\n",
    "# look at some sample predictions\n",
    "samples = np.random.choice(valid_filenames, size=10, replace=False)\n",
    "\n",
    "coords = np.arange(0, 1024, 512)\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]\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, 512)\n",
    "                    y = unnorm(y, i, 512)\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",
    "    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, 512)\n",
    "                y = unnorm(y, i, 512)\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",
    "    nms_boxes = non_max_suppression_fast(np.array(boxes), OVERLAP)\n",
    "    \n",
    "    # plot our boxes\n",
    "    for box in nms_boxes:\n",
    "        x,y,w,h = box\n",
    "        print(x, y, w, h)\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",
    "            # check if the boxes overlap at all\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",
    "            \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.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.5\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=3, 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(4):\n",
    "            for j in range(4):\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, 256)\n",
    "                        y = unnorm(y, i, 256)\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 \"YOLOv8 512x512 35 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
}