yolo_480x480_20.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "igEd7obtcxR2"
   },
   "source": [
    "## Approach\n",
    "\n",
    "* This is a simplied version of YOLO. The input image is 512x512. YOLO was intended to train on images which all contained at least one object, most multiple objects. The architecture and loss functions have been adapted for this case, where not every image contains an object and there are at most 4 ROIs per image. \n",
    "* Since the maximum number of ROIs in each image is 4 we decided to use a 2x2 grid to simplify the network. In the case where multiple ROIs occur in the same cell we combine the ROIs into a single ROI by expanding the box to include both. Each cell outputs five values:\n",
    "    * The confidence that there is pneumonia present\n",
    "    * The x, y, w, and h of the bounding box\n",
    "    * We have removed the entire classification section, using the confidence to indicate whether there is an ROI in the cell instead.\n",
    "* A sigmoid is applied to the output of the network to result in values between 0 and 1    \n",
    "* The x and y coordinates are offsets from the upper left corner of each cell, the w and h are percentage of the total width.\n",
    "* The loss function is based on YOLO with some differences:\n",
    "    * The weights of the components have been updated\n",
    "    * The classification loss has been removed.\n",
    "    * The \"objectness\" loss of YOLO tries to make the confidence match the actual IOU. Since our model only outputs one box per cell, this doesn't really make any sense and ends up driving the confidence down to the IOU. We replace this by an IOU loss which subtracts the IOU for that cell from 1.\n",
    "\n",
    "## Network\n",
    "\n",
    "* The network consists of a number of residual blocks with convolutions and downsampling blocks with max pooling.\n",
    "* There are two dilated convolutions at the end of the network to provide context.\n",
    "* The dilated convolutions are followed by two strided convolutions which downsize to 2x2.\n",
    "* Finally there are a series of 1x1 convolutions which output a 2x2x5 tensor.\n",
    "* We are using an Adam optimizer with gradient clipping to avoid exploding gradients, which had been a problem.\n",
    "\n",
    "## Predictions\n",
    "* To generate our predictions we loop through each cell of the output\n",
    "* We unnormalize the output to get the actual values\n",
    "* If the confidence is greater than 0.5 and the box has width and height we append it to the list of candidates.\n",
    "* We apply non-max suppression to remove duplicates.\n",
    "* The remaining boxes are concatenated onto the output string.\n",
    "\n",
    "**Change Log:**\n",
    "* v3 - changing output to 8x8 grid from 16x16; changed model to downsample one more time; adjusted network accordingly. \n",
    "* v4 - changed output to 4x4 grid, no image has more than 3 ROIs so this may work better? \n",
    "    * Using center point of ROI to predict instead of upper left corner.\n",
    "* v5 - We only calculate MSE loss for boxes with a confidence over 0.5 or actual truth since we don't care about predictions for boxes that are not ROIs. This will prevent the network from being constrained by outputting 0s for boxes that don't exist.\n",
    "* v6 - centering input data so maybe bboxes can be output more accurately? Also centering the image\n",
    "* v6.2 - labels have ROI centered in center of cell by default instead of mean location.\n",
    "* v8 - using custom loss function based on YOLO loss. Set default height and width to 1 px because 0 sent the gradients to -inf which screwed everything up.\n",
    "* v10 - using 2x2 grid as output\n",
    "* v12 - changing layout of model slightly - replaced final pools with convs with strides; reduced number of params; ensure that there is at least one positive image per batch.\n",
    "* v13 - tweaking architecture slightly.\n",
    "* v14 - going to 5x5 grid\n",
    "* v15 - downsizing with average and max pool followed by 1x1 convolutions\n",
    "* v15.2 - fixed an error where a layer of the network was being bypassed, removed some layers to prevent overfitting, training on labels with jittered ROIs\n",
    "* v18 - separating fc branches for confidence and boxes, might work better?\n",
    "* v19 - removing double pools and just using a simple max pool, using sigmoid activation and removing redundant sigmoids from code\n",
    "* v20 - adding some more layers, model may be too small"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "BvHm1yolU3Kn"
   },
   "outputs": [],
   "source": [
    "# install dependencies not included by Colab\n",
    "# use pip3 to ensure compatibility w/ Google Deep Learning Images \n",
    "!pip install -q pydicom \n",
    "!pip install -q imgaug \n",
    "!pip install -q kaggle\n",
    "!pip install -q tqdm "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "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",
    "from skimage.filters import gaussian\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": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "NDaVZIJ4U3_L"
   },
   "outputs": [],
   "source": [
    "from google.colab import auth\n",
    "from googleapiclient.http import MediaFileUpload\n",
    "from googleapiclient.discovery import build\n",
    "\n",
    "auth.authenticate_user()\n",
    "\n",
    "drive_service = build('drive', 'v3')\n",
    "\n",
    "def save_file_to_drive(name, path):\n",
    "    file_metadata = {\n",
    "      'name': name,\n",
    "        'mimeType': 'application/octet-stream',\n",
    "     }\n",
    "\n",
    "    media = MediaFileUpload(path, \n",
    "                    mimetype='application/octet-stream',\n",
    "                    resumable=True)\n",
    "\n",
    "    created = drive_service.files().create(body=file_metadata,\n",
    "                                   media_body=media,\n",
    "                                   fields='id').execute()\n",
    "\n",
    "    print('File ID: {}'.format(created.get('id')))\n",
    "\n",
    "    return created"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "jQ34pL62UyT4"
   },
   "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",
    "# edited to do nothing, in case we forget to remove any calls to this function\n",
    "def sigmoid(x):\n",
    "    # to avoid NaNs set a lower floor on x values\n",
    "#     y = np.maximum(x, -700)\n",
    "#     y = np.minimum(y, 700)\n",
    "#     return 1 / (1 + np.exp(-y))    \n",
    "    return x\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\n",
    "\n",
    "def get_intersect(box1, box2):\n",
    "    # unpack each box\n",
    "    x1, y1, w1, h1, c1 = box1\n",
    "    x2, y2, w2, h2, x2 = box2\n",
    "    \n",
    "    # get the far corners\n",
    "    x1_f = x1 + w1\n",
    "    y1_f = y1 + h1\n",
    "    x2_f = x2 + w2\n",
    "    y2_f = y2 + h2\n",
    "    \n",
    "    # get corners of intersection\n",
    "    x1_i = np.maximum(x1, x2)\n",
    "    y1_i = np.maximum(y1, y2)\n",
    "    x2_i = np.minimum(x1_f, x2_f)\n",
    "    y2_i = np.minimum(y1_f, y2_f)\n",
    "    \n",
    "    w_i = x2_i - x1_i\n",
    "    h_i = y2_i - y1_i\n",
    "    \n",
    "    w_i = np.maximum(w_i, 0)\n",
    "    h_i = np.maximum(h_i, 0)\n",
    "\n",
    "    intersect_area = w_i * h_i\n",
    "    \n",
    "    return intersect_area  \n",
    "  \n",
    "# helper function to calculate IoU\n",
    "def calculate_iou(box1, box2):\n",
    "    x11, y11, w1, h1 = box1[:4]\n",
    "    x21, y21, w2, h2 = box2[:4]\n",
    "    assert w1 * h1 > 0\n",
    "    assert w2 * h2 > 0\n",
    "    x12, y12 = x11 + w1, y11 + h1\n",
    "    x22, y22 = x21 + w2, y21 + h2\n",
    "\n",
    "    area1, area2 = w1 * h1, w2 * h2\n",
    "    xi1, yi1, xi2, yi2 = max([x11, x21]), max([y11, y21]), min([x12, x22]), min([y12, y22])\n",
    "    \n",
    "    if xi2 <= xi1 or yi2 <= yi1:\n",
    "        return 0\n",
    "    else:\n",
    "        intersect = (xi2-xi1) * (yi2-yi1)\n",
    "        union = area1 + area2 - intersect\n",
    "        return intersect / union  \n",
    "\n",
    "# calculate IOU where there can be multiple overlapping truths and predictions\n",
    "def calc_iou_better(y_true, y_pred):\n",
    "    label_array = np.zeros((1024, 1024))\n",
    "    \n",
    "    for truth in y_true:\n",
    "        x,y,w,h = truth[:4].astype(int)  \n",
    "        label_array[y:y+h, x:x+w] = 1\n",
    "                \n",
    "    pred_array = np.zeros((1024, 1024))\n",
    "    for pred in y_pred:\n",
    "        x,y,w,h = pred[:4].astype(int)\n",
    "\n",
    "        # update the pixels\n",
    "        pred_array[y:y+h, x:x+w] = 1\n",
    "                \n",
    "    true_area = np.sum(label_array)\n",
    "    pred_area = np.sum(pred_array)\n",
    "    \n",
    "    intersect_area = np.sum((pred_array == 1) & (label_array == 1))\n",
    "    \n",
    "    union = true_area + pred_area - intersect_area\n",
    "    \n",
    "    iou = intersect_area / (union + 1e-6)\n",
    "    \n",
    "    return iou      \n",
    "      \n",
    "def map_iou(boxes_true, boxes_pred, scores, thresholds = [0.4, 0.45, 0.5, 0.55, 0.6, 0.65, 0.7, 0.75]):\n",
    "    \"\"\"\n",
    "    Mean average precision at differnet intersection over union (IoU) threshold\n",
    "    \n",
    "    input:\n",
    "        boxes_true: Mx4 numpy array of ground true bounding boxes of one image. \n",
    "                    bbox format: (x1, y1, w, h)\n",
    "        boxes_pred: Nx4 numpy array of predicted bounding boxes of one image. \n",
    "                    bbox format: (x1, y1, w, h)\n",
    "        scores:     length N numpy array of scores associated with predicted bboxes\n",
    "        thresholds: IoU shresholds to evaluate mean average precision on\n",
    "    output: \n",
    "        map: mean average precision of the image\n",
    "    \"\"\"\n",
    "    \n",
    "    # According to the introduction, images with no ground truth bboxes will not be \n",
    "    # included in the map score unless there is a false positive detection (?)\n",
    "        \n",
    "    # return None if both are empty, don't count the image in final evaluation (?)\n",
    "    if len(boxes_true) == 0 and len(boxes_pred) == 0:\n",
    "        return None\n",
    "    elif (len(boxes_true) == 0) and (len(boxes_pred) != 0):\n",
    "        return 0\n",
    "    \n",
    "    assert boxes_true.shape[1] == 4 or boxes_pred.shape[1] == 4, \"boxes should be 2D arrays with shape[1]=4\"\n",
    "    if len(boxes_pred):\n",
    "        assert len(scores) == len(boxes_pred), \"boxes_pred and scores should be same length\"\n",
    "        # sort boxes_pred by scores in decreasing order\n",
    "        boxes_pred = boxes_pred[np.argsort(scores)[::-1], :]\n",
    "    \n",
    "    map_total = 0\n",
    "    \n",
    "    # loop over thresholds\n",
    "    for t in thresholds:\n",
    "        matched_bt = set()\n",
    "        tp, fn = 0, 0\n",
    "        for i, bt in enumerate(boxes_true):\n",
    "            matched = False\n",
    "            for j, bp in enumerate(boxes_pred):\n",
    "                miou = calculate_iou(bt, bp)\n",
    "                if miou >= t and not matched and j not in matched_bt:\n",
    "                    matched = True\n",
    "                    tp += 1 # bt is matched for the first time, count as TP\n",
    "                    matched_bt.add(j)\n",
    "            if not matched:\n",
    "                fn += 1 # bt has no match, count as FN\n",
    "                \n",
    "        fp = len(boxes_pred) - len(matched_bt) # FP is the bp that not matched to any bt\n",
    "        m = tp / (tp + fn + fp)\n",
    "        map_total += m\n",
    "    \n",
    "    return map_total / len(thresholds)      \n",
    "\n",
    "# given the label or predictions, return a list containing the actual coordinates\n",
    "# for each box\n",
    "def unnorm_label(label, threshold=0.5):\n",
    "    grid_size = label.shape[0]\n",
    "    cell_size = 1024 / grid_size\n",
    "    \n",
    "    boxes = []\n",
    "    # loop through the cells of our label\n",
    "    for i in range(grid_size):\n",
    "        for j in range(grid_size):\n",
    "            # if we have a box in a cell\n",
    "            if label[i,j,0] > threshold:\n",
    "                # unpack the data\n",
    "                c, x, y, w, h = label[i, j]\n",
    "                \n",
    "                # unnormalize the data\n",
    "                w, h = int(w*1024), int(h*1024)\n",
    "                x = unnorm(x, j, cell_size)\n",
    "                y = unnorm(y, i, cell_size)\n",
    "                \n",
    "                # get the corners\n",
    "                x = int(x - (w / 2))\n",
    "                y = int(y - (h / 2))\n",
    "                \n",
    "                # it's possible the corner is in another cell, if so catch the error\n",
    "                x = np.maximum(x, 0)\n",
    "                y = np.maximum(y, 0)\n",
    "                    \n",
    "                print(x, y, w, h, \"box:\", i, j)\n",
    "                boxes.append([x, y, w, h, c])\n",
    "    \n",
    "    return boxes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "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": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 323
    },
    "colab_type": "code",
    "id": "uj9MJL1lU8zT",
    "outputId": "2678c476-1915-4032-b31a-76f8b3c605a9"
   },
   "outputs": [],
   "source": [
    "# If you are unable to download the competition dataset, check to see if you have \n",
    "# accepted the user agreement on the competition website. \n",
    "if not os.path.exists(\"./stage_1_detailed_class_info.csv.zip\"):\n",
    "    !kaggle competitions download -c rsna-pneumonia-detection-challenge"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "LuurZeb6U9J6",
    "outputId": "2bf30f89-3cae-4795-c969-63d74b52803c"
   },
   "outputs": [],
   "source": [
    "# unzipping takes a few minutes\n",
    "if not os.path.exists(\"./stage_1_train_labels.csv\"):\n",
    "    print(\"Unzipping...\")\n",
    "    !unzip -q -o stage_1_test_images.zip -d stage_1_test_images\n",
    "    !unzip -q -o stage_1_train_images.zip -d stage_1_train_images\n",
    "    !unzip -q -o stage_1_train_labels.csv.zip"
   ]
  },
  {
   "cell_type": "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": 36,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "m_0ZGQyLdQS8",
    "outputId": "78704fa9-8a91-473e-b610-741951df68bd"
   },
   "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": 37,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 85
    },
    "colab_type": "code",
    "id": "KeFB-TfoVDZf",
    "outputId": "2743930e-6cfb-4ca1-ea37-c3851da87f79"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Updated property [core/project].\n",
      "Your credentials are invalid. Please run\n",
      "$ gcloud auth login\n"
     ]
    }
   ],
   "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 gs://{bucket_name}/yolo_labels_centered_5x5_10a.p ./yolo_labels_centered_5x5_10a.p"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "hopsRmICcxSs"
   },
   "outputs": [],
   "source": [
    "with open('yolo_labels_centered_5x5_10a.p', 'rb') as handle:\n",
    "    pneumonia_locations = pickle.load(handle)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "UoB6GBhIcxS6"
   },
   "source": [
    "# Load filenames"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 51
    },
    "colab_type": "code",
    "id": "b_iB-IftcxTE",
    "outputId": "4c7107e8-2527-46f5-883c-f53612c40478"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "n train samples 451\n",
      "n valid samples 79\n"
     ]
    }
   ],
   "source": [
    "random.seed(1702)\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.15)\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": 40,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "j9qMFILJUyUR"
   },
   "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": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "dFbLdFxVehhB"
   },
   "outputs": [],
   "source": [
    "BATCH_SIZE = 16\n",
    "IMAGE_SIZE = 480\n",
    "GRID_SIZE = 5\n",
    "CELL_SIZE = 1024 / GRID_SIZE\n",
    "CHECKPOINT_PATH = \"yolo19_1_480.h5\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 68
    },
    "colab_type": "code",
    "id": "QnPfymkWV5rY",
    "outputId": "cef9ce5b-f2b0-4aaf-fa0a-6883fb920465"
   },
   "outputs": [],
   "source": [
    "!gsutil cp gs://{bucket_name}/{CHECKPOINT_PATH} ./{CHECKPOINT_PATH}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "TueY1bVlcxTg"
   },
   "outputs": [],
   "source": [
    "class generator(keras.utils.Sequence):\n",
    "    \n",
    "    def __init__(self, folder, filenames, pneumonia_locations=None, batch_size=BATCH_SIZE, image_size=IMAGE_SIZE, shuffle=True, augment=False, predict=False):\n",
    "        self.folder = folder\n",
    "        self.filenames = filenames\n",
    "        self.pneumonia_locations = pneumonia_locations\n",
    "        self.batch_size = batch_size\n",
    "        self.image_size = image_size\n",
    "        self.shuffle = shuffle\n",
    "        self.augment = augment\n",
    "        self.predict = predict\n",
    "        self.on_epoch_end()\n",
    "    \n",
    "    def augment_imgs(self, img, confs, boxes):\n",
    "        # flip the image half the time\n",
    "        if random.random() > 0.5:\n",
    "            img = np.fliplr(img)\n",
    "            \n",
    "            # update our x coords\n",
    "            mask = (confs != 0)\n",
    "            \n",
    "            # flip\n",
    "            boxes[mask, 0] = (1 - boxes[mask,0])\n",
    "            \n",
    "            # flip our boxes lr on axis 0\n",
    "            boxes = np.flip(boxes, axis=1)   \n",
    "            \n",
    "            # flip the confidences lr as well\n",
    "            confs = np.flip(confs, axis=1)            \n",
    "        \n",
    "        # slight shape skew\n",
    "        if random.random() > 0.5:\n",
    "            x_scale = np.random.uniform(low=0.96, high=1.04, size=None)\n",
    "            y_scale = np.random.uniform(low=0.96, high=1.04, size=None)\n",
    "            \n",
    "            # resize the image\n",
    "            img = resize(img, (int(1024 * y_scale), int(1024 * x_scale)), mode='reflect')\n",
    "            \n",
    "            # resize the boxes\n",
    "            boxes[...,2] = (boxes[...,2] * x_scale)\n",
    "            boxes[...,3] = (boxes[...,3] * y_scale)\n",
    "         \n",
    "        # adjust contrast or contrast\n",
    "        if random.random() > 0.5:  \n",
    "            # adjust brightness\n",
    "            if random.random() > 0.5:  \n",
    "                factor = np.random.uniform(low=0.85, high=1.15)\n",
    "                img = img * factor\n",
    "            # adjust contrast\n",
    "            else:\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",
    "        # gaussian blur\n",
    "        if random.random() > 0.5:\n",
    "            sigma = np.random.uniform(low=0, high=1.25)\n",
    "            img = gaussian(img, sigma=sigma)\n",
    "        \n",
    "        # small shifts\n",
    "        if random.random() > 0.5:  \n",
    "            ## small random shifts\n",
    "            h_offset = np.random.randint(low=0, high=5)\n",
    "            v_offset = np.random.randint(low=0, high=5)\n",
    "\n",
    "            # crop the images\n",
    "            img = img[v_offset:,h_offset:]\n",
    "        \n",
    "        return img, confs, boxes\n",
    "            \n",
    "            \n",
    "    def __load__(self, filename):\n",
    "        # load dicom file as numpy array\n",
    "        img = pydicom.dcmread(os.path.join(self.folder, filename)).pixel_array\n",
    "        # get filename without extension\n",
    "        filename = filename.split('.')[0]\n",
    "        label = pneumonia_locations[filename].copy()\n",
    "        \n",
    "        # remove the confidence and bboxes because they will be flipped separately\n",
    "        # round the confidences since some are 0.99 or maybe even 1e-6\n",
    "        confs = np.round(label[:,:,0])\n",
    "        boxes = label[:,:,1:]\n",
    "        \n",
    "        ## augment the data with flips, small shifts and contrast adjustment\n",
    "        if self.augment:\n",
    "            img, confs, boxes = self.augment_imgs(img, confs, boxes)\n",
    "            \n",
    "        # resize both image and mask\n",
    "        img = resize(img, (self.image_size, self.image_size), mode='reflect')\n",
    "        \n",
    "        # scale and center the image\n",
    "        img = (img - np.mean(img)) / np.max(img)\n",
    "        \n",
    "        boxes = np.concatenate([confs.reshape((GRID_SIZE,GRID_SIZE,1)), boxes], axis=2)\n",
    "        \n",
    "        # add trailing channel dimension\n",
    "        img = np.expand_dims(img, -1)\n",
    "        return img, boxes\n",
    "    \n",
    "    def __loadpredict__(self, filename):\n",
    "        # load dicom file as numpy array\n",
    "        img = pydicom.dcmread(os.path.join(self.folder, filename)).pixel_array\n",
    "        # resize image\n",
    "        img = resize(img, (self.image_size, self.image_size), mode='reflect')\n",
    "        # add trailing channel dimension\n",
    "        img = np.expand_dims(img, -1)\n",
    "        return img\n",
    "        \n",
    "    def __getitem__(self, index):\n",
    "        # select batch\n",
    "        filenames = self.filenames[index*self.batch_size:(index+1)*self.batch_size]\n",
    "        # predict mode: return images and filenames\n",
    "        if self.predict:\n",
    "            # load files\n",
    "            imgs = [self.__loadpredict__(filename) for filename in filenames]\n",
    "            # create numpy batch\n",
    "            imgs = np.array(imgs)\n",
    "            return imgs, filenames\n",
    "        # train mode: return images and masks\n",
    "        else:\n",
    "            # load files\n",
    "            items = [self.__load__(filename) for filename in filenames]\n",
    "            \n",
    "            # unzip images and masks\n",
    "            imgs, bboxes = zip(*items)\n",
    "            \n",
    "            # create numpy batch\n",
    "            imgs = np.array(imgs)\n",
    "            bboxes = np.array(bboxes)\n",
    "            \n",
    "            # make sure there is at least one positive image in the batch\n",
    "            pos = np.max(bboxes[:,:,:,0])\n",
    "            if pos < 1:\n",
    "                # pick a random positive image\n",
    "                filename = np.random.choice(positive_images)\n",
    "                img, label = self.__load__(filename)\n",
    "                \n",
    "                # add the positive image to our batch\n",
    "                imgs[-1] = img\n",
    "                bboxes[-1] = label\n",
    "                \n",
    "            labels = bboxes\n",
    "            return imgs, labels\n",
    "        \n",
    "    def on_epoch_end(self):\n",
    "        if self.shuffle:\n",
    "            random.shuffle(self.filenames)\n",
    "        \n",
    "    def __len__(self):\n",
    "        if self.predict:\n",
    "            # return everything\n",
    "            return int(np.ceil(len(self.filenames) / self.batch_size))\n",
    "        else:\n",
    "            # return full batches only\n",
    "            return int(len(self.filenames) / self.batch_size)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "3_GgDL7ncxT3"
   },
   "source": [
    "# Network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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",
    "    \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",
    "    \n",
    "    x_2 = keras.layers.Conv2D(channels, 3, padding='same', use_bias=False)(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(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.add([x, x_1])\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.add([x, x_2])\n",
    "    return x\n",
    "\n",
    "def create_network(input_size, channels, n_blocks=2, depth=4):\n",
    "    # input - 512x512x3\n",
    "    inputs = keras.Input(shape=(input_size, input_size, 1))\n",
    "    \n",
    "    # downsample to 256x256x24\n",
    "    x = keras.layers.Conv2D(channels, 3, strides=(2,2), padding='same', use_bias=False)(inputs)\n",
    "    x = keras.layers.BatchNormalization(momentum=0.9)(x)\n",
    "    x = keras.layers.LeakyReLU(0.1)(x)\n",
    "    \n",
    "    x = keras.layers.Conv2D(channels, 3, padding='same', use_bias=False)(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",
    "    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",
    "    # residual blocks\n",
    "    for d in range(depth):\n",
    "        \n",
    "        x = create_downsample(channels, x)\n",
    "        channels = channels * 2\n",
    "        \n",
    "        for b in range(n_blocks):\n",
    "            x = create_resblock(channels, x)\n",
    "        \n",
    "    x_2 = x\n",
    "    x = keras.layers.BatchNormalization(momentum=0.9)(x)\n",
    "    x = keras.layers.LeakyReLU(0)(x)\n",
    "    \n",
    "    # dilated convolutions for context - 15x15x512\n",
    "    x = keras.layers.Conv2D(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)\n",
    "    x = keras.layers.LeakyReLU(0.1)(x)\n",
    "    \n",
    "    x = keras.layers.Conv2D(512, (3,3), padding='same', strides=(1,1), activation=None, name=\"last_conv\")(x)\n",
    "    x = keras.layers.BatchNormalization(momentum=0.9)(x)\n",
    "    x = keras.layers.LeakyReLU(0.1)(x)\n",
    "    x = keras.layers.Dropout(0.15)(x)\n",
    "    \n",
    "    # downsample to 8x8 with stride 2\n",
    "    x = keras.layers.Conv2D(768, (3,3), padding='same', strides=(3,3), activation=None, name=\"downsample_1\", kernel_regularizer=keras.regularizers.l2(l=0.002))(x)\n",
    "    x = keras.layers.BatchNormalization(momentum=0.9)(x)\n",
    "    x = keras.layers.LeakyReLU(0.1)(x)\n",
    "    x = keras.layers.Dropout(0.15)(x)\n",
    "    \n",
    "    # confidence output branch\n",
    "    c = keras.layers.Conv2D(1024, (1,1), padding='same', activation=None, name=\"fc_1_c\", kernel_regularizer=keras.regularizers.l2(l=0.005))(x)\n",
    "    c = keras.layers.BatchNormalization(momentum=0.9)(c)\n",
    "    c = keras.layers.LeakyReLU(0.01)(c)\n",
    "    c = keras.layers.Dropout(0.10)(c)\n",
    "    \n",
    "    confidence = keras.layers.Conv2D(1, (1,1), padding='same', activation=\"sigmoid\", name=\"confidence_output\")(c)\n",
    "    \n",
    "    # bounding box branch\n",
    "    b = keras.layers.Conv2D(1568, (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.25)(b)\n",
    "    \n",
    "    b = keras.layers.Conv2D(1568, (1,1), padding='same', activation=None, name=\"fc_2_b\", kernel_regularizer=keras.regularizers.l2(l=0.003))(b)\n",
    "    b = keras.layers.BatchNormalization(momentum=0.9)(b)\n",
    "    b = keras.layers.LeakyReLU(0.01)(b)\n",
    "    \n",
    "    boxes = keras.layers.Conv2D(4, (1,1), padding='same', activation=\"sigmoid\")(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": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Abs2CM3jUyUt"
   },
   "outputs": [],
   "source": [
    "# after each epoch run our function to calculate the actual IOU for half of the validation images. \n",
    "# I was not able to figure out how to implement this in tensorflow since I don't think we can iterate\n",
    "# through tensors. So we do it as a callback instead. Hopefully this will approximate the actual IOU score.\n",
    "class Calc_IOU_CB(keras.callbacks.Callback):\n",
    "    \n",
    "    def on_epoch_end(self, epoch, logs={}):\n",
    "        valid_gen2 = generator(folder, valid_filenames, pneumonia_locations, batch_size=32, image_size=IMAGE_SIZE, shuffle=True, predict=False)\n",
    "        ious = []\n",
    "        tps = 0\n",
    "        fps = 0\n",
    "        fns = 0\n",
    "        \n",
    "        counter = 0\n",
    "        for imgs, labels in valid_gen2:\n",
    "            preds = self.model.predict(imgs)\n",
    "            \n",
    "            for pred, label in zip(preds, labels):\n",
    "                yhat = unnorm_label(pred)\n",
    "                y_true = unnorm_label(label)\n",
    "                \n",
    "                if len(yhat) or len(y_true):\n",
    "                    iou = calc_iou_better(y_true, yhat)\n",
    "                    # calculate the IOU\n",
    "#                     iou = overlap_iou2(labels, preds)\n",
    "                    ious.append(iou)\n",
    "            \n",
    "            # calculate the precision and recall\n",
    "            tp, fp, fn = calc_tps(labels, preds)\n",
    "            \n",
    "            tps += tp\n",
    "            fps += fp\n",
    "            fns += fn\n",
    "            \n",
    "            counter += BATCH_SIZE\n",
    "            \n",
    "            if counter > 900:\n",
    "                break\n",
    "        \n",
    "        prec = tps / (tps + fps)\n",
    "        recall = tps / (tps + fns)\n",
    "        f1_score = 2 * (prec * recall) / (prec + recall)\n",
    "        print(\"Epoch\", epoch, \": Mean IOU:\", np.mean(ious), \"Precision:\", prec, \"Recall:\", recall, \"F1:\", f1_score)\n",
    "        \n",
    "        return\n",
    "      \n",
    "# count the true pos, false pos and false negatives so we can calculate precision and recall\n",
    "def calc_tps(y_true, y_pred):\n",
    "    true_confs = np.round(y_true[...,0])\n",
    "    pred_confs = np.round(y_pred[...,0])\n",
    "    \n",
    "    tps = np.sum((pred_confs == 1) & (true_confs == 1))\n",
    "    fps = np.sum((pred_confs == 1) & (true_confs != 1))\n",
    "    fns = np.sum((pred_confs != 1) & (true_confs == 1))\n",
    "    \n",
    "    return tps, fps, fns      "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 4250
    },
    "colab_type": "code",
    "id": "DCgIz8CwcxUO",
    "outputId": "cae72837-c86c-49de-e20a-c8bda7799c45"
   },
   "outputs": [],
   "source": [
    "GRID_SIZE = 5\n",
    "CELL_SIZE = 1024 / GRID_SIZE\n",
    "\n",
    "# calculate f1 score with smoothing to avoid NaNs\n",
    "def f1_score(y_true, y_pred):\n",
    "    # apply sigmoid and round our labels\n",
    "    y_pred = tf.round(y_pred[...,0])\n",
    "    y_true = tf.round(y_true[...,0])\n",
    "\n",
    "    tps = tf.reduce_sum(tf.cast(tf.logical_and(tf.equal(y_pred, 1), tf.equal(y_true, 1)), dtype=tf.int16)) + 1\n",
    "    fps = tf.reduce_sum(tf.cast(tf.logical_and(tf.equal(y_pred, 1), tf.equal(y_true, 0)), dtype=tf.int16)) + 1\n",
    "    fns = tf.reduce_sum(tf.cast(tf.logical_and(tf.equal(y_pred, 0), tf.equal(y_true, 1)), dtype=tf.int16)) + 1\n",
    "\n",
    "    precision = tps / (tps + fps)\n",
    "    recall = tps / (tps + fns)\n",
    "\n",
    "    f1_score = 2 * (precision * recall) / (precision + recall)\n",
    "\n",
    "    return f1_score\n",
    "\n",
    "def binary_accuracy(y_true, y_pred):\n",
    "    # apply sigmoid to our predictions\n",
    "#     y_pred = tf.sigmoid(y_pred)\n",
    "    # round both since our negative truths are 1e-16 instead of 0\n",
    "    y_true = tf.round(tf.reshape(y_true, [-1, 5]))\n",
    "    y_pred = tf.round(tf.reshape(y_pred, [-1, 5]))\n",
    "    \n",
    "    acc = tf.reduce_mean(tf.cast(tf.equal(y_true[:,0], y_pred[:,0]), dtype=tf.float32))\n",
    "    return acc\n",
    "\n",
    "def overlap_iou(y_true, y_pred):\n",
    "    \"\"\"\n",
    "    Args:\n",
    "        bboxes1: shape (batch_size, 16, 16, 4)\n",
    "            with x1, y1, x2, y2 point order.\n",
    "        bboxes2: shape (batch_size, 16, 16, 4)\n",
    "            with x1, y1, x2, y2 point order.\n",
    "        p1 *-----\n",
    "           |     |\n",
    "           |_____* p2\n",
    "    Returns:\n",
    "        Tensor with shape (total_bboxes1, total_bboxes2)\n",
    "        with the IoU (intersection over union) of bboxes1[i] and bboxes2[j]\n",
    "        in [i, j].\n",
    "    \"\"\"\n",
    "#     y_pred = tf.sigmoid(y_pred)\n",
    "    \n",
    "    # flatten the data because it's easier that way\n",
    "    bboxes1 = tf.reshape(y_true, (-1, 5))\n",
    "    bboxes2 = tf.reshape(y_pred, (-1, 5))\n",
    "    \n",
    "    # split the components out\n",
    "    true_boxes, x11, y11, w1, h1 = tf.split(bboxes1, 5, axis=1)\n",
    "    pred_conf, x21, y21, w2, h2 = tf.split(bboxes2, 5, axis=1)\n",
    "    \n",
    "    # uncenter the data - make sure the numbers are positive\n",
    "    x11 = x11 * CELL_SIZE\n",
    "    x21 = x21 * CELL_SIZE\n",
    "    y11 = y11 * CELL_SIZE\n",
    "    y21 = y21 * CELL_SIZE\n",
    "    \n",
    "    w1 = w1 * 1024\n",
    "    w2 = w2 * 1024\n",
    "    h1 = h1 * 1024\n",
    "    h2 = h2 * 1024\n",
    "    \n",
    "    # is there either a box predicted here or a box actually here?\n",
    "    mask = (pred_conf >= 0.5) | (tf.round(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 + 1e-6), 0)\n",
    "    \n",
    "    # apply the mask\n",
    "    iou = tf.boolean_mask(iou, mask)\n",
    "    \n",
    "    # reduce the mean so we have mean iou for our inputs\n",
    "    return tf.reduce_mean(iou)\n",
    "\n",
    "def overlap_iou2(y_true, y_pred):\n",
    "    \"\"\"\n",
    "    Args:\n",
    "        y_true and y_pred - arrays of boxes containing center points, h and w of boxes\n",
    "        p1 *-----\n",
    "           |     |\n",
    "           |_____* p2\n",
    "    Returns:\n",
    "        Average IOU over boxes\n",
    "    \"\"\"\n",
    "    OVERLAP = 0.3\n",
    "    \n",
    "    # apply the sigmoid\n",
    "#     y_pred = sigmoid(y_pred)\n",
    "    \n",
    "    ious = []\n",
    "    \n",
    "    # loop over both sets of boxes\n",
    "    for truth, pred in zip(y_true, y_pred):\n",
    "        # get true area\n",
    "        true_area = 0\n",
    "        true_boxes = []\n",
    "        for i in range(GRID_SIZE):\n",
    "            for j in range(GRID_SIZE):\n",
    "                c, x, y, w, h = truth[i,j,:]\n",
    "                # if we have an ROI\n",
    "                if c == 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",
    "        \n",
    "        # non-max suppression?\n",
    "        pred_boxes = non_max_suppression_fast(np.array(pred_boxes), OVERLAP)\n",
    "        \n",
    "        # get the intersection\n",
    "        for pred_box in pred_boxes:\n",
    "            x1_p, y1_p, x2_p, y2_p = pred_box\n",
    "            \n",
    "            for true_box in true_boxes:\n",
    "                x1_t, y1_t, x2_t, y2_t = true_box\n",
    "                \n",
    "                # if the boxes overlap at all\n",
    "                if (x1_p >= x1_t and y1_p >= y1_t) or (x1_p <= x1_t and y1_p <= y1_t):\n",
    "                    # get the intersection corners\n",
    "                    x1_i = np.maximum(x1_p, x1_t)\n",
    "                    y1_i = np.maximum(y1_p, y1_t)\n",
    "                    x2_i = np.minimum(x2_p, x2_t)\n",
    "                    y2_i = np.minimum(y2_p, y2_t)\n",
    "                    \n",
    "                    # get area of intersect\n",
    "                    i_w, i_h = x2_i - x1_i, y2_i - y1_i\n",
    "\n",
    "                    # trap for negative numbers\n",
    "                    i_w = np.maximum(i_w, 0)\n",
    "                    i_h = np.maximum(i_h, 0)\n",
    "\n",
    "                    intersection = i_w * i_h\n",
    "                    intersect_area += intersection\n",
    "\n",
    "        union = true_area + pred_area - intersect_area\n",
    "        \n",
    "        iou = intersect_area / (union + 1e-16)\n",
    "        \n",
    "        # only count the IOU if there are truths or predictions\n",
    "        if len(true_boxes) or len(pred_boxes):\n",
    "            ious.append(iou)\n",
    "    \n",
    "    iou_area = np.mean(ious)\n",
    "    return iou_area\n",
    "                    \n",
    "def loss_fn(y_true, y_pred):\n",
    "    # get the iou loss\n",
    "#     iou_loss = iou_loss_fn(y_true, y_pred)\n",
    "    \n",
    "    # get the xe loss\n",
    "    xe_loss = binary_cross_entropy(y_true, y_pred)\n",
    "    \n",
    "    # get the box loss\n",
    "    box_loss = adj_mse(y_true, y_pred)\n",
    "    \n",
    "    # add the losses and return them\n",
    "    return (1.0 * xe_loss) + (box_loss * 2.5) #+ (iou_loss * 0.5)\n",
    "\n",
    "# only apply mse to layers with high confidence that there is an ROI or if there actually is an ROI\n",
    "def adj_mse(y_true, y_pred):\n",
    "    # coefficients\n",
    "    lam_coord = 5\n",
    "    lam_noobj = 0.5\n",
    "    \n",
    "#     y_pred = tf.sigmoid(y_pred)\n",
    "        \n",
    "    # flatten the inputs\n",
    "    y_true = tf.reshape(y_true, (-1, 5))\n",
    "    y_pred = tf.reshape(y_pred, (-1, 5))\n",
    "\n",
    "    # separate the confidence from the boxes\n",
    "    conf_true, x_true, y_true, w_true, h_true = tf.split(y_true, 5, axis=1)\n",
    "    conf_pred, x_pred, y_pred, w_pred, h_pred = tf.split(y_pred, 5, axis=1)\n",
    "\n",
    "    # center squared error\n",
    "    center_se = tf.square(x_true - y_pred) + tf.square(y_true - y_pred)\n",
    "    size_se =  tf.square(tf.sqrt(w_true) - tf.sqrt(w_pred)) + tf.square(tf.sqrt(h_true) - tf.sqrt(h_pred))\n",
    "    \n",
    "    box_se = center_se + size_se\n",
    "    \n",
    "    # only get loss for boxes which are actually positive\n",
    "    mask = tf.greater(conf_true, 0.5)\n",
    "    box_se = tf.boolean_mask(box_se, mask)\n",
    "    \n",
    "    # weight the loss higher\n",
    "    box_se = tf.multiply(box_se, 5.0)\n",
    "    \n",
    "    loss = tf.reduce_sum(box_se)\n",
    "    \n",
    "    return loss\n",
    "    \n",
    "# use weight of 0.5 for negative cells, 19 for positive ones\n",
    "def binary_cross_entropy(y_true, y_pred):\n",
    "    conf_true = tf.round(y_true[...,0])\n",
    "    conf_pred = y_pred[...,0]\n",
    "    \n",
    "    weight = 8.5\n",
    "    weights = tf.multiply(conf_true, weight) + 0.5\n",
    "    \n",
    "    xe = tf.multiply(tf.square(conf_true - conf_pred), weights)\n",
    "    \n",
    "    return tf.reduce_sum(xe)\n",
    "\n",
    "def objectness_loss_fn(y_true, y_pred):\n",
    "#     y_pred = tf.sigmoid(y_pred)\n",
    "    \n",
    "    # pred box conf\n",
    "    pred_box_conf = y_pred[...,0]\n",
    "    \n",
    "    # separate the x, y from the w, h\n",
    "    true_box_xy = y_true[...,1:3] * CELL_SIZE\n",
    "    true_box_wh = y_true[...,3:] * 1024\n",
    "    \n",
    "    pred_box_xy = y_pred[...,1:3] * CELL_SIZE\n",
    "    pred_box_wh = y_pred[...,3:] * 1024\n",
    "    \n",
    "    # get the corners of the boxes by subtracting or adding half of h, w\n",
    "    true_wh_half = true_box_wh / 2.\n",
    "    true_mins    = true_box_xy - true_wh_half\n",
    "    true_maxes   = true_box_xy + true_wh_half\n",
    "    \n",
    "    pred_wh_half = pred_box_wh / 2.\n",
    "    pred_mins    = pred_box_xy - pred_wh_half\n",
    "    pred_maxes   = pred_box_xy + pred_wh_half  \n",
    "    \n",
    "    # get the corners of the intersect area\n",
    "    intersect_mins  = tf.maximum(pred_mins,  true_mins)\n",
    "    intersect_maxes = tf.minimum(pred_maxes, true_maxes)\n",
    "    intersect_wh    = tf.maximum(intersect_maxes - intersect_mins, 0.)\n",
    "    intersect_areas = intersect_wh[..., 0] * intersect_wh[..., 1]\n",
    "    \n",
    "    true_areas = true_box_wh[..., 0] * true_box_wh[..., 1]\n",
    "    pred_areas = pred_box_wh[..., 0] * pred_box_wh[..., 1]\n",
    "\n",
    "    union_areas = pred_areas + true_areas - intersect_areas\n",
    "    iou_scores  = tf.truediv(intersect_areas, union_areas + 1e-16)\n",
    "    \n",
    "    true_box_conf = iou_scores * y_true[..., 0]\n",
    "    \n",
    "    conf_mask  = tf.zeros_like(iou_scores)\n",
    "    conf_mask = conf_mask + tf.cast((iou_scores < 0.6), dtype=tf.float32) * (1 - y_true[...,0])\n",
    "    conf_mask = conf_mask + y_true[..., 0] * 5.0\n",
    "    \n",
    "    nb_conf_box  = tf.reduce_sum(tf.cast((conf_mask > 0.0), dtype=tf.float32)) \n",
    "    \n",
    "    loss_conf  =  tf.reduce_sum(tf.square(true_box_conf-pred_box_conf) * conf_mask)  / (nb_conf_box + 1e-6) / 2.\n",
    "    \n",
    "    return tf.reduce_sum(loss_conf)\n",
    "\n",
    "def iou_loss_fn(y_true, y_pred):\n",
    "    # apply sigmoid to predictions\n",
    "#     y_pred = tf.sigmoid(y_pred)\n",
    "    \n",
    "    # separate the x, y from the w, h and unnormalize them\n",
    "    true_box_xy = y_true[...,1:3] * CELL_SIZE\n",
    "    true_box_wh = y_true[...,3:] * 1024\n",
    "    \n",
    "    pred_box_xy = y_pred[...,1:3] * CELL_SIZE\n",
    "    pred_box_wh = y_pred[...,3:] * 1024\n",
    "    \n",
    "    # get the corners of the boxes by subtracting or adding half of h, w\n",
    "    true_wh_half = true_box_wh / 2.\n",
    "    true_mins    = true_box_xy - true_wh_half\n",
    "    true_maxes   = true_box_xy + true_wh_half\n",
    "    \n",
    "    pred_wh_half = pred_box_wh / 2.\n",
    "    pred_mins    = pred_box_xy - pred_wh_half\n",
    "    pred_maxes   = pred_box_xy + pred_wh_half  \n",
    "    \n",
    "    # get the corners of the intersect area\n",
    "    intersect_mins  = tf.maximum(pred_mins,  true_mins)\n",
    "    intersect_maxes = tf.minimum(pred_maxes, true_maxes)\n",
    "    intersect_wh    = tf.maximum(intersect_maxes - intersect_mins, 0.)\n",
    "    intersect_areas = intersect_wh[..., 0] * intersect_wh[..., 1]\n",
    "    \n",
    "    # get the area of the boxes\n",
    "    true_areas = true_box_wh[..., 0] * true_box_wh[..., 1]\n",
    "    pred_areas = pred_box_wh[..., 0] * pred_box_wh[..., 1]\n",
    "\n",
    "    # calculate the IOU\n",
    "    union_areas = pred_areas + true_areas - intersect_areas\n",
    "    # multiple the IOU by the confidence so we have a more useful loss function\n",
    "    iou_scores  = tf.truediv(intersect_areas, union_areas + 1e-16) * y_pred[...,0]\n",
    "    \n",
    "    # only use the IOU from boxes which actually have ROIs\n",
    "    mask = tf.greater(y_true[...,0], 0.5)\n",
    "    # subtract IOU from 1 and apply the mask\n",
    "    use_iou = (1 - iou_scores)\n",
    "    use_iou = tf.boolean_mask(use_iou, mask)\n",
    "    \n",
    "    return tf.reduce_sum(use_iou)\n",
    "    \n",
    "# create network and compiler\n",
    "model = create_network(input_size=IMAGE_SIZE, channels=24, n_blocks=1, depth=4)\n",
    "\n",
    "optimizer = keras.optimizers.Adam(lr=0.001, beta_1=0.9, beta_2=0.999, decay=1e-7, clipnorm=2.0, clipvalue=0.75)\n",
    "\n",
    "model.compile(optimizer=optimizer,\n",
    "              loss=loss_fn,\n",
    "              metrics=[binary_accuracy, adj_mse, binary_cross_entropy, overlap_iou, f1_score])\n",
    "\n",
    "# cosine learning rate annealing\n",
    "def exp_decay(x):\n",
    "    lr0 = 0.005\n",
    "    epochs_drop = 15\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",
    "calc_iou = Calc_IOU_CB()\n",
    "\n",
    "# create train and validation generators\n",
    "folder = './stage_1_train_images'\n",
    "train_gen = generator(folder, train_filenames, pneumonia_locations, batch_size=BATCH_SIZE, image_size=IMAGE_SIZE, shuffle=True, augment=True, predict=False)\n",
    "valid_gen = generator(folder, valid_filenames, pneumonia_locations, batch_size=BATCH_SIZE, image_size=IMAGE_SIZE, shuffle=False, predict=False)\n",
    "\n",
    "print(model.summary())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "jSSSpmNE-L1L"
   },
   "outputs": [],
   "source": [
    "model.load_weights(CHECKPOINT_PATH)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 445
    },
    "colab_type": "code",
    "id": "fM5VeiuDcxUY",
    "outputId": "20ad5ff1-7019-4717-cfc9-56d86a1af948"
   },
   "outputs": [],
   "source": [
    "history = model.fit_generator(train_gen, validation_data=valid_gen, callbacks=[learning_rate, checkpoint, calc_iou], epochs=20, shuffle=True, verbose=1, initial_epoch=6)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "7tyG7P79Y42U"
   },
   "outputs": [],
   "source": [
    "history = model.fit_generator(train_gen, validation_data=valid_gen, callbacks=[learning_rate, checkpoint, calc_iou], epochs=30, shuffle=True, verbose=1, initial_epoch=20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "y_bLpJPBcxUh"
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(20,5))\n",
    "plt.subplot(141)\n",
    "plt.plot(history.epoch, history.history[\"loss\"], label=\"Train loss\")\n",
    "plt.plot(history.epoch, history.history[\"val_loss\"], label=\"Valid loss\")\n",
    "plt.legend()\n",
    "plt.subplot(142)\n",
    "plt.plot(history.epoch, history.history[\"adj_mse\"], label=\"Train MSE loss\")\n",
    "plt.plot(history.epoch, history.history[\"val_adj_mse\"], label=\"Valid MSE loss\")\n",
    "plt.legend()\n",
    "plt.subplot(143)\n",
    "plt.plot(history.epoch, history.history[\"binary_accuracy\"], label=\"Train Confidence accuracy\")\n",
    "plt.plot(history.epoch, history.history[\"val_binary_accuracy\"], label=\"Valid Confidence accuracy\")\n",
    "plt.legend()\n",
    "plt.subplot(144)\n",
    "plt.plot(history.epoch, history.history[\"overlap_iou\"], label=\"Train iou\")\n",
    "plt.plot(history.epoch, history.history[\"val_overlap_iou\"], label=\"Valid iou\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "NIdDp8YOxevU"
   },
   "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": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 8538
    },
    "colab_type": "code",
    "id": "A03zhJyfUyV2",
    "outputId": "55d86a38-8c94-4e3d-9e50-909171f71de4"
   },
   "outputs": [],
   "source": [
    "THRESHOLD = 0.75\n",
    "OVERLAP = 0.3\n",
    "# look at some sample predictions\n",
    "samples = np.random.choice(valid_filenames, size=20, replace=False)\n",
    "# samples = problem_names\n",
    "\n",
    "coords = np.arange(0, 1024, CELL_SIZE)\n",
    "overall_ious = []\n",
    "map_ious = []\n",
    "\n",
    "for filename in samples:\n",
    "    # load the image\n",
    "    img = pydicom.dcmread(os.path.join(train_dicom_dir, filename)).pixel_array\n",
    "    \n",
    "    filename = filename.split('.')[0]\n",
    "    \n",
    "    f, ax = plt.subplots(figsize=(6, 6))\n",
    "    plt.imshow(img)\n",
    "    \n",
    "    # initialize our lists\n",
    "    ious = []\n",
    "    truths = []\n",
    "    boxes = []\n",
    "    \n",
    "    # draw the truth boxes\n",
    "    if filename in pneumonia_locations:\n",
    "        locs = pneumonia_locations[filename].copy()\n",
    "        for i in range(GRID_SIZE):\n",
    "            for j in range(GRID_SIZE):\n",
    "                pixel_data = locs[i,j,:]\n",
    "                if pixel_data[0] > 0.5:\n",
    "                    x, y, w, h = pixel_data[1:]\n",
    "                    \n",
    "                    # unnormalize the data\n",
    "                    w = w * 1024\n",
    "                    h = h * 1024\n",
    "                    \n",
    "                    x = unnorm(x, j, CELL_SIZE)\n",
    "                    y = unnorm(y, i, CELL_SIZE)\n",
    "                    \n",
    "                    # get the corners\n",
    "                    x = x - (w // 2)\n",
    "                    y = y - (h // 2)\n",
    "                    \n",
    "                    x = int(x)\n",
    "                    y = int(y)\n",
    "                    w = int(w)\n",
    "                    h = int(h)\n",
    "                    locs[i,j,:] = [1, x, y, w, h]\n",
    "                    print(\"Truth:\", i, j, x, y, w, h)\n",
    "                    truths.append([x, y, w, h])\n",
    "                    \n",
    "                    rect = patches.Rectangle((x,y),w,h,linewidth=1,edgecolor='b',facecolor='none')\n",
    "                    ax.add_patch(rect)\n",
    "                    \n",
    "    # predict the image\n",
    "    img = resize(img, (IMAGE_SIZE, IMAGE_SIZE), mode='reflect')\n",
    "    yhat = model.predict(img.reshape(1, IMAGE_SIZE, IMAGE_SIZE, 1))\n",
    "#     yhat = sigmoid(yhat)\n",
    "    conf = np.squeeze(yhat)[:,:,0]\n",
    "    bboxes = np.squeeze(yhat)[:,:,1:]    \n",
    "    pred_boxes = np.zeros_like(bboxes)\n",
    "\n",
    "    # loop through our predictions\n",
    "    for i in range(GRID_SIZE):\n",
    "        for j in range(GRID_SIZE):\n",
    "            conf_ = conf[i,j]\n",
    "            # if we have a prediction\n",
    "            if conf_ > THRESHOLD:\n",
    "                x,y,w,h = bboxes[i,j,:]\n",
    "                \n",
    "                # unnormalize the data\n",
    "                w = w * 1024\n",
    "                h = h * 1024\n",
    "\n",
    "                x = unnorm(x, j, CELL_SIZE)\n",
    "                y = unnorm(y, i, CELL_SIZE)\n",
    "                \n",
    "                # convert to upper left corner from center\n",
    "                x = np.maximum(x - (w // 2), 0)\n",
    "                y = np.maximum(y - (h // 2), 0)\n",
    "                \n",
    "                x = int(x)\n",
    "                y = int(y)\n",
    "                w = int(w)\n",
    "                h = int(h)\n",
    "                \n",
    "                pred_boxes[i,j,:] = [x, y, w, h]\n",
    "                \n",
    "                print(\"Pred:\", i, j, \"conf:\", conf_, x, y, w, h)\n",
    "                \n",
    "                # if the boxes have width and height add them to our list\n",
    "                if w > 30 and h > 30:\n",
    "                    boxes.append([x,y,w,h, conf_])\n",
    "    \n",
    "    ## if we have ground truths OR predictions calculate the IOUs:\n",
    "    if len(boxes) or len(truths):\n",
    "        # do non-max suppression of our boxes\n",
    "        nms_boxes = non_max_suppression_fast(np.array(boxes), OVERLAP)\n",
    "\n",
    "        # only count the IOU if there are either ground truths or predictions\n",
    "        if len(nms_boxes) or len(truths):\n",
    "            iou = calc_iou_better(np.array(truths), np.array(nms_boxes))\n",
    "            print(\"IOU post-nms:\", iou)\n",
    "            overall_ious.append(iou)\n",
    "\n",
    "        # plot our boxes\n",
    "        for box in nms_boxes:\n",
    "            x,y,w,h,c = box\n",
    "            rect = patches.Rectangle((x,y), w, h, linewidth=1, edgecolor='r', facecolor='none')\n",
    "            ax.add_patch(rect)\n",
    "\n",
    "        if len(nms_boxes):\n",
    "            scores = nms_boxes[...,4]\n",
    "        else:\n",
    "            scores = np.array([])\n",
    "\n",
    "        map_iou_score = map_iou(np.array(truths), np.array(nms_boxes), scores)\n",
    "\n",
    "        if map_iou_score is not None:\n",
    "            map_ious.append(map_iou_score)\n",
    "\n",
    "        for item in coords:\n",
    "            plt.axvline(item, linewidth=0.5)\n",
    "            plt.axhline(item, linewidth=0.5)\n",
    "\n",
    "        print(\"IOU:\", iou)\n",
    "        print(\"Map IOU:\", map_iou_score)\n",
    "    \n",
    "    plt.title(filename)\n",
    "    plt.show()\n",
    "    \n",
    "print(\"Overall Mean IOU:\", np.mean(overall_ious))    \n",
    "print(\"Overall Map IOU:\", np.mean(map_ious))    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "3tAntxescxUq"
   },
   "source": [
    "# Predict test images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "fnu7ebfScxUu"
   },
   "outputs": [],
   "source": [
    "THRESHOLD = 0.75\n",
    "\n",
    "# load and shuffle filenames\n",
    "folder = './stage_1_test_images'\n",
    "test_filenames = os.listdir(folder)\n",
    "print('n test samples:', len(test_filenames))\n",
    "\n",
    "# create test generator with predict flag set to True\n",
    "test_gen = generator(folder, test_filenames, None, batch_size=24, image_size=IMAGE_SIZE, shuffle=False, predict=True)\n",
    "\n",
    "# create submission dictionary\n",
    "submission_dict = {}\n",
    "# loop through testset\n",
    "for imgs, filenames in test_gen:\n",
    "    \n",
    "    # predict batch of images\n",
    "    yhats = model.predict(imgs)\n",
    "    \n",
    "    # apply sigmoid\n",
    "#     yhats = sigmoid(yhats)\n",
    "    \n",
    "    # loop through batch\n",
    "    for yhat, filename in zip(yhats, filenames):\n",
    "        predictionString = \"\"\n",
    "        boxes = []\n",
    "        for i in range(GRID_SIZE):\n",
    "            for j in range(GRID_SIZE):\n",
    "                conf = yhat[i, j, 0]\n",
    "                if conf > THRESHOLD:\n",
    "                    x, y, w, h = yhat[i,j, 1:]\n",
    "                    \n",
    "                    # possible thresholds to keep our boxes within reasonable sizes?\n",
    "                    if True: #w < 600 and h < 1000:\n",
    "                        w = w * 1024\n",
    "                        h = h * 1024\n",
    "\n",
    "                        x = unnorm(x, j, CELL_SIZE)\n",
    "                        y = unnorm(y, i, CELL_SIZE)\n",
    "\n",
    "                        # convert to upper left corner from center\n",
    "                        x = x - (w // 2)\n",
    "                        y = y - (h // 2)\n",
    "                        \n",
    "                        if w > 20 and h > 20:\n",
    "                            # make sure our boxes don't run off the edges of the images\n",
    "                            w = np.minimum(w, 1024 - x)\n",
    "                            h = np.minimum(h, 1024 - y)\n",
    "                            boxes.append([x,y,w,h])\n",
    "\n",
    "        # do our non-max suppression here\n",
    "        boxes = non_max_suppression_fast(np.array(boxes), 0.3)\n",
    "        \n",
    "        # loop through our suppressed boxes and creat the prediction string\n",
    "        for box in boxes:\n",
    "            x,y,w,h = box\n",
    "            \n",
    "            x = int(x)\n",
    "            y = int(y)\n",
    "            w = int(w)\n",
    "            h = int(h)\n",
    "        \n",
    "            # create the prediction string\n",
    "            predictionString += str(0.9) + ' ' + str(x) + ' ' + str(y) + ' ' + str(w) + ' ' + str(h) + ' '\n",
    "            \n",
    "        # add filename and predictionString to dictionary\n",
    "        filename = filename.split('.')[0]\n",
    "        submission_dict[filename] = predictionString\n",
    "\n",
    "    # stop if we've got them all\n",
    "    if len(submission_dict) >= len(test_filenames):\n",
    "        break\n",
    "    \n",
    "# save dictionary as csv file\n",
    "sub = pd.DataFrame.from_dict(submission_dict,orient='index')\n",
    "sub.index.names = ['patientId']\n",
    "sub.columns = ['PredictionString']\n",
    "\n",
    "now = datetime.datetime.now()\n",
    "today = str(now)[:10]\n",
    "submission_file = today + \"_yolo_submission.csv\" \n",
    "sub.to_csv(submission_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "U1NEg9QCcxU4"
   },
   "outputs": [],
   "source": [
    "!kaggle competitions submit -c rsna-pneumonia-detection-challenge -f {submission_file} -m \"YOLOv19.1 480x480 20 epochs\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 85
    },
    "colab_type": "code",
    "id": "yf06y04TVRi2",
    "outputId": "892bb3a4-e84b-44d3-b15a-8065950e3be8"
   },
   "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": "yolo_480x480_18.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.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}