yolo_512x512_8_working.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",
    "* 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 confidence lass has been greatly weighted to account for the fact that not every image contains an ROI so most cells are empty. Without a high weight on the confidence the model predicts all cells as negative which drives the gradients to -inf\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",
    "\n",
    "## Training\n",
    "* The IOU/Confidence loss tries to push the confidence to 0 if the IOU is low. With only one bounding box per cell this doesn't really make sense, and it will counteract the classification loss. So we trained without adding this to the loss. Once the boxes have decent predictions we may add it back in.\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."
   ]
  },
  {
   "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\")"
   ]
  },
  {
   "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": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sigmoid(x):\n",
    "    y = np.maximum(x, -700)\n",
    "    return 1 / (1 + np.exp(-y))"
   ]
  },
  {
   "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": 5,
   "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": 6,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "hopsRmICcxSs"
   },
   "outputs": [],
   "source": [
    "with open('yolo_labels_centered_4x4_6.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": 7,
   "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": "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 = 14\n",
    "IMAGE_SIZE = 512\n",
    "CHECKPOINT_PATH = \"yolo8_2_512.h5\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "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": 10,
   "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",
    "        # resize both image and mask\n",
    "        img = resize(img, (self.image_size, self.image_size), mode='reflect')\n",
    "        \n",
    "        # center the image\n",
    "        img = img - 0.5\n",
    "        \n",
    "        boxes = np.concatenate([confs.reshape((4,4,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",
    "            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": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/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"
     ]
    }
   ],
   "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": 12,
   "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",
    "#     # 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.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": 18,
   "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_3 (InputLayer)            (None, 512, 512, 1)  0                                            \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_43 (Conv2D)              (None, 256, 256, 24) 216         input_3[0][0]                    \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_59 (BatchNo (None, 256, 256, 24) 96          conv2d_43[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_61 (LeakyReLU)      (None, 256, 256, 24) 0           batch_normalization_59[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_44 (Conv2D)              (None, 256, 256, 24) 5184        leaky_re_lu_61[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_60 (BatchNo (None, 256, 256, 24) 96          conv2d_44[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_62 (LeakyReLU)      (None, 256, 256, 24) 0           batch_normalization_60[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_45 (Conv2D)              (None, 256, 256, 24) 5184        leaky_re_lu_62[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_61 (BatchNo (None, 256, 256, 24) 96          conv2d_45[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_63 (LeakyReLU)      (None, 256, 256, 24) 0           batch_normalization_61[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_46 (Conv2D)              (None, 256, 256, 24) 576         leaky_re_lu_63[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_62 (BatchNo (None, 256, 256, 24) 96          conv2d_46[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_64 (LeakyReLU)      (None, 256, 256, 24) 0           batch_normalization_62[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "max_pooling2d_9 (MaxPooling2D)  (None, 128, 128, 24) 0           leaky_re_lu_64[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_63 (BatchNo (None, 128, 128, 24) 96          max_pooling2d_9[0][0]            \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_65 (LeakyReLU)      (None, 128, 128, 24) 0           batch_normalization_63[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_47 (Conv2D)              (None, 128, 128, 48) 10368       leaky_re_lu_65[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_64 (BatchNo (None, 128, 128, 48) 192         conv2d_47[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_66 (LeakyReLU)      (None, 128, 128, 48) 0           batch_normalization_64[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_48 (Conv2D)              (None, 128, 128, 48) 20736       leaky_re_lu_66[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_65 (BatchNo (None, 128, 128, 48) 192         conv2d_48[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_67 (LeakyReLU)      (None, 128, 128, 48) 0           batch_normalization_65[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_49 (Conv2D)              (None, 128, 128, 48) 20736       leaky_re_lu_67[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_66 (BatchNo (None, 128, 128, 48) 192         conv2d_49[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_68 (LeakyReLU)      (None, 128, 128, 48) 0           batch_normalization_66[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_50 (Conv2D)              (None, 128, 128, 48) 20736       leaky_re_lu_68[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "add_11 (Add)                    (None, 128, 128, 48) 0           conv2d_50[0][0]                  \n",
      "                                                                 conv2d_47[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_67 (BatchNo (None, 128, 128, 48) 192         add_11[0][0]                     \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_69 (LeakyReLU)      (None, 128, 128, 48) 0           batch_normalization_67[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "max_pooling2d_10 (MaxPooling2D) (None, 64, 64, 48)   0           leaky_re_lu_69[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_68 (BatchNo (None, 64, 64, 48)   192         max_pooling2d_10[0][0]           \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_70 (LeakyReLU)      (None, 64, 64, 48)   0           batch_normalization_68[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_51 (Conv2D)              (None, 64, 64, 96)   41472       leaky_re_lu_70[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_69 (BatchNo (None, 64, 64, 96)   384         conv2d_51[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_71 (LeakyReLU)      (None, 64, 64, 96)   0           batch_normalization_69[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_52 (Conv2D)              (None, 64, 64, 96)   82944       leaky_re_lu_71[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_70 (BatchNo (None, 64, 64, 96)   384         conv2d_52[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_72 (LeakyReLU)      (None, 64, 64, 96)   0           batch_normalization_70[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_53 (Conv2D)              (None, 64, 64, 96)   82944       leaky_re_lu_72[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_71 (BatchNo (None, 64, 64, 96)   384         conv2d_53[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_73 (LeakyReLU)      (None, 64, 64, 96)   0           batch_normalization_71[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_54 (Conv2D)              (None, 64, 64, 96)   82944       leaky_re_lu_73[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "add_12 (Add)                    (None, 64, 64, 96)   0           conv2d_54[0][0]                  \n",
      "                                                                 conv2d_51[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_72 (BatchNo (None, 64, 64, 96)   384         add_12[0][0]                     \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_74 (LeakyReLU)      (None, 64, 64, 96)   0           batch_normalization_72[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "max_pooling2d_11 (MaxPooling2D) (None, 32, 32, 96)   0           leaky_re_lu_74[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_73 (BatchNo (None, 32, 32, 96)   384         max_pooling2d_11[0][0]           \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_75 (LeakyReLU)      (None, 32, 32, 96)   0           batch_normalization_73[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_55 (Conv2D)              (None, 32, 32, 192)  165888      leaky_re_lu_75[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_74 (BatchNo (None, 32, 32, 192)  768         conv2d_55[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_76 (LeakyReLU)      (None, 32, 32, 192)  0           batch_normalization_74[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_56 (Conv2D)              (None, 32, 32, 192)  331776      leaky_re_lu_76[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_75 (BatchNo (None, 32, 32, 192)  768         conv2d_56[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_77 (LeakyReLU)      (None, 32, 32, 192)  0           batch_normalization_75[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_57 (Conv2D)              (None, 32, 32, 192)  331776      leaky_re_lu_77[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_76 (BatchNo (None, 32, 32, 192)  768         conv2d_57[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_78 (LeakyReLU)      (None, 32, 32, 192)  0           batch_normalization_76[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_58 (Conv2D)              (None, 32, 32, 192)  331776      leaky_re_lu_78[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "add_13 (Add)                    (None, 32, 32, 192)  0           conv2d_58[0][0]                  \n",
      "                                                                 conv2d_55[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_77 (BatchNo (None, 32, 32, 192)  768         add_13[0][0]                     \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_79 (LeakyReLU)      (None, 32, 32, 192)  0           batch_normalization_77[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "max_pooling2d_12 (MaxPooling2D) (None, 16, 16, 192)  0           leaky_re_lu_79[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_78 (BatchNo (None, 16, 16, 192)  768         max_pooling2d_12[0][0]           \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_80 (LeakyReLU)      (None, 16, 16, 192)  0           batch_normalization_78[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_59 (Conv2D)              (None, 16, 16, 384)  663552      leaky_re_lu_80[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_79 (BatchNo (None, 16, 16, 384)  1536        conv2d_59[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_81 (LeakyReLU)      (None, 16, 16, 384)  0           batch_normalization_79[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_60 (Conv2D)              (None, 16, 16, 384)  1327104     leaky_re_lu_81[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_80 (BatchNo (None, 16, 16, 384)  1536        conv2d_60[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_82 (LeakyReLU)      (None, 16, 16, 384)  0           batch_normalization_80[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_61 (Conv2D)              (None, 16, 16, 384)  1327104     leaky_re_lu_82[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_81 (BatchNo (None, 16, 16, 384)  1536        conv2d_61[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_83 (LeakyReLU)      (None, 16, 16, 384)  0           batch_normalization_81[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_62 (Conv2D)              (None, 16, 16, 384)  1327104     leaky_re_lu_83[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "add_14 (Add)                    (None, 16, 16, 384)  0           conv2d_62[0][0]                  \n",
      "                                                                 conv2d_59[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_82 (BatchNo (None, 16, 16, 384)  1536        add_14[0][0]                     \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_84 (LeakyReLU)      (None, 16, 16, 384)  0           batch_normalization_82[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "dilated_conv_1 (Conv2D)         (None, 16, 16, 512)  1769984     leaky_re_lu_84[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_83 (BatchNo (None, 16, 16, 512)  2048        dilated_conv_1[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_85 (LeakyReLU)      (None, 16, 16, 512)  0           batch_normalization_83[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "dilated_conv_2 (Conv2D)         (None, 16, 16, 512)  2359808     leaky_re_lu_85[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_84 (BatchNo (None, 16, 16, 512)  2048        dilated_conv_2[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_86 (LeakyReLU)      (None, 16, 16, 512)  0           batch_normalization_84[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "last_conv (Conv2D)              (None, 16, 16, 512)  2359808     leaky_re_lu_86[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "add_15 (Add)                    (None, 16, 16, 512)  0           last_conv[0][0]                  \n",
      "                                                                 dilated_conv_1[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_85 (BatchNo (None, 16, 16, 512)  2048        add_15[0][0]                     \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_87 (LeakyReLU)      (None, 16, 16, 512)  0           batch_normalization_85[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "downsample_1 (Conv2D)           (None, 4, 4, 512)    4194816     leaky_re_lu_87[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_86 (BatchNo (None, 4, 4, 512)    2048        downsample_1[0][0]               \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_88 (LeakyReLU)      (None, 4, 4, 512)    0           batch_normalization_86[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "fc_1_b (Conv2D)                 (None, 4, 4, 1024)   525312      leaky_re_lu_88[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_87 (BatchNo (None, 4, 4, 1024)   4096        fc_1_b[0][0]                     \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_89 (LeakyReLU)      (None, 4, 4, 1024)   0           batch_normalization_87[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "dropout_3 (Dropout)             (None, 4, 4, 1024)   0           leaky_re_lu_89[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "fc_2_b (Conv2D)                 (None, 4, 4, 1024)   1049600     dropout_3[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_90 (LeakyReLU)      (None, 4, 4, 1024)   0           fc_2_b[0][0]                     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_63 (Conv2D)              (None, 4, 4, 5)      5125        leaky_re_lu_90[0][0]             \n",
      "==================================================================================================\n",
      "Total params: 18,470,205\n",
      "Trainable params: 18,457,389\n",
      "Non-trainable params: 12,816\n",
      "__________________________________________________________________________________________________\n",
      "None\n"
     ]
    }
   ],
   "source": [
    "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 a box predicted here? maybe use this for the IOU?\n",
    "    mask = (pred_conf >= 0.5) | (true_boxes == 1)\n",
    "    \n",
    "    # these labels are center points of the boxes so we need to calculate the corners\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",
    "    # reduce the mean so we have mean iou for our inputs\n",
    "    return tf.reduce_mean(tf.maximum(inter_area / union, 0))\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",
    "    return xe_loss + (box_loss * 5.0)  + (iou_loss * 1.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(w_true - w_pred) + tf.square(h_true - h_pred)\n",
    "    \n",
    "    box_se = center_se + size_se\n",
    "    \n",
    "    # only get loss for boxes which are actually positive\n",
    "    mask = tf.cast(tf.equal(conf_true, 1), tf.float32) * 5.0\n",
    "    \n",
    "    box_se = tf.multiply(box_se, mask)\n",
    "    \n",
    "    loss = tf.reduce_sum(box_se) + 1e-16\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.nn.sigmoid_cross_entropy_with_logits(logits=y_pred, labels=y_true)\n",
    "#     xe = tf.multiply(xe, weights)\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",
    "    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] * 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",
    "    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",
    "    nb_conf_box  = tf.cast((conf_mask > 1e-3), 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",
    "    loss_conf  =  (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",
    "# create network and compiler\n",
    "model = create_network(input_size=IMAGE_SIZE, channels=24, n_blocks=1, depth=4)\n",
    "\n",
    "model.compile(optimizer='adam',\n",
    "              loss=loss_fn,\n",
    "              metrics=[binary_accuracy, overlap_iou, iou_loss_fn, adj_mse, binary_cross_entropy])\n",
    "\n",
    "# cosine learning rate annealing\n",
    "def exp_decay(x):\n",
    "    lr0 = 0.001\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": 17,
   "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": [
      "/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",
      " 677/1651 [===========>..................] - ETA: 10:47 - loss: 71.4646 - binary_accuracy: 0.8996 - overlap_iou: 0.0070 - iou_loss_fn: 7.4146 - adj_mse: 5.6607 - binary_cross_entropy: 31.5767"
     ]
    }
   ],
   "source": [
    "history = model.fit_generator(train_gen, validation_data=valid_gen, callbacks=[learning_rate, checkpoint], 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], epochs=20, shuffle=True, verbose=1, initial_epoch=10)"
   ]
  },
  {
   "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(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": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "7dpxfbugzcLT"
   },
   "source": [
    "## Sample Predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def unnorm(val, idx, cell_size=256):\n",
    "    x = (val * cell_size) + (cell_size * idx)\n",
    "    \n",
    "    return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "c-eU6Le2zdtF"
   },
   "outputs": [],
   "source": [
    "THRESHOLD = 0.5\n",
    "OVERLAP = 0.3\n",
    "# look at some sample predictions\n",
    "samples = np.random.choice(valid_filenames, size=8, replace=False)\n",
    "\n",
    "coords = np.arange(0, 1024, 256)\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",
    "    # get boxes for the truth\n",
    "    if filename in pneumonia_locations:\n",
    "        locs = pneumonia_locations[filename]\n",
    "        for i in range(4):\n",
    "            for j in range(4):\n",
    "                pixel_data = locs[i,j,:]\n",
    "#                 print(pixel_data)\n",
    "                if pixel_data[0] > 0.5:\n",
    "                    x, y, w, h = pixel_data[1:]\n",
    "                    print(\"Truth:\", i, j)\n",
    "                    print(\"Truth:\", x, y, w, h)\n",
    "                    \n",
    "                    # unnormalize the data\n",
    "                    w = w * 1024\n",
    "                    h = h * 1024\n",
    "                    \n",
    "                    x = unnorm(x, j, 256)\n",
    "                    y = unnorm(y, i, 256)\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",
    "                    \n",
    "                    print(\"Truth:\", 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",
    "    print(\"Raw:\",yhat)\n",
    "    yhat = sigmoid(yhat)\n",
    "    print(\"Sig:\", yhat)\n",
    "    conf = np.squeeze(yhat)[:,:,0]\n",
    "    bboxes = np.squeeze(yhat)[:,:,1:]\n",
    "    boxes = []\n",
    "    \n",
    "    # loop through our predictions\n",
    "    for i in range(4):\n",
    "        for j in range(4):\n",
    "            conf_ = conf[i,j]\n",
    "#             print(\"C:\", conf_)\n",
    "            # if we have a prediction\n",
    "            if conf_ > THRESHOLD:\n",
    "                x,y,w,h = bboxes[i,j,:]\n",
    "                print(\"Pred:\", i, j, \"conf:\", conf_)\n",
    "                \n",
    "                # unnormalize the data\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 = 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",
    "                # if the boxes have width and height\n",
    "                if w > 30 and h > 30:\n",
    "                    boxes.append([x,y,w,h])\n",
    "    \n",
    "    # do non-max suppression of our boxes\n",
    "    boxes = non_max_suppression_fast(np.array(boxes), OVERLAP)\n",
    "    \n",
    "    # plot our boxes\n",
    "    for box in 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",
    "    for item in coords:\n",
    "        plt.axvline(item, linewidth=0.5)\n",
    "        plt.axhline(item, linewidth=0.5)\n",
    "    \n",
    "    plt.show()"
   ]
  },
  {
   "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",
    "    conf_preds, bbox_preds = model.predict(imgs)\n",
    "    \n",
    "    # loop through batch\n",
    "    for confs, bboxes, filename in zip(conf_preds, bbox_preds, filenames):\n",
    "        predictionString = \"\"\n",
    "        print(filename)\n",
    "        boxes = []\n",
    "        for i in range(4):\n",
    "            for j in range(4):\n",
    "                conf = confs[i, j, 0]\n",
    "                if conf > THRESHOLD:\n",
    "                    x, y, w, h = bboxes[i,j]\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(conf) + ' ' + 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",
    "    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.csv -m \"YOLOv6 512x512 10 epochs\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 68
    },
    "colab_type": "code",
    "id": "NTQ_vDaVfvZR",
    "outputId": "a9eb81fb-572d-4236-b497-6bc44130aa89"
   },
   "outputs": [],
   "source": [
    "save_file_to_drive(\"submission.csv\", \"submission.csv\")\n",
    "save_file_to_drive(CHECKPOINT_PATH, CHECKPOINT_PATH)"
   ]
  },
  {
   "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}/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# upload checkpoint to GCS\n",
    "project_id = 'mammography-198911'\n",
    "bucket_name = 'pneumonia'\n",
    "\n",
    "!gcloud config set project {project_id}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!gsutil cp ./yolo_labels.p 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
}