Colab_notebooks/RetinaNet_ZeroCostDL4Mic.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "IkSguVy8Xv83"
   },
   "source": [
    "# **RetinaNet**\n",
    "---\n",
    "\n",
    "<font size = 4> RetinaNet is a deep-learning network designed to perform object detection and classification of objects in images, published by [Tsung-Yi Lin, Priya Goyal, Ross Girshick, Kaiming He, Piotr Dollár](https://arxiv.org/abs/1708.02002). \n",
    "It is trained on images with class annotations in the form of bounding boxes drawn around the objects of interest. RetinaNet is a single, unified network composed of a backbone network and two task-specific subnetworks. The backbone is responsible for computing a convolutional feature map over an entire input image and is an off-the-self convolution network. The first subnet performs classification on the backbones output; the second subnet performs convolution bounding box regression. Aplying this mechanic, RetinaNet learns classification and object detection simultaneously by taking the whole input image into account, predicting many possible bounding box solutions, and then using regression to find the best bounding boxes and classifications for each object.\n",
    "\n",
    "\n",
    "<img width=\"40%\" alt =\"Example of image detection with retinanet.\" src=\"http://findmycells.org/uploads/3/4/9/3/34939463/fmcback4_orig.png\"><figcaption> Example of image detection with RetinaNet. </figcaption>\n",
    "\n",
    "<font size = 4>**This particular notebook enables object detection and classification on 2D images given ground truth bounding boxes. If you are interested in image segmentation, you should use our U-Net or Stardist notebooks instead.**\n",
    "\n",
    "---\n",
    "<font size = 4>*Disclaimer*:\n",
    "\n",
    "<font size = 4>This notebook is part of the Zero-Cost Deep-Learning to Enhance Microscopy project (https://github.com/HenriquesLab/DeepLearning_Collab/wiki) jointly developed by [Erlantz Calvo](https://github.com/ErlantzCalvo) and [Ignacio Arganda-Carreras](https://sites.google.com/site/iargandacarreras/).\n",
    "\n",
    " <font size = 4>This notebook is based on the following paper: \n",
    "\n",
    "**Feature Pyramid Networks for Object Detection** from  Yin Cui, Menglin Jia, Tsung-Yi Lin, Yang Song, Serge Belongie in IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2019, (https://openaccess.thecvf.com/content_CVPR_2019/html/Cui_Class-Balanced_Loss_Based_on_Effective_Number_of_Samples_CVPR_2019_paper.html) \n",
    "\n",
    "<font size = 4>**Note: The source code for this notebook is adapted for Tensorflow 2 and can be found in: (https://github.com/tensorflow/models)**\n",
    "\n",
    "\n",
    "<font size = 4>**Please also cite these original papers when using or developing this notebook.**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "jWAz2i7RdxUV"
   },
   "source": [
    "# **How to use this notebook?**\n",
    "\n",
    "---\n",
    "\n",
    "<font size = 4>Video describing how to use our notebooks are available on youtube:\n",
    "  - [**Video 1**](https://www.youtube.com/watch?v=GzD2gamVNHI&feature=youtu.be): Full run through of the workflow to obtain the notebooks and the provided test datasets as well as a common use of the notebook\n",
    "  - [**Video 2**](https://www.youtube.com/watch?v=PUuQfP5SsqM&feature=youtu.be): Detailed description of the different sections of the notebook\n",
    "\n",
    "\n",
    "---\n",
    "\n",
    "\n",
    "### **Structure of a notebook**\n",
    "\n",
    "<font size = 4>The notebook contains two types of cells:  \n",
    "\n",
    "<font size = 4>**Text cells** provide information and can be modified by douple-clicking the cell. You are currently reading the text cell. You can create a new text by clicking `+ Text`.\n",
    "\n",
    "<font size = 4>**Code cells** contain code and the code can be modfied by selecting the cell. To execute the cell, move your cursor on the `[ ]`-mark on the left side of the cell (play button appears). Click to execute the cell. After execution is done the animation of play button stops. You can create a new coding cell by clicking `+ Code`.\n",
    "\n",
    "---\n",
    "### **Table of contents, Code snippets** and **Files**\n",
    "\n",
    "<font size = 4>On the top left side of the notebook you find three tabs which contain from top to bottom:\n",
    "\n",
    "<font size = 4>*Table of contents* = contains structure of the notebook. Click the content to move quickly between sections.\n",
    "\n",
    "<font size = 4>*Code snippets* = contain examples how to code certain tasks. You can ignore this when using this notebook.\n",
    "\n",
    "<font size = 4>*Files* = contain all available files. After mounting your google drive (see section 1.) you will find your files and folders here. \n",
    "\n",
    "<font size = 4>**Remember that all uploaded files are purged after changing the runtime.** All files saved in Google Drive will remain. You do not need to use the Mount Drive-button; your Google Drive is connected in section 1.2.\n",
    "\n",
    "<font size = 4>**Note:** The \"sample data\" in \"Files\" contains default files. Do not upload anything in here!\n",
    "\n",
    "---\n",
    "### **Making changes to the notebook**\n",
    "\n",
    "<font size = 4>**You can make a copy** of the notebook and save it to your Google Drive. To do this click file -> save a copy in drive.\n",
    "\n",
    "<font size = 4>To **edit a cell**, double click on the text. This will show you either the source code (in code cells) or the source text (in text cells).\n",
    "You can use the `#`-mark in code cells to comment out parts of the code. This allows you to keep the original code piece in the cell as a comment."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "gKDLkLWUd-YX"
   },
   "source": [
    "# **0. Before getting started**\n",
    "---\n",
    "<font size = 4> Preparing the dataset carefully is essential to make this RetinaNet notebook work. This model requires as input a set of images and as target a list of annotation files in Pascal VOC format. The annotation files should have the exact same name as the input files, except with an .xml instead of the .jpg extension. The annotation files contain the class labels and all bounding boxes for the objects for each image in your dataset. Most datasets will give the option of saving the annotations in this format or using software for hand-annotations will automatically save the annotations in this format. \n",
    "\n",
    "<font size=4> If you want to assemble your own dataset we recommend using the open source https://www.makesense.ai/ resource. You can follow our instructions on how to label your dataset with this tool on our [wiki](https://github.com/HenriquesLab/ZeroCostDL4Mic/wiki/Object-Detection-(YOLOv2)).\n",
    "\n",
    "<font size = 4>**We strongly recommend that you generate extra paired images. These images can be used to assess the quality of your trained model (Quality control dataset)**. The quality control assessment can be done directly in this notebook.\n",
    "\n",
    "<font size = 4> **Additionally, the corresponding input and output files need to have the same name**.\n",
    "\n",
    "<font size = 4> Please note that you currently can **only use .png or .jpg files!**\n",
    "\n",
    "\n",
    "<font size = 4>Here's a common data structure that can work:\n",
    "*   Experiment A\n",
    "    - **Training dataset**\n",
    "      - Input images (Training_source)\n",
    "        - img_1.png, img_2.png, ...\n",
    "      - Input annotationsAnnotationsg_1.xml, img_2.xml, ...\n",
    "    - **Quality control dataset**\n",
    "     - Input images\n",
    "        - img_1.png, img_2.png\n",
    "      - Input annotations\n",
    "        - img_1.xml, img_2.xml\n",
    "    - **Data to be predicted**\n",
    "    - **Results**\n",
    "\n",
    "---\n",
    "<font size = 4>**Important note**\n",
    "\n",
    "<font size = 4>- If you wish to **Train a network from scratch** using your own dataset (and we encourage everyone to do that), you will need to run **sections 1 - 4**, then use **section 5** to assess the quality of your model and **section 6** to run predictions using the model that you trained.\n",
    "\n",
    "<font size = 4>- If you wish to **Evaluate your model** using a model previously generated and saved on your Google Drive, you will only need to run **sections 1 and 2** to set up the notebook, then use **section 5** to assess the quality of your model.\n",
    "\n",
    "<font size = 4>- If you only wish to **run predictions** using a model previously generated and saved on your Google Drive, you will only need to run **sections 1 and 2** to set up the notebook, then use **section 6** to run the predictions on the desired model.\n",
    "---\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "AdN8B91xZO0x"
   },
   "source": [
    "# **1. Install RetinaNet and Dependencies**\n",
    "---\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "fq21zJVFNASx"
   },
   "outputs": [],
   "source": [
    "Notebook_version = '1.14.1'\n",
    "Network = 'RetinaNet'\n",
    "\n",
    "#@markdown ##Install Network and Dependencies\n",
    "!pip install pascal-voc-writer\n",
    "!pip install fpdf2\n",
    "!pip install PTable\n",
    "!pip install colorama\n",
    "\n",
    "# this line exist because in june of 2021 tensorflow was updated causing internal dependecies error\n",
    "# !pip3 install --upgrade tensorflow --user\n",
    "\n",
    "\n",
    "from pascal_voc_writer import Writer\n",
    "from __future__ import division\n",
    "from __future__ import print_function\n",
    "from __future__ import absolute_import\n",
    "import csv\n",
    "import random\n",
    "import pprint\n",
    "import sys\n",
    "import time\n",
    "import numpy as np\n",
    "from optparse import OptionParser\n",
    "import pickle\n",
    "import math\n",
    "import cv2\n",
    "import copy\n",
    "import math\n",
    "from matplotlib import pyplot as plt\n",
    "import matplotlib.patches as patches\n",
    "import tensorflow as tf\n",
    "import pandas as pd\n",
    "import os\n",
    "import shutil\n",
    "from skimage import io\n",
    "from colorama import Fore, Back, Style\n",
    "from sklearn.metrics import average_precision_score\n",
    "\n",
    "#Create a variable to get and store relative base path\n",
    "base_path = os.getcwd()\n",
    "\n",
    "from keras import backend as K\n",
    "import xml.etree.ElementTree as ET\n",
    "from collections import OrderedDict, Counter\n",
    "import json\n",
    "import imageio\n",
    "import imgaug as ia\n",
    "from imgaug import augmenters as iaa\n",
    "import copy\n",
    "import cv2\n",
    "from tqdm import tqdm\n",
    "from tempfile import mkstemp\n",
    "from shutil import move, copymode\n",
    "from os import fdopen, remove\n",
    "from fpdf import FPDF, HTMLMixin\n",
    "from datetime import datetime\n",
    "from pip._internal.operations.freeze import freeze\n",
    "import subprocess as sp\n",
    "\n",
    "from prettytable import from_csv\n",
    "\n",
    "%matplotlib inline\n",
    "\n",
    "ia.seed(1)\n",
    "# imgaug uses matplotlib backend for displaying images\n",
    "from imgaug.augmentables.bbs import BoundingBox, BoundingBoxesOnImage\n",
    "import re\n",
    "import glob\n",
    "\n",
    "\n",
    "#Here, we import the main github repo for this notebook and move it to the gdrive\n",
    "\n",
    "\n",
    "!git clone https://github.com/tensorflow/models.git\n",
    "\n",
    "%cd models/research/\n",
    "\n",
    "!protoc object_detection/protos/*.proto --python_out=.\n",
    "\n",
    "%cp object_detection/packages/tf2/setup.py .\n",
    "!python -m pip install .\n",
    "\n",
    "%cd \"${base_path}\"\n",
    "\n",
    "\n",
    "import pathlib\n",
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "\n",
    "from PIL import Image\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from tqdm import tqdm\n",
    "from object_detection.utils import ops as utils_ops\n",
    "from object_detection.utils import label_map_util\n",
    "from object_detection.utils import config_util\n",
    "from object_detection.utils import visualization_utils as vis_util\n",
    "from object_detection.builders import model_builder\n",
    "\n",
    "import imageio\n",
    "from six import BytesIO\n",
    "\n",
    "utils_ops.tf = tf.compat.v1\n",
    "tf.gfile = tf.io.gfile\n",
    "\n",
    "def download_weights(model_name):\n",
    "  download_url = f'http://download.tensorflow.org/models/object_detection/tf2/20200711/{model_name}.tar.gz'\n",
    "  tar_name =f'{model_name}.tar.gz' \n",
    "  !wget -nv --show-progress $download_url\n",
    "  !tar -xf $tar_name\n",
    "  !rm -rf $tar_name\n",
    "\n",
    "  if os.path.exists(os.path.join(model_path,'checkpoint')):\n",
    "    shutil.rmtree(os.path.join(model_path,'checkpoint'))\n",
    "  \n",
    "  checkpoint_current_path = f'{model_name}/checkpoint'\n",
    "  \n",
    "  %mv $checkpoint_current_path $model_path\n",
    "\n",
    "def parse_annotation(ann_dir, img_dir, labels=[]):\n",
    "    all_imgs = []\n",
    "    seen_labels = {}\n",
    "    \n",
    "    for ann in sorted(os.listdir(ann_dir)):\n",
    "        img = {'object':[]}\n",
    "\n",
    "        tree = ET.parse(ann_dir + ann)\n",
    "        \n",
    "        for elem in tree.iter():\n",
    "            if 'filename' in elem.tag:\n",
    "                img['filename'] = img_dir + elem.text\n",
    "            if 'width' in elem.tag:\n",
    "                img['width'] = int(elem.text)\n",
    "            if 'height' in elem.tag:\n",
    "                img['height'] = int(elem.text)\n",
    "            if 'object' in elem.tag or 'part' in elem.tag:\n",
    "                obj = {}\n",
    "                \n",
    "                for attr in list(elem):\n",
    "                    if 'name' in attr.tag:\n",
    "                        obj['name'] = attr.text\n",
    "\n",
    "                        if obj['name'] in seen_labels:\n",
    "                            seen_labels[obj['name']] += 1\n",
    "                        else:\n",
    "                            seen_labels[obj['name']] = 1\n",
    "                        \n",
    "                        if len(labels) > 0 and obj['name'] not in labels:\n",
    "                            break\n",
    "                        else:\n",
    "                            img['object'] += [obj]\n",
    "                            \n",
    "                    if 'bndbox' in attr.tag:\n",
    "                        for dim in list(attr):\n",
    "                            if 'xmin' in dim.tag:\n",
    "                                obj['xmin'] = int(round(float(dim.text)))\n",
    "                            if 'ymin' in dim.tag:\n",
    "                                obj['ymin'] = int(round(float(dim.text)))\n",
    "                            if 'xmax' in dim.tag:\n",
    "                                obj['xmax'] = int(round(float(dim.text)))\n",
    "                            if 'ymax' in dim.tag:\n",
    "                                obj['ymax'] = int(round(float(dim.text)))\n",
    "\n",
    "        if len(img['object']) > 0:\n",
    "            all_imgs += [img]\n",
    "                        \n",
    "    return all_imgs, seen_labels\n",
    "\n",
    "def plt_rectangle(plt,label,x1,y1,x2,y2,fontsize=10):\n",
    "    '''\n",
    "    == Input ==\n",
    "    \n",
    "    plt   : matplotlib.pyplot object\n",
    "    label : string containing the object class name\n",
    "    x1    : top left corner x coordinate\n",
    "    y1    : top left corner y coordinate\n",
    "    x2    : bottom right corner x coordinate\n",
    "    y2    : bottom right corner y coordinate\n",
    "    '''\n",
    "    linewidth = 1\n",
    "    color = \"yellow\"\n",
    "    plt.text(x1,y1,label,fontsize=fontsize,backgroundcolor=\"magenta\")\n",
    "    plt.plot([x1,x1],[y1,y2], linewidth=linewidth,color=color)\n",
    "    plt.plot([x2,x2],[y1,y2], linewidth=linewidth,color=color)\n",
    "    plt.plot([x1,x2],[y1,y1], linewidth=linewidth,color=color)\n",
    "    plt.plot([x1,x2],[y2,y2], linewidth=linewidth,color=color)\n",
    "\n",
    "def extract_single_xml_file(tree,object_count=True):\n",
    "    Nobj = 0\n",
    "    row  = OrderedDict()\n",
    "    for elems in tree.iter():\n",
    "\n",
    "        if elems.tag == \"size\":\n",
    "            for elem in elems:\n",
    "                row[elem.tag] = int(elem.text)\n",
    "        if elems.tag == \"object\":\n",
    "            for elem in elems:\n",
    "                if elem.tag == \"name\":\n",
    "                    row[\"bbx_{}_{}\".format(Nobj,elem.tag)] = str(elem.text)              \n",
    "                if elem.tag == \"bndbox\":\n",
    "                    for k in elem:\n",
    "                        row[\"bbx_{}_{}\".format(Nobj,k.tag)] = float(k.text)\n",
    "                    Nobj += 1\n",
    "    if object_count == True:\n",
    "      row[\"Nobj\"] = Nobj\n",
    "    return(row)\n",
    "\n",
    "def count_objects(tree):\n",
    "  Nobj=0\n",
    "  for elems in tree.iter():\n",
    "    if elems.tag == \"object\":\n",
    "      for elem in elems:\n",
    "        if elem.tag == \"bndbox\":\n",
    "          Nobj += 1\n",
    "  return(Nobj)\n",
    "\n",
    "def compute_overlap(a, b):\n",
    "    \"\"\"\n",
    "    Code originally from https://github.com/rbgirshick/py-faster-rcnn.\n",
    "    Parameters\n",
    "    ----------\n",
    "    a: (N, 4) ndarray of float\n",
    "    b: (K, 4) ndarray of float\n",
    "    Returns\n",
    "    -------\n",
    "    overlaps: (N, K) ndarray of overlap between boxes and query_boxes\n",
    "    \"\"\"\n",
    "    area = (b[:, 2] - b[:, 0]) * (b[:, 3] - b[:, 1])\n",
    "\n",
    "    iw = np.minimum(np.expand_dims(a[:, 2], axis=1), b[:, 2]) - np.maximum(np.expand_dims(a[:, 0], 1), b[:, 0])\n",
    "    ih = np.minimum(np.expand_dims(a[:, 3], axis=1), b[:, 3]) - np.maximum(np.expand_dims(a[:, 1], 1), b[:, 1])\n",
    "\n",
    "    iw = np.maximum(iw, 0)\n",
    "    ih = np.maximum(ih, 0)\n",
    "\n",
    "    ua = np.expand_dims((a[:, 2] - a[:, 0]) * (a[:, 3] - a[:, 1]), axis=1) + area - iw * ih\n",
    "\n",
    "    ua = np.maximum(ua, np.finfo(float).eps)\n",
    "\n",
    "    intersection = iw * ih\n",
    "\n",
    "    return intersection / ua\n",
    "\n",
    "\n",
    "def load_image_into_numpy_array(path):\n",
    "  \"\"\"Load an image from file into a numpy array.\n",
    "\n",
    "  Puts image into numpy array to feed into tensorflow graph.\n",
    "  Note that by convention we put it into a numpy array with shape\n",
    "  (height, width, channels), where channels=3 for RGB.\n",
    "\n",
    "  Args:\n",
    "    path: a file path.\n",
    "\n",
    "  Returns:\n",
    "    uint8 numpy array with shape (img_height, img_width, 3)\n",
    "  \"\"\"\n",
    "  img_data = tf.io.gfile.GFile(path, 'rb').read()\n",
    "  image = Image.open(BytesIO(img_data))\n",
    "  (im_width, im_height) = image.size\n",
    "  image = cv2.cvtColor(image,cv2.COLOR_GRAY2RGB)\n",
    "\n",
    "  return np.array(image.getdata()).reshape(\n",
    "      (im_height, im_width, 3)).astype(np.uint8)\n",
    "\n",
    "def plot_detections(image_np,\n",
    "                    boxes,\n",
    "                    classes,\n",
    "                    scores,\n",
    "                    category_index,\n",
    "                    min_score_threshold = 0.8,\n",
    "                    figsize=(12, 16),\n",
    "                    image_name=None, show_image=True):\n",
    "  \"\"\"Wrapper function to visualize detections.\n",
    "\n",
    "  Args:\n",
    "    image_np: uint8 numpy array with shape (img_height, img_width, 3)\n",
    "    boxes: a numpy array of shape [N, 4]\n",
    "    classes: a numpy array of shape [N]. Note that class indices are 1-based,\n",
    "      and match the keys in the label map.\n",
    "    scores: a numpy array of shape [N] or None.  If scores=None, then\n",
    "      this function assumes that the boxes to be plotted are groundtruth\n",
    "      boxes and plot all boxes as black with no classes or scores.\n",
    "    category_index: a dict containing category dictionaries (each holding\n",
    "      category index `id` and category name `name`) keyed by category indices.\n",
    "    figsize: size for the figure.\n",
    "    image_name: a name for the image file.\n",
    "  \"\"\"\n",
    "  image_np_with_annotations = image_np.copy()\n",
    "  vis_util.visualize_boxes_and_labels_on_image_array(\n",
    "      image_np_with_annotations,\n",
    "      boxes,\n",
    "      classes,\n",
    "      scores,\n",
    "      category_index,\n",
    "      use_normalized_coordinates=True,\n",
    "      min_score_thresh=min_score_threshold,\n",
    "      max_boxes_to_draw=None)\n",
    "  if image_name:\n",
    "    plt.imsave(image_name, image_np_with_annotations)\n",
    "  if show_image:\n",
    "    display(Image.fromarray(image_np_with_annotations))\n",
    "  \n",
    "def prepare_validation_folder():\n",
    "  \"\"\" Move the images that are alocated in the validation folder if exists to the training images folder.\n",
    "  Also does with the annotations.\n",
    "\n",
    "  # Arguments\n",
    "      path:    Path of the validation folder\n",
    "  \n",
    "  \"\"\"\n",
    "\n",
    "  file_suffix = os.path.splitext(os.listdir(Training_Source)[0])[1]\n",
    "  file_annotation_suffix = os.path.splitext(os.listdir(Annotations)[0])[1]\n",
    "\n",
    "  # validation_number = int(train_images_number * percentage_validation/100)\n",
    "  # training_set = train_images_number - validation_number\n",
    "\n",
    "  # if the folders exist all the files into them are moved to the training folders, so that files ae nor lost (deleted)\n",
    "  if os.path.exists(Validation_Source_annotations):\n",
    "    for index, file in enumerate(glob.glob(Validation_Source_annotations+'/*'+file_annotation_suffix)):\n",
    "      filename = file.split('/')[-1]\n",
    "      shutil.move(file, os.path.join(Annotations, filename))\n",
    "    shutil.rmtree(Validation_Source_annotations)\n",
    "  os.mkdir(Validation_Source_annotations)\n",
    "\n",
    "  if os.path.exists(Validation_Source):\n",
    "    for index, file in enumerate(glob.glob(Validation_Source+'/*'+file_suffix)):\n",
    "      filename = file.split('/')[-1]\n",
    "      shutil.move(file, os.path.join(Training_Source, filename))\n",
    "    shutil.rmtree(Validation_Source)\n",
    "  os.mkdir(Validation_Source)\n",
    "\n",
    "  \n",
    "def move_to_validation(filename, img_extension = 'png'):\n",
    "\n",
    "  shutil.move(os.path.join(Training_Source,filename + '.' + img_extension), os.path.join(Validation_Source, filename+'.'+img_extension))\n",
    "  shutil.move(os.path.join(Annotations, filename+'.xml'), os.path.join(Validation_Source_annotations, filename+'.xml'))\n",
    "\n",
    "\n",
    "def compute_ap(recall, precision):\n",
    "  \"\"\" Compute the average precision, given the recall and precision curves.\n",
    "  Code originally from https://github.com/rbgirshick/py-faster-rcnn.\n",
    "\n",
    "  # Arguments\n",
    "      recall:    The recall curve (list).\n",
    "      precision: The precision curve (list).\n",
    "  # Returns\n",
    "      The average precision as computed in py-faster-rcnn.\n",
    "  \"\"\"\n",
    "  # correct AP calculation\n",
    "  # first append sentinel values at the end\n",
    "  mrec = np.concatenate(([0.], recall, [1.]))\n",
    "  mpre = np.concatenate(([0.], precision, [0.]))\n",
    "\n",
    "  # compute the precision envelope\n",
    "  for i in range(mpre.size - 1, 0, -1):\n",
    "      mpre[i - 1] = np.maximum(mpre[i - 1], mpre[i])\n",
    "\n",
    "  # to calculate area under PR curve, look for points\n",
    "  # where X axis (recall) changes value\n",
    "  i = np.where(mrec[1:] != mrec[:-1])[0]\n",
    "\n",
    "  # and sum (\\Delta recall) * prec\n",
    "  ap = np.sum((mrec[i + 1] - mrec[i]) * mpre[i + 1])\n",
    "  return ap   \n",
    "\n",
    "def load_annotation(image_folder,annotations_folder, i, labels):\n",
    "  annots = []\n",
    "  imgs, anns = parse_annotation(annotations_folder,image_folder,labels)\n",
    "  for obj in imgs[i]['object']:\n",
    "      annot = [obj['xmin'], obj['ymin'], obj['xmax'], obj['ymax'], labels.index(obj['name'])]\n",
    "      annots += [annot]\n",
    "\n",
    "  if len(annots) == 0: annots = [[]]\n",
    "\n",
    "  return np.array(annots)\n",
    "\n",
    "def _calc_avg_precisions(model, image_folder,annotations_folder,iou_threshold,score_threshold):\n",
    "\n",
    "  num_labels = len(category_index)\n",
    "  # gather all detections and annotations\n",
    "  all_detections = [[None for _ in range(num_labels)] for _ in range(len(os.listdir(image_folder)))]\n",
    "  all_annotations = [[None for _ in range(num_labels)] for _ in range(len(os.listdir(annotations_folder)))]\n",
    "\n",
    "  for i in range(len(os.listdir(image_folder))):\n",
    "      raw_image = os.path.join(image_folder,sorted(os.listdir(image_folder))[i])\n",
    "\n",
    "      raw_height, raw_width, _ = cv2.imread(raw_image).shape\n",
    "      # make the boxes and the labels\n",
    "\n",
    "      pred_boxes = predict( model,raw_image, show_image = False)[0]\n",
    "\n",
    "      score = pred_boxes['detection_scores'].numpy()\n",
    "\n",
    "      pred_labels = pred_boxes['detection_classes'][0].numpy().astype(int).tolist()\n",
    "      num_predictions = np.where(pred_boxes['detection_scores'].numpy() > score_threshold)[0].shape[0]\n",
    "\n",
    "      # copy detections to all_detections\n",
    "      for label in range(num_labels):\n",
    "        curr_label = np.array([])\n",
    "        for j in range(len(pred_labels)):\n",
    "          if label == pred_labels[j]:\n",
    "            curr_label = np.append(curr_label,[pred_boxes['detection_boxes'].numpy()[0,j,1] * raw_width, pred_boxes['detection_boxes'].numpy()[0,j,0] * raw_height, pred_boxes['detection_boxes'].numpy()[0,j,3] * raw_width,\n",
    "                                  pred_boxes['detection_boxes'].numpy()[0,j,2] * raw_height, pred_boxes['detection_scores'].numpy()[0,j]])\n",
    "        all_detections[i][label] = curr_label.reshape(-1,5)\n",
    "      # print(all_detections)\n",
    "\n",
    "\n",
    "      annotations = load_annotation(image_folder,annotations_folder,i, list(classes_dict.keys()))\n",
    "      # copy ground truth to all_annotations\n",
    "      for label in range(num_labels):\n",
    "        curr_anno = np.array([], dtype=int)\n",
    "        for anno in annotations:\n",
    "          if anno[4] == label:\n",
    "            curr_anno = np.append(curr_anno, anno[:4])\n",
    "        all_annotations[i][label] = curr_anno.reshape(-1,4)\n",
    "\n",
    "\n",
    "  # # compute mAP by comparing all detections and all annotations\n",
    "  average_precisions = {}\n",
    "  F1_scores = {}\n",
    "  total_recall = []\n",
    "  total_precision = []\n",
    "  \n",
    "  QCDataCSVpath = os.path.join(QC_model_folder,'Quality Control/QC_results.csv')\n",
    "  with open(QCDataCSVpath, \"w\", newline='') as file:\n",
    "    writer = csv.writer(file)\n",
    "    writer.writerow([\"class\", \"false positive\", \"true positive\", \"false negative\", \"recall\", \"precision\", \"accuracy\", \"f1 score\", \"average_precision\"])  \n",
    "  \n",
    "    for label in range(num_labels):\n",
    "        false_positives = np.zeros((0,))\n",
    "        true_positives = np.zeros((0,))\n",
    "        scores = np.zeros((0,))\n",
    "        num_annotations = 0.0\n",
    "\n",
    "        for i in range(len(os.listdir(image_folder))):\n",
    "            detections = all_detections[i][label]\n",
    "            annotations = all_annotations[i][label]\n",
    "            num_annotations += annotations.shape[0]\n",
    "            detected_annotations = []\n",
    "\n",
    "            for d in detections:\n",
    "                scores = np.append(scores, d[4])\n",
    "\n",
    "                if annotations.shape[0] == 0:\n",
    "                    false_positives = np.append(false_positives, 1)\n",
    "                    true_positives = np.append(true_positives, 0)\n",
    "                    continue\n",
    "\n",
    "                overlaps = compute_overlap(np.expand_dims(d, axis=0), annotations)\n",
    "                assigned_annotation = np.argmax(overlaps, axis=1)\n",
    "                max_overlap = overlaps[0, assigned_annotation]\n",
    "\n",
    "                if max_overlap >= iou_threshold and assigned_annotation not in detected_annotations:\n",
    "                    false_positives = np.append(false_positives, 0)\n",
    "                    true_positives = np.append(true_positives, 1)\n",
    "                    detected_annotations.append(assigned_annotation)\n",
    "                else:\n",
    "                    false_positives = np.append(false_positives, 1)\n",
    "                    true_positives = np.append(true_positives, 0)\n",
    "\n",
    "\n",
    "\n",
    "        # no annotations -> AP for this class is 0 (is this correct?)\n",
    "        if num_annotations == 0:\n",
    "            average_precisions[label] = 0\n",
    "            continue\n",
    "\n",
    "        # sort by score\n",
    "        indices = np.argsort(-scores)\n",
    "        false_positives = false_positives[indices]\n",
    "        true_positives = true_positives[indices]\n",
    "\n",
    "        # compute false positives and true positives\n",
    "        false_positives = np.cumsum(false_positives)\n",
    "        true_positives = np.cumsum(true_positives)\n",
    "\n",
    "        # compute recall and precision\n",
    "        recall = true_positives / num_annotations\n",
    "        precision = true_positives / np.maximum(true_positives + false_positives, np.finfo(np.float64).eps)\n",
    "        total_recall.append(recall)\n",
    "        total_precision.append(precision)\n",
    "\n",
    "        # compute average precision\n",
    "        average_precision = compute_ap(recall, precision)\n",
    "        average_precisions[label] = average_precision\n",
    "\n",
    "        if len(precision) != 0:\n",
    "          F1_score = 2*(precision[-1]*recall[-1]/(precision[-1]+recall[-1]))\n",
    "          F1_scores[label] = F1_score\n",
    "          writer.writerow([list(classes_dict.keys())[label], str(int(false_positives[-1])), str(int(true_positives[-1])), str(int(num_annotations-true_positives[-1])), str(recall[-1]), str(precision[-1]), str(true_positives[-1]/num_annotations), str(F1_scores[label]), str(average_precisions[label])])\n",
    "        else:\n",
    "          F1_score = 0\n",
    "          F1_scores[label] = F1_score\n",
    "          writer.writerow([list(classes_dict.keys())[label], str(0), str(0), str(0), str(0), str(0), str(0), str(F1_score), str(average_precisions[label])])\n",
    "    return F1_scores, average_precisions, total_recall, total_precision\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "class BoundBox:\n",
    "    \"\"\"\n",
    "    Here, we are adapting classes and functions from https://github.com/MathGaron/mean_average_precision\n",
    "    \"\"\"\n",
    "    def __init__(self, xmin, ymin, xmax, ymax, c = None, classes = None):\n",
    "        self.xmin = xmin\n",
    "        self.ymin = ymin\n",
    "        self.xmax = xmax\n",
    "        self.ymax = ymax\n",
    "        \n",
    "        self.c     = c\n",
    "        self.classes = classes\n",
    "\n",
    "        self.label = -1\n",
    "        self.score = -1\n",
    "\n",
    "    def get_label(self):\n",
    "        if self.label == -1:\n",
    "            self.label = np.argmax(self.classes)\n",
    "        \n",
    "        return self.label\n",
    "    \n",
    "    def get_score(self):\n",
    "        if self.score == -1:\n",
    "            self.score = self.classes[self.get_label()]\n",
    "            \n",
    "        return self.score\n",
    "\n",
    "class WeightReader:\n",
    "    def __init__(self, weight_file):\n",
    "        self.offset = 4\n",
    "        self.all_weights = np.fromfile(weight_file, dtype='float32')\n",
    "        \n",
    "    def read_bytes(self, size):\n",
    "        self.offset = self.offset + size\n",
    "        return self.all_weights[self.offset-size:self.offset]\n",
    "    \n",
    "    def reset(self):\n",
    "        self.offset = 4\n",
    "\n",
    "def bbox_iou(box1, box2):\n",
    "    intersect_w = _interval_overlap([box1.xmin, box1.xmax], [box2.xmin, box2.xmax])\n",
    "    intersect_h = _interval_overlap([box1.ymin, box1.ymax], [box2.ymin, box2.ymax])  \n",
    "    \n",
    "    intersect = intersect_w * intersect_h\n",
    "\n",
    "    w1, h1 = box1.xmax-box1.xmin, box1.ymax-box1.ymin\n",
    "    w2, h2 = box2.xmax-box2.xmin, box2.ymax-box2.ymin\n",
    "    \n",
    "    union = w1*h1 + w2*h2 - intersect\n",
    "    \n",
    "    return float(intersect) / union\n",
    "\n",
    "def save_boxes(image_path, boxes, scores, boxes_labels, labels):#, save_path):\n",
    "    image = cv2.imread(image_path)\n",
    "    image_h, image_w, _ = image.shape\n",
    "    save_boxes =[]\n",
    "    save_boxes_names = []\n",
    "    save_boxes.append(os.path.basename(image_path))\n",
    "    save_boxes_names.append(os.path.basename(image_path))\n",
    "    for i in range(len(boxes)):\n",
    "        # xmin = box.xmin\n",
    "        save_boxes.append(int(boxes[i,1]*image_w))\n",
    "        save_boxes_names.append(int(boxes[i,1]*image_w))\n",
    "        # ymin = box.ymin\n",
    "        save_boxes.append(int(boxes[i,0]*image_h))\n",
    "        save_boxes_names.append(int(boxes[i,0]*image_h))\n",
    "        # xmax = box.xmax\n",
    "        save_boxes.append(int(boxes[i,3]*image_w))\n",
    "        save_boxes_names.append(int(boxes[i,3]*image_w))\n",
    "        # ymax = box.ymax\n",
    "        save_boxes.append(int(boxes[i,2]*image_h))\n",
    "        save_boxes_names.append(int(boxes[i,2]*image_h))\n",
    "        \n",
    "        save_boxes.append(scores[i])\n",
    "        save_boxes_names.append(scores[i])\n",
    "        save_boxes.append(boxes_labels[i])\n",
    "        save_boxes_names.append(category_index[boxes_labels[i]+1]['name'])\n",
    "    \n",
    "    #This file will be for later analysis of the bounding boxes in imagej\n",
    "    if not os.path.exists(base_path + '/predicted_bounding_boxes.csv'):\n",
    "      with open(base_path + '/predicted_bounding_boxes.csv', 'w', newline='') as csvfile:\n",
    "        csvwriter = csv.writer(csvfile, delimiter=',')\n",
    "        specs_list = ['filename']+['xmin', 'ymin', 'xmax', 'ymax', 'confidence', 'class']*len(boxes)\n",
    "        csvwriter.writerow(specs_list)\n",
    "        csvwriter.writerow(save_boxes)\n",
    "    else:\n",
    "      with open(base_path + '/predicted_bounding_boxes.csv', 'a+', newline='') as csvfile:\n",
    "        csvwriter = csv.writer(csvfile)\n",
    "        csvwriter.writerow(save_boxes)\n",
    "    \n",
    "    if not os.path.exists(base_path + '/predicted_bounding_boxes_names.csv'):\n",
    "      with open(base_path + '/predicted_bounding_boxes_names.csv', 'w', newline='') as csvfile_names:\n",
    "        csvwriter = csv.writer(csvfile_names, delimiter=',')\n",
    "        specs_list = ['filename']+['xmin', 'ymin', 'xmax', 'ymax', 'confidence', 'class']*len(boxes)\n",
    "        csvwriter.writerow(specs_list)\n",
    "        csvwriter.writerow(save_boxes_names)\n",
    "    else:\n",
    "      with open(base_path + '/predicted_bounding_boxes_names.csv', 'a+', newline='') as csvfile_names:\n",
    "        csvwriter = csv.writer(csvfile_names)\n",
    "        csvwriter.writerow(save_boxes_names)\n",
    "    # #This file is to create a nicer display for the output images\n",
    "    # if not os.path.exists(base_path + '/predicted_bounding_boxes_display.csv'):\n",
    "    #   with open(base_path + '/predicted_bounding_boxes_display.csv', 'w', newline='') as csvfile_new:\n",
    "    #     csvwriter2 = csv.writer(csvfile_new, delimiter=',')\n",
    "    #     specs_list = ['filename','width','height','class','xmin','ymin','xmax','ymax']\n",
    "    #     csvwriter2.writerow(specs_list)\n",
    "    # else:\n",
    "    #   with open(base_path + '/predicted_bounding_boxes_display.csv','a+',newline='') as csvfile_new:\n",
    "    #     csvwriter2 = csv.writer(csvfile_new)\n",
    "    #     for box in boxes:\n",
    "    #       row = [os.path.basename(image_path),image_w,image_h,box.get_label(),int(box.xmin*image_w),int(box.ymin*image_h),int(box.xmax*image_w),int(box.ymax*image_h)]\n",
    "    #       csvwriter2.writerow(row)\n",
    "\n",
    "def add_header(inFilePath,outFilePath):\n",
    "  header = ['filename']+['xmin', 'ymin', 'xmax', 'ymax', 'confidence', 'class']*100#max(n_objects)\n",
    "  with open(inFilePath, newline='') as inFile, open(outFilePath, 'w', newline='') as outfile:\n",
    "      r = csv.reader(inFile)\n",
    "      w = csv.writer(outfile)\n",
    "      next(r, None)  # skip the first row from the reader, the old header\n",
    "      # write new header\n",
    "      w.writerow(header)\n",
    "      # copy the rest\n",
    "      for row in r:\n",
    "          w.writerow(row)\n",
    " \n",
    "def load_train_images(train_path):\n",
    "  train_images_np = []\n",
    "\n",
    "  class_images = sorted(os.listdir(train_path))\n",
    "  for img in class_images:\n",
    "    image_path = os.path.join(os.path.join(train_path, img))\n",
    "    img = np.array(Image.open(path))\n",
    "    img = cv2.cvtColor(img,cv2.COLOR_GRAY2RGB)\n",
    "    train_images_np.append(img)\n",
    "  return train_images_np\n",
    "\n",
    "\n",
    "def load_model_checkpoints( checkpoint_path, pipeline_config, num_classes, iou_threshold = 0.6, classification_weight= 1.0, localization_weight= 1.0):\n",
    "    configs = config_util.get_configs_from_pipeline_file(pipeline_config)\n",
    "    model_config = configs['model']\n",
    "    model_config.ssd.num_classes = num_classes\n",
    "    model_config.ssd.post_processing.batch_non_max_suppression.iou_threshold  = iou_threshold\n",
    "    model_config.ssd.loss.classification_weight = classification_weight\n",
    "    model_config.ssd.loss.localization_weight = localization_weight\n",
    "\n",
    "    model_config.ssd.freeze_batchnorm = True\n",
    "    detection_model = model_builder.build(\n",
    "    model_config=model_config, is_training=True)\n",
    "\n",
    "    fake_box_predictor = tf.compat.v2.train.Checkpoint(\n",
    "    _base_tower_layers_for_heads=detection_model._box_predictor._base_tower_layers_for_heads,\n",
    "    _box_prediction_head=detection_model._box_predictor._box_prediction_head,\n",
    "    )\n",
    "\n",
    "    fake_model = tf.compat.v2.train.Checkpoint(\n",
    "          _feature_extractor=detection_model._feature_extractor,\n",
    "          _box_predictor=fake_box_predictor)\n",
    "    ckpt = tf.compat.v2.train.Checkpoint(model=fake_model)\n",
    "    ckpt.restore(checkpoint_path).expect_partial()\n",
    "\n",
    "    # Run model through a dummy image so that variables are created\n",
    "    image, shapes = detection_model.preprocess(tf.zeros([1, 640, 640, 3]))\n",
    "    prediction_dict = detection_model.predict(image, shapes)\n",
    "    _ = detection_model.postprocess(prediction_dict, shapes)\n",
    "\n",
    "    configs['model'] = model_config\n",
    "    return detection_model, configs, ckpt\n",
    "\n",
    "def save_model_config(config_to_save):\n",
    "  os.makedirs(saved_model_path+\"/config\")\n",
    "  os.makedirs(saved_model_path+\"/checkpoints\")\n",
    "\n",
    "  new_pipeline_proto = config_util.create_pipeline_proto_from_configs(config_to_save)\n",
    "  config_util.save_pipeline_config(new_pipeline_proto, saved_model_path + '/config/')\n",
    "\n",
    "def save_category_index():\n",
    "  if os.path.exists(saved_model_path+\"/category_index\"):\n",
    "    shutil.rmtree(saved_model_path+\"/category_index\")\n",
    "  os.makedirs(saved_model_path+\"/category_index\")\n",
    "  with open(saved_model_path+\"/category_index/category_index.json\", 'w', encoding='utf-8') as f:\n",
    "    json.dump(category_index, f, ensure_ascii=False, indent=4)\n",
    "\n",
    "def load_pretrained_model(model_path):\n",
    "  \"\"\" Load model's checkpoints from one previously saved model folder. The models are saved everytime \n",
    "        the training is finished.\n",
    "\n",
    "  # Arguments\n",
    "      model_path: the path to the model's 'saved_model' folder.\n",
    "  # Returns\n",
    "      The pretrained model, category index dictionary and classes dictionary.\n",
    "  \"\"\"\n",
    "  pipeline_config =  os.path.join(model_path, 'config', 'pipeline.config')\n",
    "  configs = config_util.get_configs_from_pipeline_file(pipeline_config)\n",
    "  model_config = configs['model']\n",
    "\n",
    "  category_index_path = os.path.join(model_path, \"category_index\", \"category_index.json\")\n",
    "  with open(category_index_path) as json_file:\n",
    "    category_index_string = json.load(json_file)\n",
    "\n",
    "  classes_dict = dict()\n",
    "  category_index = dict()\n",
    "  for key,val in category_index_string.items():\n",
    "    classes_dict[val['name']] = int(key)\n",
    "    category_index[int(key)] = val \n",
    "\n",
    "  strategy = tf.compat.v2.distribute.get_strategy()\n",
    "  with strategy.scope():\n",
    "    pretrained_model = model_builder.build(model_config=configs['model'], is_training=True)\n",
    "\n",
    "  checkpoints = tf.compat.v2.train.Checkpoint( model=pretrained_model)\n",
    "\n",
    "  latest_checkpoint = tf.train.latest_checkpoint(os.path.join(model_path, \"checkpoints\"))\n",
    "  checkpoints.restore(latest_checkpoint)\n",
    "  return pretrained_model, category_index, classes_dict, configs\n",
    "\n",
    "def prepare_data_to_train(training_img_path, df, data_gt_boxes, data_gt_classes):\n",
    "  train_image_tensors = []\n",
    "  gt_classes_one_hot_tensors = []\n",
    "  gt_box_tensors = []\n",
    "\n",
    "  for index, row in df.iterrows():\n",
    "    gt_box_tensors.append(tf.convert_to_tensor(data_gt_boxes[index], dtype=tf.float32))\n",
    "\n",
    "    img_dir = training_img_path\n",
    "    path = os.path.join(img_dir, row[\"fileID\"] + file_suffix)\n",
    "    img = np.array(Image.open(path))\n",
    "    img = cv2.cvtColor(img,cv2.COLOR_GRAY2RGB)\n",
    "    train_image_tensors.append(tf.expand_dims(tf.convert_to_tensor(\n",
    "      img, dtype=tf.float32), axis=0))\n",
    "    \n",
    "    predicted_classes = np.zeros(shape=[data_gt_boxes[index].shape[0]], dtype=np.int32)\n",
    "    for i in range(data_gt_classes[index].shape[0]):\n",
    "      predicted_classes[i] = classes_dict[data_gt_classes[index][i,0]]\n",
    "\n",
    "    zero_indexed_groundtruth_classes = tf.convert_to_tensor(predicted_classes-1)\n",
    "    gt_classes_one_hot_tensors.append(tf.one_hot(\n",
    "     zero_indexed_groundtruth_classes , len(class_count)))\n",
    "\n",
    "  return train_image_tensors, gt_classes_one_hot_tensors, gt_box_tensors\n",
    "\n",
    "def train(model, verbose = True):\n",
    "  train_image_tensors, gt_classes_one_hot_tensors, gt_box_tensors = prepare_data_to_train(augmented_training_source, df_anno, gt_boxes, gt_classes)\n",
    "  print('Done training data preprocessing.')\n",
    "\n",
    "  validation_image_tensors, validation_gt_classes_one_hot_tensors, validation_gt_box_tensors = prepare_data_to_train(Validation_Source, df_anno_valid, gt_boxes_valid, gt_classes_valid)\n",
    "  print('Done validation data preprocessing.')\n",
    "\n",
    "  batches_per_epochs = len(train_image_tensors)//batch_size\n",
    "  num_batches = number_of_epochs * batches_per_epochs\n",
    "\n",
    "\n",
    "  trainable_variables = model.trainable_variables\n",
    "  to_fine_tune = []\n",
    "  prefixes_to_train = [\n",
    "  'WeightSharedConvolutionalBoxPredictor/WeightSharedConvolutionalBoxHead',\n",
    "  'WeightSharedConvolutionalBoxPredictor/WeightSharedConvolutionalClassHead']\n",
    "\n",
    "  for var in trainable_variables:\n",
    "    if any([var.name.startswith(prefix) for prefix in prefixes_to_train]):\n",
    "      to_fine_tune.append(var)\n",
    "\n",
    "\n",
    "  def get_model_train_step_function(model, optimizer, vars_to_fine_tune):\n",
    "    \"\"\"Get a tf.function for training step.\"\"\"\n",
    "\n",
    "    def evaluate(image_tensors,\n",
    "                        groundtruth_boxes_list,\n",
    "                        groundtruth_classes_list):\n",
    "        \"\"\"A single validation iteration.\n",
    "\n",
    "        Args:\n",
    "          image_tensors: A list of [1, height, width, 3] Tensor of type tf.float32.\n",
    "            Note that the height and width can vary across images, as they are\n",
    "            reshaped within this function to be 640x640.\n",
    "          groundtruth_boxes_list: A list of Tensors of shape [N_i, 4] with type\n",
    "            tf.float32 representing groundtruth boxes for each image in the batch.\n",
    "          groundtruth_classes_list: A list of Tensors of shape [N_i, num_classes]\n",
    "            with type tf.float32 representing groundtruth boxes for each image in\n",
    "            the batch.\n",
    "\n",
    "        Returns:\n",
    "          A scalar tensor representing the total loss for the input batch.\n",
    "        \"\"\"\n",
    "        shapes = tf.constant(len(image_tensors) * [[640, 640, 3]], dtype=tf.int32)\n",
    "        model.provide_groundtruth(\n",
    "              groundtruth_boxes_list=groundtruth_boxes_list,\n",
    "              groundtruth_classes_list=groundtruth_classes_list)\n",
    "        \n",
    "        preprocessed_images = tf.concat(\n",
    "            [model.preprocess(image_tensor)[0]\n",
    "            for image_tensor in image_tensors], axis=0)\n",
    "        prediction_dict = model.predict(preprocessed_images, shapes)\n",
    "\n",
    "\n",
    "        losses_dict = model.loss(prediction_dict, shapes)\n",
    "        total_loss = losses_dict['Loss/localization_loss'] + losses_dict['Loss/classification_loss']\n",
    "        return total_loss, losses_dict['Loss/localization_loss'], losses_dict['Loss/classification_loss']\n",
    "\n",
    "    # Use tf.function for a bit of speed.\n",
    "    # Comment out the tf.function decorator if you want the inside of the\n",
    "    # function to run eagerly.\n",
    "    # @tf.function\n",
    "    def train_step_fn(image_tensors,\n",
    "                        groundtruth_boxes_list,\n",
    "                        groundtruth_classes_list):\n",
    "        \"\"\"A single training iteration.\n",
    "\n",
    "        Args:\n",
    "          image_tensors: A list of [1, height, width, 3] Tensor of type tf.float32.\n",
    "            Note that the height and width can vary across images, as they are\n",
    "            reshaped within this function to be 640x640.\n",
    "          groundtruth_boxes_list: A list of Tensors of shape [N_i, 4] with type\n",
    "            tf.float32 representing groundtruth boxes for each image in the batch.\n",
    "          groundtruth_classes_list: A list of Tensors of shape [N_i, num_classes]\n",
    "            with type tf.float32 representing groundtruth boxes for each image in\n",
    "            the batch.\n",
    "\n",
    "        Returns:\n",
    "          A scalar tensor representing the total loss for the input batch.\n",
    "        \"\"\"\n",
    "        shapes = tf.constant(batch_size * [[640, 640, 3]], dtype=tf.int32)\n",
    "        model.provide_groundtruth(\n",
    "              groundtruth_boxes_list=groundtruth_boxes_list,\n",
    "              groundtruth_classes_list=groundtruth_classes_list)\n",
    "        with tf.GradientTape() as tape:\n",
    "          preprocessed_images = tf.concat(\n",
    "              [model.preprocess(image_tensor)[0]\n",
    "              for image_tensor in image_tensors], axis=0)\n",
    "          prediction_dict = model.predict(preprocessed_images, shapes)\n",
    "\n",
    "\n",
    "          losses_dict = model.loss(prediction_dict, shapes)\n",
    "          total_loss = losses_dict['Loss/localization_loss'] + losses_dict['Loss/classification_loss']\n",
    "          gradients = tape.gradient(total_loss, vars_to_fine_tune)\n",
    "          optimizer.apply_gradients(zip(gradients, vars_to_fine_tune))\n",
    "        return total_loss, losses_dict['Loss/localization_loss'], losses_dict['Loss/classification_loss']\n",
    "\n",
    "    return train_step_fn, evaluate                                                    \n",
    "\n",
    "  # optimizer = tf.keras.optimizers.SGD(learning_rate=learning_rate, momentum=0.9)\n",
    "  optimizer = tf.keras.optimizers.Adam(learning_rate=learning_rate)\n",
    "  train_step_fn, evaluate = get_model_train_step_function(\n",
    "  model, optimizer, to_fine_tune)\n",
    "\n",
    "\n",
    "  ckpt = tf.compat.v2.train.Checkpoint(model=model)\n",
    "  manager = tf.compat.v2.train.CheckpointManager(\n",
    "        ckpt, saved_model_path+\"/checkpoints\", max_to_keep=None)\n",
    "  \n",
    "\n",
    "  lossDataCSVpath = os.path.join(full_model_path,'Quality Control/training_evaluation.csv')\n",
    "  with open(lossDataCSVpath, 'w') as f1:\n",
    "    writer = csv.writer(f1)\n",
    "    writer.writerow(['loss','val_loss','mAP','learning rate'])\n",
    "\n",
    "    total_losses, localization_losses, classification_losses, epochNumber, validation_total_losses, val_localization_losses, val_classification_losses, val_mAP = [], [], [], [], [], [], [], []\n",
    "    batch_index = 0\n",
    "    for idx in range(num_batches):\n",
    "\n",
    "      # Grab keys for subset of data\n",
    "      all_keys = list(range(len(train_image_tensors)))\n",
    "      batch_init = batch_size * (batch_index//batch_size)\n",
    "      batch_end = batch_size * (batch_index//batch_size + 1)\n",
    "\n",
    "      example_keys = all_keys[batch_init:batch_end]\n",
    "\n",
    "      gt_boxes_list = [gt_box_tensors[key] for key in example_keys]\n",
    "      gt_classes_list = [gt_classes_one_hot_tensors[key] for key in example_keys]\n",
    "      image_tensors = [train_image_tensors[key] for key in example_keys]\n",
    "\n",
    "      batch_index +=1\n",
    "      if batch_end >= len(train_image_tensors):\n",
    "        batch_index = 0\n",
    "\n",
    "      # Training step (forward pass + backwards pass)\n",
    "      total_loss, localization_loss, classification_loss = train_step_fn(image_tensors, gt_boxes_list, gt_classes_list)\n",
    "\n",
    "\n",
    "      if idx % batches_per_epochs  == 0:\n",
    "        current_epoch = idx // batches_per_epochs\n",
    "        if (current_epoch +1) % number_of_epochs == 0:\n",
    "          print(f\"Model saved to {saved_model_path }\")\n",
    "          manager.save()\n",
    "      \n",
    "        #validation\n",
    "        val_total_loss, val_localization_loss, val_classification_loss = evaluate( validation_image_tensors, validation_gt_box_tensors, validation_gt_classes_one_hot_tensors)\n",
    "        \n",
    "        print('Epoch ' + str(current_epoch+1) + '/' + str(number_of_epochs)\n",
    "        + ', Total loss = ' +  str(total_loss.numpy()) + ' Localization loss = ' + str(localization_loss.numpy()) + ' Classification loss = ' + str(classification_loss.numpy()) + ' Validation loss = ' + str(val_total_loss.numpy()), flush=True)\n",
    "\n",
    "        if verbose:\n",
    "          # mAP calculator\n",
    "          f1, ap, total_recall, total_precision= _calc_avg_precisions(model, Validation_Source, Validation_Source_annotations,  0.3, 0.6)\n",
    "          mAP = np.sum(list(ap.values()))/len(ap)\n",
    "\n",
    "          print(\"\\nAverage Precision:\")\n",
    "          for key in ap:\n",
    "            print(f\"{category_index[key+1]['name']}: {ap[key]}\", flush=True)\n",
    "\n",
    "          print(f\"\\nmAP: {mAP}\\n\")\n",
    "\n",
    "        if verbose:\n",
    "          writer.writerow([total_loss.numpy(), val_total_loss.numpy(), mAP, learning_rate])\n",
    "        else:\n",
    "          writer.writerow([total_loss.numpy(), val_total_loss.numpy(), learning_rate])\n",
    "\n",
    "  manager.save()\n",
    "\n",
    "  return manager\n",
    "\n",
    "def predict(model, images_path, save_path = None, show_image = True, min_score_threshold=0.6):\n",
    "  test_images_np = []\n",
    "  if os.path.isdir(images_path):\n",
    "    for img_name in os.listdir(images_path):\n",
    "      path = os.path.join(images_path, img_name)\n",
    "      img = np.array(Image.open(path))\n",
    "      img = cv2.cvtColor(img,cv2.COLOR_GRAY2RGB)\n",
    "      test_images_np.append(np.expand_dims(\n",
    "        img, axis=0))\n",
    "  else:\n",
    "    img = np.array(Image.open(images_path))\n",
    "    img = cv2.cvtColor(img,cv2.COLOR_GRAY2RGB)\n",
    "    test_images_np.append(np.expand_dims(\n",
    "      img, axis=0))\n",
    "\n",
    "  def detect(input_tensor):\n",
    "    \"\"\"Run detection on an input image.\n",
    "\n",
    "    Args:\n",
    "      input_tensor: A [1, height, width, 3] Tensor of type tf.float32.\n",
    "        Note that height and width can be anything since the image will be\n",
    "        immediately resized according to the needs of the model within this\n",
    "        function.\n",
    "\n",
    "    Returns:\n",
    "      A dict containing 3 Tensors (`detection_boxes`, `detection_classes`,\n",
    "        and `detection_scores`).\n",
    "    \"\"\"\n",
    "    preprocessed_image, shapes = model.preprocess(input_tensor)\n",
    "    prediction_dict = model.predict(preprocessed_image, shapes)\n",
    "    return model.postprocess(prediction_dict, shapes)\n",
    "\n",
    "# Note that the first frame will trigger tracing of the tf.function, which will\n",
    "# take some time, after which inference should be fast.\n",
    "\n",
    "  label_id_offset = 1\n",
    "  predictions = []\n",
    "  for i in range(len(test_images_np)):\n",
    "    input_tensor = tf.convert_to_tensor(test_images_np[i], dtype=tf.float32)\n",
    "    detections = detect(input_tensor)\n",
    "\n",
    "    # Supress overlapping\n",
    "    selected_indices = tf.image.non_max_suppression(\n",
    "      detections['detection_boxes'][0].numpy(),  detections['detection_scores'][0].numpy(), max_output_size=detections['detection_scores'][0].numpy().shape[0], iou_threshold=0.3)\n",
    "    \n",
    "    sboxes = tf.gather(detections['detection_boxes'][0], selected_indices)\n",
    "    sscores = tf.gather(detections['detection_scores'][0], selected_indices)\n",
    "    sclasses = tf.gather(detections['detection_classes'][0], selected_indices)\n",
    "\n",
    "\n",
    "    plot_detections(\n",
    "        test_images_np[i][0],\n",
    "        sboxes.numpy(),\n",
    "        sclasses.numpy().astype(np.uint32)\n",
    "        + label_id_offset,\n",
    "        sscores.numpy(),\n",
    "        category_index, \n",
    "        min_score_threshold = min_score_threshold, image_name=save_path, show_image = show_image)\n",
    "\n",
    "    save_boxes(images_path,sboxes.numpy(), sscores.numpy(),sclasses.numpy().astype(np.uint32), list(classes_dict.keys()))\n",
    "    predictions.append(detections)\n",
    "\n",
    "  return predictions\n",
    "\n",
    "# function to convert BoundingBoxesOnImage object into DataFrame\n",
    "def bbs_obj_to_df(bbs_object):\n",
    "#     convert BoundingBoxesOnImage object into array\n",
    "    bbs_array = bbs_object.to_xyxy_array()\n",
    "#     convert array into a DataFrame ['xmin', 'ymin', 'xmax', 'ymax'] columns\n",
    "    df_bbs = pd.DataFrame(bbs_array, columns=['xmin', 'ymin', 'xmax', 'ymax'])\n",
    "    return df_bbs\n",
    "\n",
    "# Function that will extract column data for our CSV file\n",
    "def xml_to_csv(path):\n",
    "    xml_list = []\n",
    "    for xml_file in glob.glob(path + '/*.xml'):\n",
    "        tree = ET.parse(xml_file)\n",
    "        root = tree.getroot()\n",
    "        for member in root.findall('object'):\n",
    "            value = (root.find('filename').text,\n",
    "                     int(root.find('size')[0].text),\n",
    "                     int(root.find('size')[1].text),\n",
    "                     member[0].text,\n",
    "                     int(member[4][0].text),\n",
    "                     int(member[4][1].text),\n",
    "                     int(member[4][2].text),\n",
    "                     int(member[4][3].text)\n",
    "                     )\n",
    "            xml_list.append(value)\n",
    "    column_name = ['filename', 'width', 'height', 'class', 'xmin', 'ymin', 'xmax', 'ymax']\n",
    "    xml_df = pd.DataFrame(xml_list, columns=column_name)\n",
    "    return xml_df\n",
    "\n",
    "\n",
    "\n",
    "def image_aug(df, images_path, aug_images_path, image_prefix, augmentor):\n",
    "    # create data frame which we're going to populate with augmented image info\n",
    "    aug_bbs_xy = pd.DataFrame(columns=\n",
    "                              ['filename','width','height','class', 'xmin', 'ymin', 'xmax', 'ymax']\n",
    "                             )\n",
    "    grouped = df.groupby('filename')\n",
    "    \n",
    "    for filename in df['filename'].unique():\n",
    "    #   get separate data frame grouped by file name\n",
    "        group_df = grouped.get_group(filename)\n",
    "        group_df = group_df.reset_index()\n",
    "        group_df = group_df.drop(['index'], axis=1)   \n",
    "    #   read the image\n",
    "        image = imageio.imread(images_path+filename)\n",
    "    #   get bounding boxes coordinates and write into array        \n",
    "        bb_array = group_df.drop(['filename', 'width', 'height', 'class'], axis=1).values\n",
    "    #   pass the array of bounding boxes coordinates to the imgaug library\n",
    "        bbs = BoundingBoxesOnImage.from_xyxy_array(bb_array, shape=image.shape)\n",
    "    #   apply augmentation on image and on the bounding boxes\n",
    "        image_aug, bbs_aug = augmentor(image=image, bounding_boxes=bbs)\n",
    "    #   disregard bounding boxes which have fallen out of image pane    \n",
    "        bbs_aug = bbs_aug.remove_out_of_image()\n",
    "    #   clip bounding boxes which are partially outside of image pane\n",
    "        bbs_aug = bbs_aug.clip_out_of_image()\n",
    "        \n",
    "    #   don't perform any actions with the image if there are no bounding boxes left in it    \n",
    "        if re.findall('Image...', str(bbs_aug)) == ['Image([]']:\n",
    "            pass\n",
    "        \n",
    "    #   otherwise continue\n",
    "        else:\n",
    "        #   write augmented image to a file\n",
    "            imageio.imwrite(aug_images_path+image_prefix+filename, image_aug)  \n",
    "        #   create a data frame with augmented values of image width and height\n",
    "            info_df = group_df.drop(['xmin', 'ymin', 'xmax', 'ymax'], axis=1)    \n",
    "            for index, _ in info_df.iterrows():\n",
    "                info_df.at[index, 'width'] = image_aug.shape[1]\n",
    "                info_df.at[index, 'height'] = image_aug.shape[0]\n",
    "        #   rename filenames by adding the predifined prefix\n",
    "            info_df['filename'] = info_df['filename'].apply(lambda x: image_prefix+x)\n",
    "        #   create a data frame with augmented bounding boxes coordinates using the function we created earlier\n",
    "            bbs_df = bbs_obj_to_df(bbs_aug)\n",
    "        #   concat all new augmented info into new data frame\n",
    "            aug_df = pd.concat([info_df, bbs_df], axis=1)\n",
    "        #   append rows to aug_bbs_xy data frame\n",
    "            aug_bbs_xy = pd.concat([aug_bbs_xy, aug_df])            \n",
    "    \n",
    "    # return dataframe with updated images and bounding boxes annotations \n",
    "    aug_bbs_xy = aug_bbs_xy.reset_index()\n",
    "    aug_bbs_xy = aug_bbs_xy.drop(['index'], axis=1)\n",
    "    return aug_bbs_xy\n",
    "\n",
    "\n",
    "print('-------------------------------------------')\n",
    "print(\"Depencies installed and imported.\")\n",
    "\n",
    "\n",
    "# Colors for the warning messages\n",
    "class bcolors:\n",
    "  WARNING = '\\033[31m'\n",
    "  NORMAL = '\\033[0m'\n",
    "  OKGREEN = '\\033[92m'\n",
    "\n",
    "# Check if this is the latest version of the notebook\n",
    "All_notebook_versions = pd.read_csv(\"https://raw.githubusercontent.com/HenriquesLab/ZeroCostDL4Mic/master/Colab_notebooks/Latest_Notebook_versions.csv\", dtype=str)\n",
    "print('Notebook version: '+Notebook_version)\n",
    "Latest_Notebook_version = All_notebook_versions[All_notebook_versions[\"Notebook\"] == Network]['Version'].iloc[0]\n",
    "print('Latest notebook version: '+Latest_Notebook_version)\n",
    "if Notebook_version == Latest_Notebook_version:\n",
    "  print(\"This notebook is up-to-date.\")\n",
    "else:\n",
    "  print(bcolors.WARNING +\"A new version of this notebook has been released. We recommend that you download it at https://github.com/HenriquesLab/ZeroCostDL4Mic/wiki\")\n",
    "\n",
    "\n",
    "\n",
    "print(Back.BLACK +bcolors.OKGREEN + \"Everything has set up correctly, please don't worry about the possible error messages.\")\n",
    "print(Style.RESET_ALL)\n",
    "\n",
    "#Create a pdf document with training summary\n",
    "\n",
    "# save FPDF() class into a  \n",
    "# variable pdf \n",
    "\n",
    "def pdf_export(trained = False, augmentation = False, pretrained_model = False):\n",
    "  class MyFPDF(FPDF, HTMLMixin):\n",
    "    pass\n",
    "\n",
    "  pdf = MyFPDF()\n",
    "  pdf.add_page()\n",
    "  pdf.set_right_margin(-1)\n",
    "  pdf.set_font(\"Arial\", size = 11, style='B') \n",
    "\n",
    "  Network = 'RetinaNet'\n",
    "  day = datetime.now()\n",
    "  datetime_str = str(day)[0:10]\n",
    "\n",
    "  Header = 'Training report for '+Network+' model ('+model_name+'):\\nDate: '+datetime_str\n",
    "  pdf.multi_cell(180, 5, txt = Header, align = 'L') \n",
    "  pdf.ln(1)\n",
    "\n",
    "  # add another cell\n",
    "  if trained:\n",
    "    training_time = \"Training time: \"+str(hour)+ \"hour(s) \"+str(mins)+\"min(s) \"+str(round(sec))+\"sec(s)\"\n",
    "    pdf.cell(190, 5, txt = training_time, ln = 1, align='L')\n",
    "  pdf.ln(1)\n",
    "\n",
    "  Header_2 = 'Information for your materials and methods:'\n",
    "  pdf.cell(190, 5, txt=Header_2, ln=1, align='L')\n",
    "\n",
    "  all_packages = ''\n",
    "  for requirement in freeze(local_only=True):\n",
    "    all_packages = all_packages+requirement+', '\n",
    "  #print(all_packages)\n",
    "\n",
    "  #Main Packages\n",
    "  main_packages = ''\n",
    "  version_numbers = []\n",
    "  for name in ['tensorflow','numpy','Keras']:\n",
    "    find_name=all_packages.find(name)\n",
    "    main_packages = main_packages+all_packages[find_name:all_packages.find(',',find_name)]+', '\n",
    "    #Version numbers only here:\n",
    "    version_numbers.append(all_packages[find_name+len(name)+2:all_packages.find(',',find_name)])\n",
    "\n",
    "  try:\n",
    "    cuda_version = subprocess.run([\"nvcc\",\"--version\"],stdout=subprocess.PIPE)\n",
    "    cuda_version = cuda_version.stdout.decode('utf-8')\n",
    "    cuda_version = cuda_version[cuda_version.find(', V')+3:-1]\n",
    "  except:\n",
    "    cuda_version = ' - No cuda found - '\n",
    "  try:\n",
    "    gpu_name = subprocess.run([\"nvidia-smi\"],stdout=subprocess.PIPE)\n",
    "    gpu_name = gpu_name.stdout.decode('utf-8')\n",
    "    gpu_name = gpu_name[gpu_name.find('Tesla'):gpu_name.find('Tesla')+10]\n",
    "  except:\n",
    "    gpu_name = ' - No GPU found - '\n",
    "  #print(cuda_version[cuda_version.find(', V')+3:-1])\n",
    "  #print(gpu_name)\n",
    "\n",
    "  shape = io.imread(Training_Source+'/'+os.listdir(Training_Source)[1]).shape\n",
    "  dataset_size = len(os.listdir(Training_Source))\n",
    "\n",
    "  text = 'The '+Network+' model was trained from scratch for '+str(number_of_epochs)+' epochs on '+str(dataset_size)+' labelled images (image dimensions: '+str(shape)+') with a batch size of '+str(batch_size)+' and a custom loss function combining MSE and crossentropy losses, using the '+Network+' ZeroCostDL4Mic notebook (v '+Notebook_version[0]+') (von Chamier & Laine et al., 2020). Key python packages used include tensorflow (v '+version_numbers[0]+'), Keras (v '+version_numbers[2]+'), numpy (v '+version_numbers[1]+'), cuda (v '+cuda_version+'). The training was accelerated using a '+gpu_name+'GPU.'\n",
    "\n",
    "  if pretrained_model:\n",
    "    text = 'The '+Network+' model was trained for '+str(number_of_epochs)+' epochs on '+str(dataset_size)+' labelled images (image dimensions: '+str(shape)+') with a batch size of '+str(batch_size)+' and a custom loss function combining MSE and crossentropy losses, using the '+Network+' ZeroCostDL4Mic notebook (v '+Notebook_version[0]+') (von Chamier & Laine et al., 2020). The model was retrained from a pretrained model. Key python packages used include tensorflow (v '+version_numbers[0]+'), Keras (v '+version_numbers[2]+'), numpy (v '+version_numbers[1]+'), cuda (v '+cuda_version+'). The training was accelerated using a '+gpu_name+'GPU.'\n",
    "\n",
    "  pdf.set_font('')\n",
    "  pdf.set_font_size(10.)\n",
    "  pdf.multi_cell(190, 5, txt = text, align='L')\n",
    "  pdf.ln(1)\n",
    "  pdf.set_font('')\n",
    "  pdf.set_font('Arial', size = 10, style = 'B')\n",
    "  pdf.ln(1)\n",
    "  pdf.cell(28, 5, txt='Augmentation: ', ln=0)\n",
    "  pdf.set_font('')\n",
    "  if augmentation:\n",
    "    aug_text = 'The dataset was augmented by a factor of '+str(multiply_dataset_by)+' by'\n",
    "    if multiply_dataset_by >= 2:\n",
    "      aug_text = aug_text+'\\n- flipping'\n",
    "    if multiply_dataset_by > 2:\n",
    "      aug_text = aug_text+'\\n- rotation'\n",
    "  else:\n",
    "    aug_text = 'No augmentation was used for training.'\n",
    "  pdf.multi_cell(190, 5, txt=aug_text, align='L')\n",
    "  pdf.ln(1)\n",
    "  pdf.set_font('Arial', size = 11, style = 'B')\n",
    "  pdf.ln(1)\n",
    "  pdf.cell(180, 5, txt = 'Parameters', align='L', ln=1)\n",
    "  pdf.set_font('')\n",
    "  pdf.set_font_size(10.)\n",
    "  if Use_Default_Advanced_Parameters:\n",
    "    pdf.cell(200, 5, txt='Default Advanced Parameters were enabled')\n",
    "  pdf.cell(200, 5, txt='The following parameters were used for training:')\n",
    "  pdf.ln(1)\n",
    "  html = \"\"\" \n",
    "  <table width=40% style=\"margin-left:0px;\">\n",
    "    <tr>\n",
    "      <th width = 50% align=\"left\">Parameter</th>\n",
    "      <th width = 50% align=\"left\">Value</th>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td width = 50%>number_of_epochs</td>\n",
    "      <td width = 50%>{0}</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td width = 50%>batch_size</td>\n",
    "      <td width = 50%>{1}</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td width = 50%>learning_rate</td>\n",
    "      <td width = 50%>{2}</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td width = 50%>percentage_validation</td>\n",
    "      <td width = 50%>{3}</td>\n",
    "    </tr>\n",
    "  </table>\n",
    "  \"\"\".format(number_of_epochs, batch_size, learning_rate, percentage_validation)\n",
    "  pdf.write_html(html)\n",
    "\n",
    "  #pdf.multi_cell(190, 5, txt = text_2, align='L')\n",
    "  pdf.set_font(\"Arial\", size = 11, style='B')\n",
    "  pdf.ln(1)\n",
    "  pdf.cell(190, 5, txt = 'Training Dataset', align='L', ln=1)\n",
    "  pdf.set_font('')\n",
    "  pdf.set_font('Arial', size = 10, style = 'B')\n",
    "  pdf.cell(30, 5, txt= 'Training_source:', align = 'L', ln=0)\n",
    "  pdf.set_font('')\n",
    "  pdf.multi_cell(170, 5, txt = Training_Source, align = 'L')\n",
    "  pdf.ln(1)\n",
    "  pdf.set_font('')\n",
    "  pdf.set_font('Arial', size = 10, style = 'B')\n",
    "  pdf.cell(29, 5, txt= 'Training_target:', align = 'L', ln=0)\n",
    "  pdf.set_font('')\n",
    "  pdf.multi_cell(170, 5, txt = Annotations, align = 'L')\n",
    "  #pdf.cell(190, 5, txt=aug_text, align='L', ln=1)\n",
    "  pdf.ln(1)\n",
    "  pdf.set_font('')\n",
    "  pdf.set_font('Arial', size = 10, style = 'B')\n",
    "  pdf.cell(22, 5, txt= 'Model Path:', align = 'L', ln=0)\n",
    "  pdf.set_font('')\n",
    "  pdf.multi_cell(170, 5, txt = model_path+'/'+model_name, align = 'L')\n",
    "  pdf.ln(1)\n",
    "  if visualise_example == True:\n",
    "    pdf.cell(60, 5, txt = 'Example ground-truth annotation', ln=1)\n",
    "    pdf.ln(1)\n",
    "    exp_size = io.imread(base_path + '/TrainingDataExample_RetinaNet.png').shape\n",
    "    pdf.image(base_path + '/TrainingDataExample_RetinaNet.png', x = 11, y = None, w = round(exp_size[1]/8), h = round(exp_size[0]/8))\n",
    "    pdf.ln(1)\n",
    "  ref_1 = 'References:\\n - ZeroCostDL4Mic: von Chamier, Lucas & Laine, Romain, et al. \"ZeroCostDL4Mic: an open platform to simplify access and use of Deep-Learning in Microscopy.\" bioRxiv (2020).'\n",
    "  pdf.multi_cell(190, 5, txt = ref_1, align='L')\n",
    "  pdf.ln(1)\n",
    "  ref_2 = '- Feature Pyramid Networks for Object Detection from Yin Cui, Menglin Jia, Tsung-Yi Lin, Yang Song, Serge Belongie in IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2019.'\n",
    "  pdf.multi_cell(190, 5, txt = ref_2, align='L')\n",
    "  pdf.ln(1)\n",
    "  if augmentation:\n",
    "    ref_4 = '- imgaug: Jung, Alexander et al., https://github.com/aleju/imgaug, (2020)'\n",
    "    pdf.multi_cell(190, 5, txt = ref_4, align='L')\n",
    "  \n",
    "  pdf.ln(2)\n",
    "  pdf.set_font('Arial', size = 10, style = 'B')\n",
    "  developed_by = 'The RetinaNet ZeroCostDL4Mic notebook was developed by:\\n'\n",
    "  pdf.set_font('')\n",
    "\n",
    "  pdf.cell(190, 5, txt = developed_by, ln=1, align='L')\n",
    "  pdf.cell(25, 5, txt = '- Erlantz Calvo (https://github.com/ErlantzCalvo)', align='L', ln=1,  link=\"https://github.com/ErlantzCalvo\")\n",
    "  pdf.cell(25, 5, txt = '- Ignacio Arganda-Carreras (https://sites.google.com/site/iargandacarreras)', align='L', ln=1, link=\"https://sites.google.com/site/iargandacarreras/\")\n",
    "  \n",
    "  pdf.ln(3)\n",
    "  reminder = 'Important:\\nRemember to perform the quality control step on all newly trained models\\nPlease consider depositing your training dataset on Zenodo'\n",
    "  pdf.set_font('Arial', size = 11, style='B')\n",
    "  pdf.multi_cell(190, 5, txt=reminder, align='C')\n",
    "  pdf.ln(1)\n",
    "\n",
    "  pdf.output(model_path+'/'+model_name+'/'+model_name+'_training_report.pdf')\n",
    "  # pdf.output(pdf_export_path+'_training_report.pdf')\n",
    "\n",
    "  print('------------------------------')\n",
    "  print('PDF report exported in '+model_path+'/'+model_name+'/')\n",
    "\n",
    "\n",
    "def qc_pdf_export():\n",
    "  class MyFPDF(FPDF, HTMLMixin):\n",
    "    pass\n",
    "\n",
    "  pdf = MyFPDF()\n",
    "  pdf.add_page()\n",
    "  pdf.set_right_margin(-1)\n",
    "  pdf.set_font(\"Arial\", size = 11, style='B') \n",
    "\n",
    "  Network = 'RetinaNet'\n",
    "\n",
    "  day = datetime.now()\n",
    "  datetime_str = str(day)[0:16]\n",
    "\n",
    "  Header = 'Quality Control report for '+Network+' model ('+QC_model_name+')\\nDate and Time: '+datetime_str\n",
    "  pdf.multi_cell(180, 5, txt = Header, align = 'L') \n",
    "  pdf.ln(1)\n",
    "\n",
    "  all_packages = ''\n",
    "  for requirement in freeze(local_only=True):\n",
    "    all_packages = all_packages+requirement+', '\n",
    "\n",
    "  pdf.set_font('')\n",
    "  pdf.set_font('Arial', size = 11, style = 'B')\n",
    "  pdf.ln(2)\n",
    "  pdf.cell(190, 5, txt = 'Development of Training Losses', ln=1, align='L')\n",
    "  pdf.ln(1)\n",
    "  if os.path.exists(QC_model_folder+'/Quality Control/lossCurveAndmAPPlots.png'):\n",
    "    exp_size = io.imread(QC_model_folder+'/Quality Control/lossCurveAndmAPPlots.png').shape\n",
    "    pdf.image(QC_model_folder+'/Quality Control/lossCurveAndmAPPlots.png', x = 11, y = None, w = round(exp_size[1]/8), h = round(exp_size[0]/8))\n",
    "  else:\n",
    "    pdf.set_font('')\n",
    "    pdf.set_font('Arial', size=10)\n",
    "    pdf.multi_cell(190, 5, txt='If you would like to see the evolution of the loss function during training please play the first cell of the QC section in the notebook.')\n",
    "    pdf.ln(1)\n",
    "  pdf.set_font('')\n",
    "  pdf.set_font('Arial', size = 10, style = 'B')\n",
    "  pdf.cell(80, 5, txt = 'P-R curves for test dataset', ln=1, align='L')\n",
    "  pdf.ln(2)\n",
    "  for i in range(len(AP)):\n",
    "    if os.path.exists(QC_model_folder+f'/Quality Control/P-R_curve_{category_index[i+1][\"name\"]}.png'):\n",
    "      exp_size = io.imread(QC_model_folder+f'/Quality Control/P-R_curve_{category_index[i+1][\"name\"]}.png').shape\n",
    "      pdf.ln(1)\n",
    "      pdf.image(QC_model_folder+f'/Quality Control/P-R_curve_{category_index[i+1][\"name\"]}.png', x=16, y=None, w=round(exp_size[1]/4), h=round(exp_size[0]/4))\n",
    "    else:\n",
    "      pdf.cell(100, 5, txt=f'For the class {category_index[i+1][\"name\"]} the model did not predict any objects.', ln=1, align='L')\n",
    "  pdf.ln(3)\n",
    "  pdf.set_font('')\n",
    "  pdf.set_font('Arial', size = 11, style = 'B')\n",
    "  pdf.ln(1)\n",
    "  pdf.cell(180, 5, txt = 'Quality Control Metrics', align='L', ln=1)\n",
    "  pdf.set_font('')\n",
    "  pdf.set_font_size(10.)\n",
    "\n",
    "  pdf.ln(1)\n",
    "  html = \"\"\"\n",
    "  <body>\n",
    "  <font size=\"10\" face=\"Courier\" >\n",
    "  <table width=95% style=\"margin-left:0px;\">\"\"\"\n",
    "  with open(QC_model_folder+'/Quality Control/QC_results.csv', 'r') as csvfile:\n",
    "    metrics = csv.reader(csvfile)\n",
    "    header = next(metrics)\n",
    "    class_name = header[0]\n",
    "    fp = header[1]\n",
    "    tp = header[2]\n",
    "    fn = header[3]\n",
    "    recall = header[4]\n",
    "    precision = header[5]\n",
    "    acc = header[6]\n",
    "    f1 = header[7]\n",
    "    AP_score = header[8]\n",
    "    header = \"\"\"\n",
    "    <tr>\n",
    "    <th width = 11% align=\"left\">{0}</th>\n",
    "    <th width = 13% align=\"left\">{1}</th>\n",
    "    <th width = 13% align=\"left\">{2}</th>\n",
    "    <th width = 13% align=\"left\">{3}</th>\n",
    "    <th width = 8% align=\"left\">{4}</th>\n",
    "    <th width = 9% align=\"left\">{5}</th>\n",
    "    <th width = 9% align=\"left\">{6}</th>\n",
    "    <th width = 9% align=\"left\">{7}</th>\n",
    "    <th width = 14% align=\"left\">{8}</th>\n",
    "    </tr>\"\"\".format(class_name,fp,tp,fn,recall,precision,acc,f1,AP_score)\n",
    "    html = html+header\n",
    "    i=0\n",
    "    for row in metrics:\n",
    "      i+=1\n",
    "      class_name = row[0]\n",
    "      fp = row[1]\n",
    "      tp = row[2]\n",
    "      fn = row[3]\n",
    "      recall = row[4]\n",
    "      precision = row[5]\n",
    "      acc = row[6]\n",
    "      f1 = row[7]\n",
    "      AP_score = row[8]\n",
    "      cells = \"\"\"\n",
    "        <tr>\n",
    "          <td width = 11% align=\"left\">{0}</td>\n",
    "          <td width = 13% align=\"left\">{1}</td>\n",
    "          <td width = 13% align=\"left\">{2}</td>\n",
    "          <td width = 13% align=\"left\">{3}</td>\n",
    "          <td width = 8% align=\"left\">{4}</td>\n",
    "          <td width = 9% align=\"left\">{5}</td>\n",
    "          <td width = 9% align=\"left\">{6}</td>\n",
    "          <td width = 9% align=\"left\">{7}</td>\n",
    "          <td width = 14% align=\"left\">{8}</td>\n",
    "        </tr>\"\"\".format(class_name,fp,tp,fn,str(round(float(recall),3)),str(round(float(precision),3)),str(round(float(acc),3)),str(round(float(f1),3)),str(round(float(AP_score),3)))\n",
    "      html = html+cells\n",
    "    html = html+\"\"\"</body></table>\"\"\"\n",
    "\n",
    "  pdf.write_html(html)\n",
    "  pdf.cell(180, 5, txt='Mean average precision (mAP) over the all classes is: '+str(round(mAP_score,3)), ln=1, align='L')\n",
    "  pdf.set_font('')\n",
    "  pdf.set_font('Arial', size = 11, style = 'B')\n",
    "  pdf.ln(3)\n",
    "  pdf.cell(80, 5, txt = 'Example Quality Control Visualisation', ln=1)\n",
    "  pdf.ln(3)\n",
    "  exp_size = io.imread(QC_model_folder+'/Quality Control/QC_example_data.png').shape\n",
    "  pdf.image(QC_model_folder+'/Quality Control/QC_example_data.png', x = 16, y = None, w = round(exp_size[1]/10), h = round(exp_size[0]/10))\n",
    "  \n",
    "  pdf.ln(2)\n",
    "  pdf.set_font('Arial', size = 10, style = 'B')\n",
    "  developed_by = 'The RetinaNet ZeroCostDL4Mic notebook was developed by:\\n'\n",
    "  pdf.set_font('')\n",
    "\n",
    "  pdf.cell(190, 5, txt = developed_by, ln=1, align='L')\n",
    "  pdf.cell(25, 5, txt = '- Erlantz Calvo (https://github.com/ErlantzCalvo)', align='L', ln=1,  link=\"https://github.com/ErlantzCalvo\")\n",
    "  pdf.cell(25, 5, txt = '- Ignacio Arganda-Carreras (https://sites.google.com/site/iargandacarreras)', align='L', ln=1, link=\"https://sites.google.com/site/iargandacarreras/\")\n",
    "  \n",
    "  pdf.set_font('')\n",
    "  pdf.set_font_size(10.)\n",
    "  pdf.ln(3)\n",
    "  ref_1 = 'References:\\n - ZeroCostDL4Mic: von Chamier, Lucas & Laine, Romain, et al. \"ZeroCostDL4Mic: an open platform to simplify access and use of Deep-Learning in Microscopy.\" bioRxiv (2020).'\n",
    "  pdf.multi_cell(190, 5, txt = ref_1, align='L')\n",
    "  pdf.ln(1)\n",
    "  ref_2 = '- Feature Pyramid Networks for Object Detection from Yin Cui, Menglin Jia, Tsung-Yi Lin, Yang Song, Serge Belongie in IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2019'\n",
    "  pdf.multi_cell(190, 5, txt = ref_2, align='L')\n",
    "  pdf.ln(1)\n",
    "\n",
    "\n",
    "  pdf.ln(3)\n",
    "  reminder = 'To find the parameters and other information about how this model was trained, go to the training_report.pdf of this model which should be in the folder of the same name.'\n",
    "\n",
    "  pdf.set_font('Arial', size = 11, style='B')\n",
    "  pdf.multi_cell(190, 5, txt=reminder, align='C')\n",
    "  pdf.ln(1)\n",
    "\n",
    "  pdf.output(QC_model_folder+'/Quality Control/'+QC_model_name+'_QC_report.pdf')\n",
    "\n",
    "\n",
    "  print('------------------------------')\n",
    "  print('PDF report exported in '+QC_model_folder+'/Quality Control/')\n",
    "# Exporting requirements.txt for local run\n",
    "# !pip freeze > requirements.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "n4yWFoJNnoin"
   },
   "source": [
    "# **2. Initialise the Colab session**\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "---\n",
    "\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DMNHVZfHmbKb"
   },
   "source": [
    "\n",
    "## **2.1. Check for GPU access**\n",
    "---\n",
    "\n",
    "By default, the session should be using Python 3 and GPU acceleration, but it is possible to ensure that these are set properly by doing the following:\n",
    "\n",
    "<font size = 4>Go to **Runtime -> Change the Runtime type**\n",
    "\n",
    "<font size = 4>**Runtime type: Python 3** *(Python 3 is programming language in which this program is written)*\n",
    "\n",
    "<font size = 4>**Accelator: GPU** *(Graphics processing unit)*\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "zCvebubeSaGY"
   },
   "outputs": [],
   "source": [
    "#@markdown ##Run this cell to check if you have GPU access\n",
    "import tensorflow as tf\n",
    "if tf.test.gpu_device_name()=='':\n",
    "  print('You do not have GPU access.') \n",
    "  print('Did you change your runtime ?') \n",
    "  print('If the runtime setting is correct then Google did not allocate a GPU for your session')\n",
    "  print('Expect slow performance. To access GPU try reconnecting later')\n",
    "\n",
    "else:\n",
    "  print('You have GPU access')\n",
    "  !nvidia-smi"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "sNIVx8_CLolt"
   },
   "source": [
    "## **2.2. Mount your Google Drive**\n",
    "---\n",
    "<font size = 4> To use this notebook on the data present in your Google Drive, you need to mount your Google Drive to this notebook.\n",
    "\n",
    "<font size = 4> Play the cell below to mount your Google Drive and follow the link. In the new browser window, select your drive and select 'Allow', copy the code, paste into the cell and press enter. This will give Colab access to the data on the drive. \n",
    "\n",
    "<font size = 4> Once this is done, your data are available in the **Files** tab on the top left of notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "01Djr8v-5pPk"
   },
   "outputs": [],
   "source": [
    "#@markdown ##Play the cell to connect your Google Drive to Colab\n",
    "\n",
    "#@markdown * Click on the URL. \n",
    "\n",
    "#@markdown * Sign in your Google Account. \n",
    "\n",
    "#@markdown * Copy the authorization code. \n",
    "\n",
    "#@markdown * Enter the authorization code. \n",
    "\n",
    "#@markdown * Click on \"Files\" site on the right. Refresh the site. Your Google Drive folder should now be available here as \"drive\". \n",
    "\n",
    "# mount user's Google Drive to Google Colab.\n",
    "from google.colab import drive\n",
    "drive.mount(base_path + '/gdrive')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "HLYcZR9gMv42"
   },
   "source": [
    "# **3. Select your paths and parameters**\n",
    "\n",
    "---\n",
    "\n",
    "<font size = 4>The code below allows the user to enter the paths to where the training data is and to define the training parameters.\n",
    "\n",
    "<font size = 4>After playing the cell will display some quantitative metrics of your dataset, including a count of objects per image and the number of instances per class.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "RGqhBI4JEnFe"
   },
   "source": [
    "## In case you want to try this notebook with one of our datasets (optional)\n",
    "An already annotated dataset will be downloaded locally. This means that it won't be saved in your Google Drive and, once the execution is stopped, it will be lost until the next time it's downloaded. <a name=\"example_dataset\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "KdfB8FMTEmt4"
   },
   "outputs": [],
   "source": [
    "#@markdown ###Download an already annotated example dataset.\n",
    "!wget -nv --show-progress -A.zip https://zenodo.org/record/3941908//files/YoloV2%20dataset_v1.zip -O base_path + \"/example_dataset.zip\" \n",
    "!unzip -qq -o base_path + \"/example_dataset.zip\" -d base_path + \"/example_dataset\"\n",
    "!rm -rf \"example_dataset.zip\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "FQ_QxtSWQ7CL"
   },
   "source": [
    "## **3.1. Setting main training parameters**\n",
    "---\n",
    "<font size = 4> "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "AuESFimvMv43"
   },
   "source": [
    "<font size = 5> **Paths for training, predictions and results**\n",
    "\n",
    "\n",
    "* <font size = 4>**`Training_source:`, `Annotations`:** These are the paths to your folders containing the Training_source and the annotation data respectively. To find the paths of the folders containing the respective datasets, go to your Files on the left of the notebook, navigate to the folder containing your files and copy the path by right-clicking on the folder, **Copy path** and pasting it into the right box below. **Note**: If you are using the [example dataset](#example_dataset) downloaded in the previous cell, the **`Training_source`** field would be *`base_folder + /example_dataset/Training_Images`* and the **`Annotations`** field would be *`base_folder + /example_dataset/Training_Annotations`*.\n",
    "\n",
    "* <font size = 4>**`model_name`:** Enter here a name for your model (e.g.: `RetinaNet`). Use only my_model -style, not my-model (Use \"_\" not \"-\"). Do not use spaces in the name. Avoid using the name of an existing model (saved in the same folder) as it will be overwritten.\n",
    "\n",
    "* <font size = 4>**`model_path`:** Enter the path where your model will be saved once trained (for instance your result folder). For instance, for saving the model trained in the root folder of the Google Drive account chose in the section 1 this field would be *base_folder + /gdrive/MyDrive*.  \n",
    "\n",
    "<font size = 5>**Training Parameters**\n",
    "\n",
    "* <font size = 4>**`number_of_epochs`:** Give estimates for training performance given a number of epochs and provide a default value. **Default value: 50**\n",
    "\n",
    "\n",
    "<font size = 5>**Advanced Parameters - experienced users only**\n",
    "\n",
    "* <font size = 4>**`batch_size:`** This parameter defines the number of patches seen in each training step. Reducing or increasing the **batch size** may slow or speed up your training, respectively, and can influence network performance. **Default value: 8**\n",
    "\n",
    "* <font size = 4>**`learning_rate:`** Input the initial value to be used as learning rate. **Default value: 0.0001**\n",
    "\n",
    "\n",
    "* <font size = 4>**`percentage_validation:`**  Input the percentage of your training dataset you want to use to validate the network during training. **Default value: 10** "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "ewpNJ_I0Mv47"
   },
   "outputs": [],
   "source": [
    "#@markdown ###Path to training images:\n",
    "%matplotlib inline\n",
    "Training_Source = \"\" #@param {type:\"string\"}\n",
    "\n",
    "# Ground truth images\n",
    "Annotations = \"\" #@param {type:\"string\"}\n",
    "\n",
    "# sanity check of the path in order to avoid possible errors\n",
    "if Training_Source[-1] == '/':\n",
    "  Training_Source = Training_Source[:-1]\n",
    "\n",
    "if Annotations[-1] == '/':\n",
    "  Annotations = Annotations[:-1]\n",
    "# model name and path\n",
    "#@markdown ###Name of the model and path to model folder:\n",
    "model_name = \"\" #@param {type:\"string\"}\n",
    "model_path = \"\" #@param {type:\"string\"}\n",
    "\n",
    "if model_path[-1] == '/':\n",
    "  model_path = model_path[:-1]\n",
    "\n",
    "full_model_path = os.path.join(model_path,model_name)\n",
    "if os.path.exists(full_model_path):\n",
    "  print(bcolors.WARNING+'Model folder already exists and will be overwritten.'+bcolors.NORMAL)\n",
    "  shutil.rmtree(full_model_path)\n",
    "os.mkdir(full_model_path)\n",
    "\n",
    "\n",
    "# other parameters for training.\n",
    "# @markdown ###Training Parameters\n",
    "# @markdown Number of epochs:\n",
    "\n",
    "number_of_epochs =  100#@param {type:\"number\"}\n",
    "number_of_epochs = int(number_of_epochs) # prevent the user of inserting float number\n",
    "\n",
    "# #@markdown ###Advanced Parameters\n",
    "\n",
    "Use_Default_Advanced_Parameters = True #@param {type:\"boolean\"}\n",
    "\n",
    "#@markdown ###If not, please input:\n",
    "batch_size =  8#@param {type:\"number\"}\n",
    "learning_rate = 1e-4 #@param{type:\"number\"}\n",
    "percentage_validation =  10#@param{type:\"number\"}\n",
    "\n",
    "batch_size = int(batch_size)\n",
    "if (Use_Default_Advanced_Parameters): \n",
    "  print(\"Default advanced parameters enabled\")\n",
    "  batch_size = 8\n",
    "  learning_rate = 1e-4\n",
    "  percentage_validation = 10\n",
    "\n",
    "Validation_Source= os.path.join(os.path.dirname(Training_Source), 'Validation_Images/')\n",
    "Validation_Source_annotations= os.path.join(os.path.dirname(Training_Source), 'Validation_Annotations/')\n",
    "\n",
    "prepare_validation_folder()\n",
    "dir_anno = Annotations\n",
    "training_images = os.listdir(dir_anno)\n",
    "train_images_number = len(training_images)\n",
    "validation_number = int(train_images_number * percentage_validation/100)\n",
    "training_images = os.listdir(dir_anno)\n",
    "\n",
    "training_set = train_images_number - validation_number\n",
    "\n",
    "\n",
    "random_val_images = random.sample(range(0,train_images_number-1), validation_number)\n",
    "df_anno, df_anno_valid = [], []\n",
    "index = 0\n",
    "\n",
    "for fnm in training_images:  \n",
    "    if not fnm.startswith('.'): ## do not include hidden folders/files\n",
    "        tree = ET.parse(os.path.join(dir_anno,fnm))\n",
    "        row = extract_single_xml_file(tree)\n",
    "        row[\"fileID\"] = os.path.splitext(fnm)[0]\n",
    "        if index in random_val_images:\n",
    "          df_anno_valid.append(row)\n",
    "          move_to_validation( fnm.split('.')[0])\n",
    "        else:\n",
    "          df_anno.append(row)\n",
    "\n",
    "        index +=1\n",
    "\n",
    "df_anno = pd.DataFrame(df_anno)\n",
    "df_anno_valid = pd.DataFrame(df_anno_valid)\n",
    "\n",
    "maxNobj = np.max(df_anno[\"Nobj\"])\n",
    "if len(df_anno_valid) > 0:\n",
    "  maxNobjVal = np.max(df_anno_valid[\"Nobj\"])\n",
    "totalNobj = np.sum(df_anno[\"Nobj\"])\n",
    "\n",
    "###### RetinaNet's data preprocess method\n",
    "gt_boxes = []\n",
    "gt_classes = []\n",
    "for index, row in df_anno.iterrows():\n",
    "  n_objects = row['Nobj']\n",
    "  objects_coords = np.zeros((n_objects,4))\n",
    "  objects_classes = np.zeros((n_objects,1), dtype=object)\n",
    "  for i in range(n_objects):\n",
    "    current_bbx = f'bbx_{i}'\n",
    "    objects_coords[i,0] = row[current_bbx + '_ymin'] / row['height']\n",
    "    objects_coords[i,1] = row[current_bbx + '_xmin'] / row['width']\n",
    "    objects_coords[i,2] = row[current_bbx + '_ymax'] / row['height']\n",
    "    objects_coords[i,3] = row[current_bbx + '_xmax'] / row['width']\n",
    "\n",
    "    objects_classes[i,0] = row[current_bbx + '_name']\n",
    "\n",
    "  gt_boxes.append(objects_coords)\n",
    "  gt_classes.append(objects_classes)\n",
    "\n",
    "gt_boxes_valid = []\n",
    "gt_classes_valid = []\n",
    "for index, row in df_anno_valid.iterrows():\n",
    "  n_objects = row['Nobj']\n",
    "  objects_coords = np.zeros((n_objects,4))\n",
    "  objects_classes = np.zeros((n_objects,1), dtype=object)\n",
    "  for i in range(n_objects):\n",
    "    current_bbx = f'bbx_{i}'\n",
    "    objects_coords[i,0] = row[current_bbx + '_ymin'] / row['height']\n",
    "    objects_coords[i,1] = row[current_bbx + '_xmin'] / row['width']\n",
    "    objects_coords[i,2] = row[current_bbx + '_ymax'] / row['height']\n",
    "    objects_coords[i,3] = row[current_bbx + '_xmax'] / row['width']\n",
    "\n",
    "    objects_classes[i,0] = row[current_bbx + '_name']\n",
    "\n",
    "  gt_boxes_valid.append(objects_coords)\n",
    "  gt_classes_valid.append(objects_classes)\n",
    "\n",
    "\n",
    "\n",
    "class_obj = []\n",
    "for ibbx in range(maxNobj):\n",
    "    class_obj.extend(df_anno[\"bbx_{}_name\".format(ibbx)].values)\n",
    "class_obj = np.array(class_obj)\n",
    "\n",
    "count             = Counter(class_obj[class_obj != 'nan'])\n",
    "print(count)\n",
    "class_nm          = list(count.keys())\n",
    "class_labels = json.dumps(class_nm)\n",
    "class_count       = list(count.values())\n",
    "asort_class_count = np.argsort(class_count)\n",
    "\n",
    "class_nm          = np.array(class_nm)[asort_class_count]\n",
    "class_count       = np.array(class_count)[asort_class_count]\n",
    "\n",
    "xs = range(len(class_count))\n",
    "\n",
    "# Creation of the category index (classes dictionary) sorted by objects number in descending order\n",
    "classes_dict = dict()\n",
    "category_index = dict()\n",
    "\n",
    "index = 0\n",
    "for key in sorted(count, key=count.get, reverse=True):\n",
    "  category_index[index+1] = {'id': index+1, 'name': key}\n",
    "  classes_dict[key] = index + 1\n",
    "  index +=1\n",
    "\n",
    "\n",
    "#Show how many objects there are in the images\n",
    "plt.figure(figsize=(15,8))\n",
    "plt.subplot(1,2,1)\n",
    "plt.hist(df_anno[\"Nobj\"].values,bins=50)\n",
    "plt.title(\"Total number of objects in the dataset: {}\".format(totalNobj))\n",
    "plt.xlabel('Number of objects per image')\n",
    "plt.ylabel('Occurences')\n",
    "\n",
    "plt.subplot(1,2,2)\n",
    "plt.barh(xs,class_count)\n",
    "plt.yticks(xs,class_nm)\n",
    "plt.title(\"The number of objects per class: {} classes in total\".format(len(count)))\n",
    "plt.show()\n",
    "\n",
    "visualise_example = False\n",
    "Use_pretrained_model = True\n",
    "Use_Data_augmentation = True\n",
    "\n",
    "\n",
    "# full_model_path = os.path.join(model_path,model_name)\n",
    "# pdf_export_path = os.path.join(full_model_path,'training_report')\n",
    "\n",
    "# Create a new directory\n",
    "\n",
    "# visualise an example image from your dataset to make sure annotations and images are properly matched.\n",
    "print()\n",
    "print('Now an image from the training set will be shown with its annotations in order to visualise that the annotations are properly matched.')\n",
    "print()\n",
    "\n",
    "size = 1 \n",
    "ind_random = np.random.randint(0,df_anno.shape[0],size=size)\n",
    "img_dir=Training_Source\n",
    "\n",
    "file_suffix = os.path.splitext(os.listdir(Training_Source)[0])[1]\n",
    "for irow in ind_random:\n",
    "    row  = df_anno.iloc[irow,:]\n",
    "    path = os.path.join(img_dir, row[\"fileID\"] + file_suffix)\n",
    "    \n",
    "    predicted_classes = np.zeros(shape=[gt_boxes[irow].shape[0]], dtype=np.int32)\n",
    "\n",
    "    # get classes of each prediction\n",
    "    for i in range(gt_classes[irow].shape[0]):\n",
    "      predicted_classes[i] = classes_dict[gt_classes[irow][i,0]]\n",
    "    \n",
    "    # read in image and convert to rgb (3 channels)\n",
    "    img  = np.array(Image.open(path))\n",
    "    img = cv2.cvtColor(img,cv2.COLOR_GRAY2RGB)\n",
    "    dummy_scores = np.ones(gt_boxes[irow].shape[0], dtype=np.float32)  # give boxes a score of 100%\n",
    "\n",
    "    plot_detections(\n",
    "      img,\n",
    "      gt_boxes[irow],\n",
    "      predicted_classes,\n",
    "      dummy_scores, category_index)\n",
    "\n",
    "\n",
    "pdf_export()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "xyQZKby8yFME"
   },
   "source": [
    "## **3.2. Data augmentation**\n",
    "\n",
    "---\n",
    "\n",
    "<font size = 4> Data augmentation can improve training progress by amplifying differences in the dataset. This can be useful if the available dataset is small since, in this case, it is possible that a network could quickly learn every example in the dataset (overfitting), without augmentation. Augmentation is not necessary for training and if the dataset the `Use_Data_Augmentation` box can be unticked.\n",
    "\n",
    "<font size = 4>Here, the images and bounding boxes are augmented by flipping and rotation. When doubling the dataset the images are only flipped. With each higher factor of augmentation the images added to the dataset represent one further rotation to the right by 90 degrees. 8x augmentation will give a dataset that is fully rotated and flipped once."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "DMqWq5-AxnFU"
   },
   "outputs": [],
   "source": [
    "#@markdown ##**Augmentation Options**\n",
    "%matplotlib inline\n",
    "\n",
    "Use_Data_augmentation = True #@param {type:\"boolean\"}\n",
    "\n",
    "multiply_dataset_by = 2 #@param {type:\"slider\", min:2, max:8, step:1}\n",
    "\n",
    "rotation_range = 90\n",
    "\n",
    "\n",
    "file_suffix = os.path.splitext(os.listdir(Training_Source)[0])[1]\n",
    "if (Use_Data_augmentation):\n",
    "  print('Data Augmentation enabled')\n",
    "  # load images as NumPy arrays and append them to images list\n",
    "  if os.path.exists(Training_Source+'/.ipynb_checkpoints'):\n",
    "    shutil.rmtree(Training_Source+'/.ipynb_checkpoints')\n",
    "  \n",
    "  images = []\n",
    "  for index, file in enumerate(glob.glob(Training_Source+'/*'+file_suffix)):\n",
    "      images.append(imageio.imread(file))\n",
    "     \n",
    "  # how many images we have\n",
    "  print('Augmenting {} images'.format(len(images)))\n",
    "\n",
    "  # apply xml_to_csv() function to convert all XML files in images/ folder into labels.csv\n",
    "  labels_df = xml_to_csv(Annotations)\n",
    "  labels_df.to_csv((base_path + '/original_labels.csv'), index=None)\n",
    "  \n",
    "  # Apply flip augmentation\n",
    "  aug = iaa.OneOf([    \n",
    "      iaa.Fliplr(1),\n",
    "      iaa.Flipud(1)\n",
    "  ])\n",
    "  aug_2 = iaa.Affine(rotate=rotation_range, fit_output=True)\n",
    "  aug_3 = iaa.Affine(rotate=rotation_range*2, fit_output=True)\n",
    "  aug_4 = iaa.Affine(rotate=rotation_range*3, fit_output=True)\n",
    "\n",
    "  #Here we create a folder that will hold the original image dataset and the augmented image dataset\n",
    "  augmented_training_source = os.path.dirname(Training_Source)+'/'+os.path.basename(Training_Source)+'_augmentation'\n",
    "  validation_source = os.path.dirname(Training_Source)+'/'+os.path.basename(Training_Source)+'_validation'\n",
    "  if os.path.exists(augmented_training_source):\n",
    "    shutil.rmtree(augmented_training_source)\n",
    "  os.mkdir(augmented_training_source)\n",
    "\n",
    "  if os.path.exists(validation_source):\n",
    "    shutil.rmtree(validation_source)\n",
    "  os.mkdir(validation_source)\n",
    "\n",
    "  #Here we create a folder that will hold the original image annotation dataset and the augmented image annotation dataset (the bounding boxes).\n",
    "  augmented_training_source_annotation = os.path.dirname(Annotations)+'/'+os.path.basename(Annotations)+'_augmentation'\n",
    "  validation_source_annotation = os.path.dirname(Annotations)+'/'+os.path.basename(Annotations)+'_augmentation'\n",
    "  if os.path.exists(augmented_training_source_annotation):\n",
    "    shutil.rmtree(augmented_training_source_annotation)\n",
    "  os.mkdir(augmented_training_source_annotation)\n",
    "\n",
    "\n",
    "  #Create the augmentation\n",
    "  augmented_images_df = image_aug(labels_df, Training_Source+'/', augmented_training_source+'/', 'flip_', aug)\n",
    "  \n",
    "  # Concat resized_images_df and augmented_images_df together and save in a new all_labels.csv file\n",
    "  all_labels_df = pd.concat([labels_df, augmented_images_df])\n",
    "  all_labels_df.to_csv(base_path + '/combined_labels.csv', index=False)\n",
    "\n",
    "  #Here we convert the new bounding boxes for the augmented images to PASCAL VOC .xml format\n",
    "  def convert_to_xml(df,source,target_folder):\n",
    "    grouped = df.groupby('filename')\n",
    "    for file in os.listdir(source):\n",
    "      #if file in grouped.filename:\n",
    "      group_df = grouped.get_group(file)\n",
    "      group_df = group_df.reset_index()\n",
    "      group_df = group_df.drop(['index'], axis=1)\n",
    "      #group_df = group_df.dropna(axis=0)\n",
    "      writer = Writer(source+'/'+file,group_df.iloc[1]['width'],group_df.iloc[1]['height'])\n",
    "      for i, row in group_df.iterrows():\n",
    "        writer.addObject(row['class'],round(row['xmin']),round(row['ymin']),round(row['xmax']),round(row['ymax']))\n",
    "        writer.save(target_folder+'/'+os.path.splitext(file)[0]+'.xml')\n",
    "  convert_to_xml(all_labels_df,augmented_training_source,augmented_training_source_annotation)\n",
    "  \n",
    "  #Second round of augmentation\n",
    "  if multiply_dataset_by > 2:\n",
    "    aug_labels_df_2 = xml_to_csv(augmented_training_source_annotation)\n",
    "    augmented_images_2_df = image_aug(aug_labels_df_2, augmented_training_source+'/', augmented_training_source+'/', 'rot1_90_', aug_2)\n",
    "    all_aug_labels_df = pd.concat([augmented_images_df, augmented_images_2_df])\n",
    "  \n",
    "    for file in os.listdir(augmented_training_source_annotation):\n",
    "      os.remove(os.path.join(augmented_training_source_annotation,file))\n",
    "    convert_to_xml(all_aug_labels_df,augmented_training_source,augmented_training_source_annotation)\n",
    "\n",
    "  if multiply_dataset_by > 3:\n",
    "    print('Augmenting again')\n",
    "    aug_labels_df_3 = xml_to_csv(augmented_training_source_annotation)\n",
    "    augmented_images_3_df = image_aug(aug_labels_df_3, augmented_training_source+'/', augmented_training_source+'/', 'rot2_90_', aug_2)\n",
    "    all_aug_labels_df_3 = pd.concat([all_aug_labels_df, augmented_images_3_df])\n",
    "\n",
    "    for file in os.listdir(augmented_training_source_annotation):\n",
    "      os.remove(os.path.join(augmented_training_source_annotation,file))\n",
    "    convert_to_xml(all_aug_labels_df_3,augmented_training_source,augmented_training_source_annotation)\n",
    "  \n",
    "  #This is a preliminary remover of potential duplicates in the augmentation\n",
    "  #Ideally, duplicates are not even produced, but this acts as a fail safe.\n",
    "    if multiply_dataset_by==4:\n",
    "      for file in os.listdir(augmented_training_source):\n",
    "        if file.startswith('rot2_90_flip_'):\n",
    "          os.remove(os.path.join(augmented_training_source,file))\n",
    "          os.remove(os.path.join(augmented_training_source_annotation, os.path.splitext(file)[0]+'.xml'))\n",
    "\n",
    "  if multiply_dataset_by > 4:\n",
    "    print('And Again')\n",
    "    aug_labels_df_4 = xml_to_csv(augmented_training_source_annotation)\n",
    "    augmented_images_4_df = image_aug(aug_labels_df_4, augmented_training_source+'/',augmented_training_source+'/','rot3_90_', aug_2)\n",
    "    all_aug_labels_df_4 = pd.concat([all_aug_labels_df_3, augmented_images_4_df])\n",
    "\n",
    "    for file in os.listdir(augmented_training_source_annotation):\n",
    "      os.remove(os.path.join(augmented_training_source_annotation,file))\n",
    "    convert_to_xml(all_aug_labels_df_4,augmented_training_source,augmented_training_source_annotation)\n",
    "\n",
    "    for file in os.listdir(augmented_training_source):\n",
    "      if file.startswith('rot3_90_rot2_90_flip_'):\n",
    "        os.remove(os.path.join(augmented_training_source,file))\n",
    "        os.remove(os.path.join(augmented_training_source_annotation, os.path.splitext(file)[0]+'.xml'))\n",
    "      if file.startswith('rot3_90_rot1_90_flip_'):\n",
    "        os.remove(os.path.join(augmented_training_source,file))\n",
    "        os.remove(os.path.join(augmented_training_source_annotation, os.path.splitext(file)[0]+'.xml'))\n",
    "      if file.startswith('rot3_90_flip_'):\n",
    "        os.remove(os.path.join(augmented_training_source,file))\n",
    "        os.remove(os.path.join(augmented_training_source_annotation, os.path.splitext(file)[0]+'.xml'))\n",
    "      if file.startswith('rot2_90_flip_'):\n",
    "        os.remove(os.path.join(augmented_training_source,file))\n",
    "        os.remove(os.path.join(augmented_training_source_annotation, os.path.splitext(file)[0]+'.xml'))\n",
    "\n",
    "\n",
    "  if multiply_dataset_by > 5:\n",
    "    print('And again')\n",
    "    augmented_images_5_df = image_aug(labels_df, Training_Source+'/', augmented_training_source+'/', 'rot_90_', aug_2)\n",
    "    all_aug_labels_df_5 = pd.concat([all_aug_labels_df_4,augmented_images_5_df])\n",
    "\n",
    "    for file in os.listdir(augmented_training_source_annotation):\n",
    "     os.remove(os.path.join(augmented_training_source_annotation,file))\n",
    "    \n",
    "    convert_to_xml(all_aug_labels_df_5,augmented_training_source,augmented_training_source_annotation)\n",
    "\n",
    "  if multiply_dataset_by > 6:\n",
    "    print('And again')\n",
    "    augmented_images_df_6 = image_aug(labels_df, Training_Source+'/', augmented_training_source+'/', 'rot_180_', aug_3)\n",
    "    all_aug_labels_df_6 = pd.concat([all_aug_labels_df_5,augmented_images_df_6])\n",
    "    \n",
    "    for file in os.listdir(augmented_training_source_annotation):\n",
    "      os.remove(os.path.join(augmented_training_source_annotation,file))\n",
    "    convert_to_xml(all_aug_labels_df_6,augmented_training_source,augmented_training_source_annotation)\n",
    "\n",
    "  if multiply_dataset_by > 7:\n",
    "    print('And again')\n",
    "    augmented_images_df_7 = image_aug(labels_df, Training_Source+'/', augmented_training_source+'/', 'rot_270_', aug_4)\n",
    "    all_aug_labels_df_7 = pd.concat([all_aug_labels_df_6,augmented_images_df_7])\n",
    "    \n",
    "    for file in os.listdir(augmented_training_source_annotation):\n",
    "      os.remove(os.path.join(augmented_training_source_annotation,file))\n",
    "    convert_to_xml(all_aug_labels_df_7,augmented_training_source,augmented_training_source_annotation)\n",
    "\n",
    "  for file in os.listdir(Training_Source):\n",
    "    shutil.copyfile(Training_Source+'/'+file,augmented_training_source+'/'+file)\n",
    "    shutil.copyfile(Annotations+'/'+os.path.splitext(file)[0]+'.xml',augmented_training_source_annotation+'/'+os.path.splitext(file)[0]+'.xml')\n",
    "  # display new dataframe\n",
    "\n",
    "  df_anno = []\n",
    "  dir_anno = augmented_training_source_annotation\n",
    "  for fnm in os.listdir(dir_anno):  \n",
    "      if not fnm.startswith('.'): ## do not include hidden folders/files\n",
    "          tree = ET.parse(os.path.join(dir_anno,fnm))\n",
    "          row = extract_single_xml_file(tree)\n",
    "          row[\"fileID\"] = os.path.splitext(fnm)[0]\n",
    "          df_anno.append(row)\n",
    "  df_anno = pd.DataFrame(df_anno)\n",
    "\n",
    "  maxNobj = np.max(df_anno[\"Nobj\"])\n",
    "\n",
    "  gt_boxes = []\n",
    "  gt_classes = []\n",
    "  for index, row in df_anno.iterrows():\n",
    "    n_objects = row['Nobj']\n",
    "    objects_coords = np.zeros((n_objects,4))\n",
    "    objects_classes = np.zeros((n_objects,1), dtype=object)\n",
    "    for i in range(n_objects):\n",
    "      current_bbx = f'bbx_{i}'\n",
    "      objects_coords[i,0] = row[current_bbx + '_ymin'] / row['height']\n",
    "      objects_coords[i,1] = row[current_bbx + '_xmin'] / row['width']\n",
    "      objects_coords[i,2] = row[current_bbx + '_ymax'] / row['height']\n",
    "      objects_coords[i,3] = row[current_bbx + '_xmax'] / row['width']\n",
    "\n",
    "      objects_classes[i,0] = row[current_bbx + '_name']\n",
    "\n",
    "    gt_boxes.append(objects_coords)\n",
    "    gt_classes.append(objects_classes)\n",
    "\n",
    "  #Write the annotations to a csv file\n",
    "  df_anno.to_csv(model_path+'/annot.csv', index=False)#header=False, sep=',')\n",
    "\n",
    "  #Show how many objects there are in the images\n",
    "  plt.figure()\n",
    "  plt.subplot(2,1,1)\n",
    "  plt.hist(df_anno[\"Nobj\"].values,bins=50)\n",
    "  plt.title(\"max N of objects per image={}\".format(maxNobj))\n",
    "  plt.show()\n",
    "\n",
    "  #Show the classes and how many there are of each in the dataset\n",
    "  class_obj = []\n",
    "  for ibbx in range(maxNobj):\n",
    "      class_obj.extend(df_anno[\"bbx_{}_name\".format(ibbx)].values)\n",
    "  class_obj = np.array(class_obj)\n",
    "\n",
    "  count             = Counter(class_obj[class_obj != 'nan'])\n",
    "  print(count)\n",
    "  class_nm          = list(count.keys())\n",
    "  class_labels = json.dumps(class_nm)\n",
    "  class_count       = list(count.values())\n",
    "  asort_class_count = np.argsort(class_count)\n",
    "\n",
    "  class_nm          = np.array(class_nm)[asort_class_count]\n",
    "  class_count       = np.array(class_count)[asort_class_count]\n",
    "\n",
    "  xs = range(len(class_count))\n",
    "\n",
    "  plt.subplot(2,1,2)\n",
    "  plt.barh(xs,class_count)\n",
    "  plt.yticks(xs,class_nm)\n",
    "  plt.title(\"The number of objects per class: {} objects in total\".format(len(count)))\n",
    "  plt.show()\n",
    "\n",
    "else:\n",
    "  augmented_training_source = Training_Source\n",
    "  print('No augmentation will be used')\n",
    "\n",
    "# show augmented annotated examples\n",
    "\n",
    "if (Use_Data_augmentation):\n",
    "  df_anno_aug = []\n",
    "  dir_anno_aug = augmented_training_source_annotation\n",
    "  for fnm in os.listdir(dir_anno_aug):  \n",
    "      if not fnm.startswith('.'): ## do not include hidden folders/files\n",
    "          tree = ET.parse(os.path.join(dir_anno_aug,fnm))\n",
    "          row = extract_single_xml_file(tree)\n",
    "          row[\"fileID\"] = os.path.splitext(fnm)[0]\n",
    "          df_anno_aug.append(row)\n",
    "  df_anno_aug = pd.DataFrame(df_anno_aug)\n",
    "\n",
    "  print()\n",
    "  print('Now some images from the training set will be displayed with its annotations in order to visualise that the annotations are properly matched.')\n",
    "  print()\n",
    "\n",
    "  size = 3    \n",
    "  ind_random = np.random.randint(0,df_anno_aug.shape[0],size=size)\n",
    "  img_dir=augmented_training_source\n",
    "\n",
    "  file_suffix = os.path.splitext(os.listdir(augmented_training_source)[0])[1]\n",
    "  for irow in ind_random:\n",
    "      row  = df_anno_aug.iloc[irow,:]\n",
    "      path = os.path.join(img_dir, row[\"fileID\"] + file_suffix)\n",
    "     \n",
    "      # read in image\n",
    "      predicted_classes = np.zeros(shape=[gt_boxes[irow].shape[0]], dtype=np.int32)\n",
    "\n",
    "      for i in range(gt_classes[irow].shape[0]):\n",
    "        predicted_classes[i] = classes_dict[gt_classes[irow][i,0]]\n",
    "    \n",
    "    # read in image and convert to rgb (3 channels)\n",
    "      img = np.array(Image.open(path))\n",
    "      img = cv2.cvtColor(img,cv2.COLOR_GRAY2RGB)\n",
    "      dummy_scores = np.ones(gt_boxes[irow].shape[0], dtype=np.float32)  # give boxes a score of 100%\n",
    "      plot_detections(\n",
    "        img,\n",
    "        gt_boxes[irow],\n",
    "        predicted_classes,\n",
    "        dummy_scores, category_index, figsize=(12,12))\n",
    "  print('These are the augmented training images.')\n",
    "\n",
    "else:\n",
    "  print('Data augmentation disabled.')\n",
    "\n",
    "\n",
    "pdf_export(augmentation = Use_Data_augmentation, pretrained_model = Use_pretrained_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3L9zSGtORKYI"
   },
   "source": [
    "\n",
    "## **3.3. Using weights from a pre-trained model as initial weights**\n",
    "---\n",
    "<font size = 4>  Here, you create the RetinaNet model with the backbone chosen in the **Weights_choice** menu and then the pretrained model's weights will be loaded. There are 3 possible options:\n",
    "*   **`Resnet50`**: A very popular Neural Network used for image classification. This network rescales the images to 640x640 dimensions. **This option is the most tested one in this notebook.**\n",
    "*   **`Mobilenet_V1`**: One efficient neural network which performs very well in devices with little resources i.e. smartphones, raspberries... (etc.). This network rescales the images to 640x640 dimensions.\n",
    "*   **`Mobilenet_V2`**: The same NN as the previous but the last version of this one. This network rescales the images to 320x320 dimensions.\n",
    "    \n",
    "\n",
    "\n",
    " The values for the **classification_weight** and for **localization_weight** set the priority given to the accuracy of the classification over the location and vice versa in the training process. I.e. If **`localization_weight`** value is set to 10.0 and **`classification_weight`** to 1.0 the model is going to give 10 times more importance to improve the localization of the objects that its classification.\n",
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "9vC2n-HeLdiJ"
   },
   "outputs": [],
   "source": [
    "# @markdown ##Loading weights from a pretrained network\n",
    "\n",
    "Weights_choice = \"Resnet50\" #@param [\"Resnet50\",\"Mobilenet_V1\", \"Mobilenet_V2\"]\n",
    "classification_weight = 1.0 #@param{type:\"number\"}\n",
    "localization_weight = 1.0 #@param{type:\"number\"}\n",
    "\n",
    "# @markdown ##Or load a previously saved model instead\n",
    "Use_pretrained_model = False #@param {type:\"boolean\"}\n",
    "pretrained_model_folder = \"\" #@param {type:\"string\"}\n",
    "  \n",
    "QC_model_folder = full_model_path\n",
    "\n",
    "if Use_pretrained_model:\n",
    "  assert os.path.exists(pretrained_model_folder), \"Pretrained model's path does not exist\"\n",
    "  \n",
    "  print('Loading model\\'s checkpoints...')\n",
    "  pretrained_model, category_index, classes_dict, configs = load_pretrained_model(os.path.join(pretrained_model_folder,'saved_model'))\n",
    "  # a prediction is done in order to tune the weights \n",
    "  predict(pretrained_model ,images_path=os.path.join(Training_Source, os.listdir(Training_Source)[0]), show_image=False)\n",
    "  # full_model_path = pretrained_model_folder\n",
    "  # QC_model_folder = pretrained_model_folder\n",
    "  print(bcolors.OKGREEN + 'checkpoints loaded correctly.'+ bcolors.NORMAL)\n",
    "\n",
    "else:\n",
    "\n",
    "  # in order to make it more user friendly\n",
    "  if Weights_choice == \"Resnet50\":\n",
    "    Weights_choice = 'ssd_resnet50_v1_fpn_640x640_coco17_tpu-8'\n",
    "  elif Weights_choice == \"Mobilenet_V1\":\n",
    "    Weights_choice = 'ssd_mobilenet_v1_fpn_640x640_coco17_tpu-8'\n",
    "  elif Weights_choice == \"Mobilenet_V2\":\n",
    "    Weights_choice = 'ssd_mobilenet_v2_320x320_coco17_tpu-8'\n",
    "\n",
    "  QC_model_folder = full_model_path\n",
    "\n",
    "  checkpoints_path = os.path.join(model_path,'checkpoint')\n",
    "  if Weights_choice:\n",
    "    print('Downloading weights...')\n",
    "    download_weights(Weights_choice)\n",
    "    print('Weights downloaded succesfully.')\n",
    "    checkpoints_file = 'ckpt-0'\n",
    "    pipeline_config = fbase_path + '/models/research/object_detection/configs/tf2/{Weights_choice}.config'\n",
    "\n",
    "\n",
    "\n",
    "  if os.path.exists(checkpoints_path):\n",
    "    print('Loading model\\'s checkpoints...')\n",
    "    num_classes = len(class_count)\n",
    "    pretrained_model, configs, ckpt = load_model_checkpoints(os.path.join(checkpoints_path, checkpoints_file), pipeline_config, num_classes, iou_threshold = 0.2, classification_weight= classification_weight, localization_weight= localization_weight)\n",
    "    print(bcolors.OKGREEN + 'checkpoints loaded correctly.'+ bcolors.NORMAL)\n",
    "  else:\n",
    "    print('Checkpoint\\'s path does not exist.')\n",
    "\n",
    "  \n",
    "pdf_export(augmentation = Use_Data_augmentation, pretrained_model = Use_pretrained_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "MCGklf1vZf2M"
   },
   "source": [
    "# **4. Train the network**\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0Dfn8ZsEMv5d"
   },
   "source": [
    "## **4.1. Start training**\n",
    "---\n",
    "<font size = 4>When playing the cell below you should see updates after each epoch (round). Network training can take some time.\n",
    "\n",
    "<font size = 4>* **CRITICAL NOTE:** Google Colab has a time limit for processing (to prevent using GPU power for datamining). Training time must be less than 12 hours! If training takes longer than 12 hours, please decrease the number of epochs or number of patches.\n",
    "\n",
    "<!-- <font size = 4>Once training is complete, the trained model is automatically saved on your Google Drive, in the **model_path** folder that was selected in Section 3. It is however wise to download the folder as all data can be erased at the next training if using the same folder. -->"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "iwNmp1PUzRDQ",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "#@markdown ##Start training\n",
    "#@markdown Check the **`Verbose`** field if you want to show how the mAP metric change every epoch. This function slows down a bit the training process.\n",
    "Verbose = True #@param {type:\"boolean\"}\n",
    "\n",
    "\n",
    "#Write the annotations to a csv file\n",
    "df_anno.to_csv(full_model_path+'/annotations.csv', index=False)#header=False, sep=',')\n",
    "\n",
    "if os.path.exists(full_model_path+\"/Quality Control\"):\n",
    "  shutil.rmtree(full_model_path+\"/Quality Control\")\n",
    "os.makedirs(full_model_path+\"/Quality Control\")\n",
    "\n",
    "# creation of the folder where the model will be saved\n",
    "saved_model_path = os.path.join(full_model_path,'saved_model')\n",
    "if not os.path.exists(saved_model_path):\n",
    "  # shutil.rmtree(saved_model_path)\n",
    "  os.mkdir(saved_model_path)\n",
    "  save_model_config(configs)\n",
    "  save_category_index()\n",
    "\n",
    "\n",
    "start = time.time()\n",
    "\n",
    "train(pretrained_model, verbose = Verbose)\n",
    "\n",
    "# Displaying the time elapsed for training\n",
    "dt = time.time() - start\n",
    "mins, sec = divmod(dt, 60) \n",
    "hour, mins = divmod(mins, 60) \n",
    "print(\"Time elapsed:\",hour, \"hour(s)\",mins,\"min(s)\",round(sec),\"sec(s)\")\n",
    "\n",
    "#Create a pdf document with training summary\n",
    "\n",
    "pdf_export(trained = True, augmentation = Use_Data_augmentation, pretrained_model = Use_pretrained_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_0Hynw3-xHp1"
   },
   "source": [
    "# **5. Evaluate your model**\n",
    "---\n",
    "\n",
    "<font size = 4>This section allows the user to perform important quality checks on the validity and generalisability of the trained model. \n",
    "\n",
    "<font size = 4>**We highly recommend to perform quality control on all newly trained models.**\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "eAJzMwPA6tlH"
   },
   "outputs": [],
   "source": [
    "# model name and path\n",
    "#@markdown ###Do you want to assess the model you just trained ?\n",
    "Use_the_current_trained_model = True #@param {type:\"boolean\"}\n",
    "\n",
    "# @markdown ###If not, please provide the name of the model folder:\n",
    "# @markdown ***Note:***  *The model folder is the path of an already trained model in some previous usage of this notebook. E.g.:  `base_path + /gdrive/MyDrive/RetinaNet`*\n",
    "\n",
    "QC_model_folder = \"\" #@param {type:\"string\"}\n",
    "if QC_model_folder and QC_model_folder[-1] == '/':\n",
    "  QC_model_folder = QC_model_folder[:-1]\n",
    "\n",
    "if (Use_the_current_trained_model): \n",
    "  QC_model_folder = full_model_path\n",
    "else:\n",
    "  # load model from folder\n",
    "  assert os.path.exists(QC_model_folder), \"Model folder path does not exist.\"\n",
    "  pretrained_model, category_index, classes_dict, configs = load_pretrained_model(os.path.join(QC_model_folder,'saved_model'))\n",
    "QC_model_name = os.path.basename(QC_model_folder)\n",
    "\n",
    "if os.path.exists(QC_model_folder):\n",
    "  print(\"The \"+QC_model_name+\" model will be evaluated\")\n",
    "else:\n",
    "  W  = '\\033[0m'  # white (normal)\n",
    "  R  = '\\033[31m' # red\n",
    "  print(R+'!! WARNING: The chosen model does not exist !!'+W)\n",
    "  print('Please make sure you provide a valid model path before proceeding further.')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "dhJROwlAMv5o"
   },
   "source": [
    "## **5.1. Inspection of the loss function**\n",
    "---\n",
    "\n",
    "<font size = 4>First, it is good practice to evaluate the training progress by comparing the training loss with the validation loss. The latter is a metric which shows how well the network performs on a subset of unseen data which is set aside from the training dataset. For more information on this, see for example [this review](https://www.ncbi.nlm.nih.gov/pmc/articles/PMC6381354/) by Nichols *et al.*\n",
    "\n",
    "<font size = 4>**Training loss** describes an error value after each epoch for the difference between the model's prediction and its ground-truth target.\n",
    "\n",
    "<font size = 4>**Validation loss** describes the same error value between the model's prediction on a validation image and compared to it's target.\n",
    "\n",
    "<font size = 4>During training both values should decrease before reaching a minimal value which does not decrease further even after more training. Comparing the development of the validation loss with the training loss can give insights into the model's performance.\n",
    "\n",
    "<font size = 4>Decreasing **Training loss** and **Validation loss** indicates that training is still necessary and increasing the `number_of_epochs` is recommended. Note that the curves can look flat towards the right side, just because of the y-axis scaling. The network has reached convergence once the curves flatten out. After this point no further training is required. If the **Validation loss** suddenly increases again an the **Training loss** simultaneously goes towards zero, it means that the network is overfitting to the training data. In other words the network is remembering the exact patterns from the training data and no longer generalizes well to unseen data. In this case the training dataset has to be increased."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "vMzSP50kMv5p"
   },
   "outputs": [],
   "source": [
    "#@markdown ##Play the cell to show a plot of training errors vs. epoch number\n",
    "import csv\n",
    "from matplotlib import pyplot as plt\n",
    "%matplotlib inline\n",
    "lossDataFromCSV = []\n",
    "vallossDataFromCSV = []\n",
    "mAPDataFromCSV = []\n",
    "with open(QC_model_folder+'/Quality Control/training_evaluation.csv','r') as csvfile:\n",
    "    csvRead = csv.reader(csvfile, delimiter=',')\n",
    "    next(csvRead)\n",
    "    for row in csvRead:\n",
    "        lossDataFromCSV.append(float(row[0]))\n",
    "        vallossDataFromCSV.append(float(row[1]))\n",
    "        mAPDataFromCSV.append(float(row[2]))\n",
    "epochNumber = range(len(lossDataFromCSV))\n",
    "plt.figure(figsize=(20,15))\n",
    "\n",
    "plt.subplot(3,1,1)\n",
    "plt.plot(epochNumber,lossDataFromCSV, label='Training loss')\n",
    "plt.plot(epochNumber,vallossDataFromCSV, label='Validation loss')\n",
    "plt.title('Training loss and validation loss vs. epoch number (linear scale)')\n",
    "plt.ylabel('Loss')\n",
    "plt.xlabel('Epoch number')\n",
    "plt.legend()\n",
    "\n",
    "plt.subplot(3,1,2)\n",
    "plt.semilogy(epochNumber,lossDataFromCSV, label='Training loss')\n",
    "plt.semilogy(epochNumber,vallossDataFromCSV, label='Validation loss')\n",
    "plt.title('Training loss and validation loss vs. epoch number (log scale)')\n",
    "plt.ylabel('Loss')\n",
    "plt.xlabel('Epoch number')\n",
    "plt.legend()\n",
    "\n",
    "\n",
    "plt.subplot(3,1,3)\n",
    "plt.plot(epochNumber,mAPDataFromCSV, label='mAP score')\n",
    "plt.title('mean average precision (mAP) vs. epoch number (linear scale)')\n",
    "plt.ylabel('mAP score')\n",
    "plt.xlabel('Epoch number')\n",
    "plt.legend()\n",
    "plt.savefig(QC_model_folder+'/Quality Control/lossCurveAndmAPPlots.png',bbox_inches='tight', pad_inches=0)\n",
    "plt.show()\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "X5_92nL2xdP6"
   },
   "source": [
    "## **5.2. Error mapping and quality metrics estimation**\n",
    "---\n",
    "\n",
    "<font size = 4>This section will display the predictions of the model for the images contained in a folder your specify. Additionally, the below cell will show the mAP value of the model on the QC data together with plots of the Precision-Recall curves for all the classes in the dataset. If you want to read in more detail about these scores, we recommend [this brief explanation](https://medium.com/@jonathan_hui/map-mean-average-precision-for-object-detection-45c121a31173).\n",
    "\n",
    "<font size = 4> The images provided in the \"Source_QC_folder\" and \"Target_QC_folder\" should contain images (e.g. as .jpg) and annotations (.xml files)!\n",
    "\n",
    "<font size = 4>For example, if you want to evaluate your model over the example dataset from [section 3](#example_dataset), the paths you should use are *`base_path + /example_dataset/Test_Images`* for the Source_QC_folder and *`base_path + /example_dataset/Test_Annotations`* for the Annotations_QC_folder.\n",
    "\n",
    "<font size = 4>**mAP score:** This refers to the mean average precision of the model on the given dataset. This value gives an indication how precise the predictions of the classes on this dataset are when compared to the ground-truth. Values closer to 1 indicate a good fit.\n",
    "\n",
    "<font size = 4>**Precision:** This is the proportion of the correct classifications (true positives) in all the predictions made by the model.\n",
    "\n",
    "<font size = 4>**Recall:** This is the proportion of the detected true positives in all the detectable data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "w90MdriMxhjD"
   },
   "outputs": [],
   "source": [
    "#@markdown ##Choose the folders that contain your Quality Control dataset\n",
    "\n",
    "Source_QC_folder = \"\" #@param{type:\"string\"}\n",
    "Annotations_QC_folder = \"\" #@param{type:\"string\"}\n",
    "\n",
    "## @markdown ##Choose which model you want to evaluate:\n",
    "\n",
    "file_suffix = os.path.splitext(os.listdir(Source_QC_folder)[0])[1]\n",
    "\n",
    "# Create a quality control/Prediction Folder\n",
    "if os.path.exists(QC_model_folder+\"/Quality Control/Prediction\"):\n",
    "  shutil.rmtree(QC_model_folder+\"/Quality Control/Prediction\")\n",
    "\n",
    "os.makedirs(QC_model_folder+\"/Quality Control/Prediction\")\n",
    "\n",
    "#Delete old csv with box predictions if one exists\n",
    "\n",
    "if os.path.exists(base_path + '/predicted_bounding_boxes.csv'):\n",
    "  os.remove(base_path + '/predicted_bounding_boxes.csv')\n",
    "if os.path.exists(base_path + '/predicted_bounding_boxes_names.csv'):\n",
    "  os.remove(base_path + '/predicted_bounding_boxes_names.csv')\n",
    "if os.path.exists(Source_QC_folder+'/.ipynb_checkpoints'):\n",
    "  shutil.rmtree(Source_QC_folder+'/.ipynb_checkpoints')\n",
    "\n",
    "\n",
    "n_objects = []\n",
    "for img in os.listdir(Source_QC_folder):\n",
    "  full_image_path = Source_QC_folder+'/'+img\n",
    "  if not img[-12:-4] == 'detected':\n",
    "    print(img)\n",
    "    n_obj = predict(pretrained_model, full_image_path, min_score_threshold=0.3)\n",
    "    n_objects.append(len(n_obj[0]))\n",
    "    K.clear_session()\n",
    "\n",
    "for img in os.listdir(Source_QC_folder):\n",
    "  if img.endswith('detected'+file_suffix):\n",
    "    shutil.move(Source_QC_folder+'/'+img,QC_model_folder+\"/Quality Control/Prediction/\"+img)\n",
    "\n",
    "# #Make a csv file to read into imagej macro, to create custom bounding boxes\n",
    "header = ['filename']+['xmin', 'ymin', 'xmax', 'ymax', 'confidence', 'class']*max(n_objects)\n",
    "with open(base_path + '/predicted_bounding_boxes.csv', newline='') as inFile, open(base_path + '/predicted_bounding_boxes_new.csv', 'w', newline='') as outfile:\n",
    "    r = csv.reader(inFile)\n",
    "    w = csv.writer(outfile)\n",
    "    next(r, None)  # skip the first row from the reader, the old header\n",
    "    # write new header\n",
    "    w.writerow(header)\n",
    "    # copy the rest\n",
    "    for row in r:\n",
    "        w.writerow(row)\n",
    "\n",
    "df_bbox=pd.read_csv(base_path + '/predicted_bounding_boxes_new.csv',error_bad_lines=False)\n",
    "df_bbox=df_bbox.transpose()\n",
    "new_header = df_bbox.iloc[0] #grab the first row for the header\n",
    "df_bbox = df_bbox[1:] #take the data less the header row\n",
    "df_bbox.columns = new_header #set the header row as the df header\n",
    "df_bbox.sort_values(by='filename',axis=1,inplace=True)\n",
    "df_bbox.to_csv(QC_model_folder+'/Quality Control/predicted_bounding_boxes_for_custom_ROI_QC.csv')\n",
    "\n",
    "F1_scores, AP, recall, precision = _calc_avg_precisions(pretrained_model,Source_QC_folder,Annotations_QC_folder+'/',0.3,0.3)\n",
    "\n",
    "\n",
    "\n",
    "with open(QC_model_folder+\"/Quality Control/QC_results.csv\", \"r\") as file:\n",
    "    x = from_csv(file)\n",
    "  \n",
    "print(x)\n",
    "mAP_score = sum(AP.values())/len(AP)\n",
    "\n",
    "print('mAP score for QC dataset: '+str(mAP_score))\n",
    "\n",
    "for i in range(len(AP)):\n",
    "  if AP[i]!=0:\n",
    "    fig = plt.figure(figsize=(8,4))\n",
    "    if len(recall[i]) == 1:\n",
    "      new_recall = np.linspace(0,list(recall[i])[0],10)\n",
    "      new_precision = list(precision[i])*10\n",
    "      fig = plt.figure(figsize=(3,2))\n",
    "      plt.plot(new_recall,new_precision)\n",
    "      plt.axis([min(new_recall),1,0,1.02])\n",
    "      plt.xlabel('Recall',fontsize=14)\n",
    "      plt.ylabel('Precision',fontsize=14)\n",
    "      plt.title(f'{category_index[i+1][\"name\"]}, AP: '+str(round(AP[i],3)),fontsize=14)\n",
    "      plt.fill_between(new_recall,new_precision,alpha=0.3)\n",
    "      plt.savefig(QC_model_folder+f'/Quality Control/P-R_curve_{category_index[i+1][\"name\"]}.png',  bbox_inches='tight', pad_inches=0)\n",
    "      plt.show()\n",
    "    else:\n",
    "      new_recall = list(recall[i])\n",
    "      new_recall.append(new_recall[len(new_recall)-1])\n",
    "      new_precision = list(precision[i])\n",
    "      new_precision.append(0)\n",
    "      plt.plot(new_recall,new_precision)\n",
    "      plt.axis([min(new_recall),1,0,1.02])\n",
    "      plt.xlabel('Recall',fontsize=14)\n",
    "      plt.ylabel('Precision',fontsize=14)\n",
    "      plt.title(f'{category_index[i+1][\"name\"]}, AP: '+str(round(AP[i],3)),fontsize=14)\n",
    "      plt.fill_between(new_recall,new_precision,alpha=0.3)\n",
    "      plt.savefig(QC_model_folder+'/Quality Control/P-R_curve_'+f'{category_index[i+1][\"name\"]}.png', bbox_inches='tight', pad_inches=0)\n",
    "      plt.show()\n",
    "  else:\n",
    "    print(f'No object of class {category_index[i+1][\"name\"]} was detected. This will lower the mAP score. Consider adding an image containing this class to your QC dataset to see if the model can detect this class at all.')\n",
    "\n",
    "\n",
    "# --------------------------------------------------------------\n",
    "add_header(base_path + '/predicted_bounding_boxes_names.csv',base_path + '/predicted_bounding_boxes_names_new.csv')\n",
    "\n",
    "# # # This will display a randomly chosen dataset input and predicted output\n",
    "\n",
    "print('Below is an example input, prediction and ground truth annotation from your test dataset.')\n",
    "random_choice = random.choice(os.listdir(Source_QC_folder))\n",
    "file_suffix = os.path.splitext(random_choice)[1]\n",
    "\n",
    "plt.figure(figsize=(30,15))\n",
    "\n",
    "# # ### Display Raw input ###\n",
    "\n",
    "x = plt.imread(Source_QC_folder+\"/\"+random_choice)\n",
    "plt.subplot(1,3,1)\n",
    "plt.axis('off')\n",
    "plt.imshow(x, interpolation='nearest', cmap='gray')\n",
    "plt.title('Input', fontsize = 12)\n",
    "\n",
    "# # ### Display Predicted annotation ###\n",
    "\n",
    "predict(pretrained_model, os.path.join(Source_QC_folder ,random_choice), save_path=QC_model_folder+'/Quality Control/QC_Prediction.png', show_image=False, min_score_threshold=0.3)\n",
    "\n",
    "# # ### Display GT Annotation ###\n",
    "\n",
    "df_anno_QC_gt = []\n",
    "for fnm in os.listdir(Annotations_QC_folder):  \n",
    "    if not fnm.startswith('.'): ## do not include hidden folders/files\n",
    "        tree = ET.parse(os.path.join(Annotations_QC_folder,fnm))\n",
    "        row = extract_single_xml_file(tree)\n",
    "        row[\"fileID\"] = os.path.splitext(fnm)[0]\n",
    "        df_anno_QC_gt.append(row)\n",
    "df_anno_QC_gt = pd.DataFrame(df_anno_QC_gt)\n",
    "\n",
    "for i in range(0,df_anno_QC_gt.shape[0]):\n",
    "  if df_anno_QC_gt.iloc[i][\"fileID\"]+file_suffix == random_choice:\n",
    "    row = df_anno_QC_gt.iloc[i]\n",
    "\n",
    "\n",
    "gt_boxes_QC = []\n",
    "gt_classes_QC = []\n",
    "objects_coords = np.zeros((row[\"Nobj\"],4))\n",
    "objects_classes = np.zeros((row[\"Nobj\"],1), dtype=object)\n",
    "\n",
    "for iplot in range(row[\"Nobj\"]):\n",
    "  objects_coords[iplot,0] = row[\"bbx_{}_ymin\".format(iplot)] / row['height']\n",
    "  objects_coords[iplot,1] = row[\"bbx_{}_xmin\".format(iplot)] / row['width']\n",
    "  objects_coords[iplot,2] = row[\"bbx_{}_ymax\".format(iplot)] / row['height']\n",
    "  objects_coords[iplot,3] = row[\"bbx_{}_xmax\".format(iplot)] / row['width']\n",
    "  objects_classes[iplot,0] = row[\"bbx_{}_name\".format(iplot)]\n",
    "\n",
    "gt_boxes_QC.append(objects_coords)\n",
    "gt_classes_QC.append(objects_classes)\n",
    "\n",
    "  \n",
    "\n",
    "# read in image\n",
    "predicted_classes_QC = np.zeros(shape=[gt_boxes_QC[0].shape[0]], dtype=np.int32)\n",
    "\n",
    "for i in range(gt_classes_QC[0].shape[0]):\n",
    "  predicted_classes_QC[i] = classes_dict[gt_classes_QC[0][i,0]]\n",
    "\n",
    "# read in image and convert to rgb (3 channels)\n",
    "\n",
    "img = np.array(Image.open(os.path.join(Source_QC_folder ,random_choice)))\n",
    "img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)\n",
    "dummy_scores = np.ones(gt_boxes_QC[0].shape[0], dtype=np.float32)  # give boxes a score of 100%\n",
    "plot_detections(\n",
    "  img,\n",
    "  gt_boxes_QC[0],\n",
    "  predicted_classes_QC,\n",
    "  dummy_scores, category_index, figsize=(12,12), image_name=QC_model_folder+'/Quality Control/GT_QC_image.png', show_image=False)\n",
    "\n",
    "\n",
    "predicted_QC = plt.imread(QC_model_folder+'/Quality Control/QC_Prediction.png')\n",
    "plt.subplot(1,3,2)\n",
    "plt.axis('off')\n",
    "plt.imshow(predicted_QC, interpolation='nearest', cmap='gray')\n",
    "plt.title('Prediction', fontsize = 12)\n",
    "\n",
    "GT_QC = plt.imread(QC_model_folder+'/Quality Control/GT_QC_image.png')\n",
    "plt.subplot(1,3,3)\n",
    "plt.axis('off')\n",
    "plt.imshow(GT_QC, interpolation='nearest', cmap='gray')\n",
    "plt.title('Ground Truth annotations', fontsize = 12)\n",
    "\n",
    "\n",
    "# ### Show the plot ###\n",
    "plt.savefig(QC_model_folder+'/Quality Control/QC_example_data.png',bbox_inches='tight',pad_inches=0)\n",
    "plt.show()\n",
    "\n",
    "\n",
    "#Make a pdf summary of the QC results\n",
    "\n",
    "qc_pdf_export()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-tJeeJjLnRkP"
   },
   "source": [
    "# **6. Using the trained model**\n",
    "\n",
    "---\n",
    "\n",
    "<font size = 4>In this section the unseen data is processed using the trained model (in section 4). First, your unseen images are uploaded and prepared for prediction. After that your trained model from section 4 is activated and finally saved into your Google Drive."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "d8wuQGjoq6eN"
   },
   "source": [
    "## **6.1. Generate prediction(s) from unseen dataset**\n",
    "---\n",
    "\n",
    "<font size = 4>The current trained model (from section 4.2) can now be used to process images. If you want to use an older model, untick the **Use_the_current_trained_model** box and enter the name and path of the model to use. Predicted output images are saved in your **Result_folder** folder. Even though all the images will be predicted, and those predictions will be saved in the  **Result_folder** path, only a maximum number of 5 images will be displayed in this cell.\n",
    "\n",
    "<font size = 4>**`Data_folder`:** This folder should contain the images that you want to use your trained network on for processing.\n",
    "\n",
    "<font size = 4>**`Result_folder`:** This folder will contain the predicted output images. In case the folder does not exist it will be created.\n",
    "\n",
    "<font size = 4>**`Prediction_model_path`:** This should be the folder that contains your model. The model folder is the path of an already trained model in some previous usage of this notebook. E.g.: *`base_path + /gdrive/MyDrive/RetinaNet`*\n",
    "\n",
    "<font size = 4>**`Score_threshold`:** Insert the threshold of accuracy to the model, this is, the minimum percentage of confidence of the model in the predictions. Decimal number between 0 and 1. If the threshold is 1 there will be shown only the predictions with 100% of confidence (difficult) and, otherwise, if the threshold is 0 there will be shown all the predictions. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "y2TD5p7MZrEb"
   },
   "outputs": [],
   "source": [
    "#@markdown ### Provide the path to your dataset and to the folder where the predictions are saved, then play the cell to predict outputs from your unseen images with a minimum confidence established by the Score_threshold value.\n",
    "\n",
    "Data_folder = \"\" #@param {type:\"string\"}\n",
    "Result_folder = \"\" #@param {type:\"string\"}\n",
    "Score_threshold =  0.4#@param {type:\"number\"}\n",
    "\n",
    "assert os.path.exists(Data_folder), \"Data_folder path does not exist.\"\n",
    "\n",
    "if Data_folder[-1] == '/':\n",
    "  Data_folder = Data_folder[:-1]\n",
    "\n",
    "if os.path.exists(Data_folder+'/.ipynb_checkpoints'):\n",
    "  shutil.rmtree(Data_folder+'/.ipynb_checkpoints')\n",
    "\n",
    "file_suffix = os.path.splitext(os.listdir(Data_folder)[0])[1]\n",
    "# model name and path\n",
    "#@markdown ###Do you want to use the current trained model?\n",
    "Use_the_current_trained_model = True #@param {type:\"boolean\"}\n",
    "\n",
    "#@markdown ###If not, provide the path to model folder:\n",
    "\n",
    "Prediction_model_path = \"insert text here\" #@param {type:\"string\"}\n",
    "\n",
    "\n",
    "\n",
    "if not os.path.exists(Result_folder):\n",
    "  os.makedirs(Result_folder)\n",
    "\n",
    "if (Use_the_current_trained_model): \n",
    "  print(\"Using current trained network\")\n",
    "  Prediction_model_path = QC_model_folder\n",
    "else:\n",
    "  pretrained_model, category_index, classes_dict = load_pretrained_model(os.path.join(Prediction_model_path,'saved_model'))\n",
    "\n",
    "\n",
    "# Provide the code for performing predictions and saving them\n",
    "print(\"Images will be saved into folder:\", Result_folder)\n",
    "\n",
    "\n",
    "# ----- Predictions ------\n",
    "\n",
    "start = time.time()\n",
    "\n",
    "#Remove any files that might be from the prediction of QC examples.\n",
    "if os.path.exists(base_path + '/predicted_bounding_boxes.csv'):\n",
    "  os.remove(base_path + '/predicted_bounding_boxes.csv')\n",
    "if os.path.exists(base_path + '/predicted_bounding_boxes_new.csv'):\n",
    "  os.remove(base_path + '/predicted_bounding_boxes_new.csv')\n",
    "if os.path.exists(base_path + '/predicted_bounding_boxes_names.csv'):\n",
    "  os.remove(base_path + '/predicted_bounding_boxes_names.csv')\n",
    "if os.path.exists(base_path + '/predicted_bounding_boxes_names_new.csv'):\n",
    "  os.remove(base_path + '/predicted_bounding_boxes_names_new.csv')\n",
    "\n",
    "\n",
    "if os.path.exists(Data_folder+'/.ipynb_checkpoints'):\n",
    "  shutil.rmtree(Data_folder+'/.ipynb_checkpoints')\n",
    "\n",
    "n_objects = []\n",
    "# show_img = True\n",
    "show_img_index = 0 # the first 5 images are shown\n",
    "for img in os.listdir(Data_folder):\n",
    "  full_image_path = Data_folder+'/'+img\n",
    "  \n",
    "  # if show_img_index > 5:\n",
    "  #   show_img = False\n",
    "  n_obj = predict(pretrained_model,full_image_path, save_path=full_image_path[:-4] + '_detected' + full_image_path[-4:] , min_score_threshold=Score_threshold, show_image=(show_img_index < 5))#,Result_folder)\n",
    "  n_objects.append(n_obj)\n",
    "  show_img_index +=1\n",
    "  K.clear_session()\n",
    "\n",
    "for img in os.listdir(Data_folder):\n",
    "  if img.endswith('detected'+file_suffix):\n",
    "    shutil.move(Data_folder+'/'+img,Result_folder+'/'+img)\n",
    "\n",
    "if os.path.exists(base_path + '/predicted_bounding_boxes.csv'):\n",
    "  print('Bounding box labels and coordinates saved to '+ Result_folder)\n",
    "else:\n",
    "  print('For some reason the bounding box labels and coordinates were not saved. Check that your predictions look as expected.')\n",
    "\n",
    "#Make a csv file to read into imagej macro, to create custom bounding boxes\n",
    "header = ['filename']+['xmin', 'ymin', 'xmax', 'ymax', 'confidence', 'class']*100#max(n_objects)\n",
    "with open(base_path + '/predicted_bounding_boxes.csv', newline='') as inFile, open(base_path + '/predicted_bounding_boxes_new.csv', 'w', newline='') as outfile:\n",
    "    r = csv.reader(inFile)\n",
    "    w = csv.writer(outfile)\n",
    "    next(r, None)  # skip the first row from the reader, the old header\n",
    "    # write new header\n",
    "    w.writerow(header)\n",
    "    # copy the rest\n",
    "    for row in r:\n",
    "        w.writerow(row)\n",
    "\n",
    "df_bbox=pd.read_csv(base_path + '/predicted_bounding_boxes_new.csv')\n",
    "df_bbox=df_bbox.transpose()\n",
    "new_header = df_bbox.iloc[0] #grab the first row for the header\n",
    "df_bbox = df_bbox[1:] #take the data less the header row\n",
    "df_bbox.columns = new_header #set the header row as the df header\n",
    "df_bbox.sort_values(by='filename',axis=1,inplace=True)\n",
    "df_bbox.to_csv(Result_folder+'/predicted_bounding_boxes_for_custom_ROI.csv')\n",
    "\n",
    "# Displaying the time elapsed for training\n",
    "dt = time.time() - start\n",
    "mins, sec = divmod(dt, 60) \n",
    "hour, mins = divmod(mins, 60) \n",
    "print(\"Time elapsed:\",hour, \"hour(s)\",mins,\"min(s)\",round(sec),\"sec(s)\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "tP1isF0PO4C1"
   },
   "source": [
    "## **6.2. Inspect the predicted output**\n",
    "---\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "ypLeYWnzO6tv"
   },
   "outputs": [],
   "source": [
    "# @markdown ##Run this cell to display a randomly chosen input and its corresponding predicted output.\n",
    "import random\n",
    "from matplotlib.pyplot import imread\n",
    "# This will display a randomly chosen dataset input and predicted output\n",
    "random_choice = random.choice(os.listdir(Data_folder))\n",
    "print(random_choice)\n",
    "x = imread(Data_folder+\"/\"+random_choice)\n",
    "\n",
    "os.chdir(Result_folder)\n",
    "y = imread(Result_folder+\"/\"+os.path.splitext(random_choice)[0]+'_detected'+file_suffix)\n",
    "\n",
    "plt.figure(figsize=(20,8))\n",
    "\n",
    "plt.subplot(1,3,1)\n",
    "plt.axis('off')\n",
    "plt.imshow(x, interpolation='nearest', cmap='gray')\n",
    "plt.title('Input')\n",
    "\n",
    "plt.subplot(1,3,2)\n",
    "plt.axis('off')\n",
    "plt.imshow(y, interpolation='nearest')\n",
    "plt.title('Predicted output');\n",
    "\n",
    "add_header(base_path + '/predicted_bounding_boxes_names.csv',base_path + '/predicted_bounding_boxes_names_new.csv')\n",
    "\n",
    "#We need to edit this predicted_bounding_boxes_new.csv file slightly to display the bounding boxes\n",
    "df_bbox2 = pd.read_csv(base_path + '/predicted_bounding_boxes_names_new.csv')\n",
    "for img in range(0,df_bbox2.shape[0]):\n",
    "  df_bbox2.iloc[img]\n",
    "  row = pd.DataFrame(df_bbox2.iloc[img])\n",
    "  if row[img][0] == random_choice:\n",
    "    row = row.dropna()\n",
    "    image  = imageio.imread(Data_folder+'/'+row[img][0])\n",
    "    plt.subplot(1,3,3)\n",
    "    plt.axis('off')\n",
    "    plt.title('Alternative Display of Prediction')\n",
    "    plt.imshow(image, cmap='gray') # plot image\n",
    "\n",
    "    for i in range(1,int(len(row)-1),6):\n",
    "      plt_rectangle(plt,\n",
    "                    label = row[img][i+5],\n",
    "                    x1=row[img][i],\n",
    "                    y1=row[img][i+1],\n",
    "                    x2=row[img][i+2],\n",
    "                    y2=row[img][i+3])\n",
    "                    \n",
    "    \n",
    "    plt.savefig(base_path + '/detected_cells.png',bbox_inches='tight',transparent=True,pad_inches=0)\n",
    "plt.show() ## show the plot\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hvkd66PldsXB"
   },
   "source": [
    "## **6.3. Download your predictions**\n",
    "---\n",
    "\n",
    "<font size = 4>**Store your data** and ALL its results elsewhere by downloading it from Google Drive and after that clean the original folder tree (datasets, results, trained model etc.) if you plan to train or use new networks. Please note that the notebook will otherwise **OVERWRITE** all files which have the same name."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "AbvUcekhLqPj"
   },
   "source": [
    "# **7. Version log**\n",
    "---\n",
    "\n",
    "<font size = 4>**v1.14.1**:  \n",
    "*  Fixed the PDF generation (fpdf2, delimiter, Cournier)\n",
    "*  Removed the paths to /content\n",
    "\n",
    "<font size = 4>**v1.13**:  \n",
    "\n",
    "\n",
    "*   First official release of the notebook!\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "UvSlTaH14s3t"
   },
   "source": [
    "# **Thank you for using RetinaNet!**"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "machine_shape": "hm",
   "name": "RetinaNet_ZeroCostDL4Mic.ipynb",
   "provenance": []
  },
  "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}