Yolo_V2_tf_eager.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "jyG2Gd6ZeQRS"
   },
   "source": [
    "# YOLO V2 with TensorFlow - Eager execution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "0X7wmRCheQRT"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import glob\n",
    "import re\n",
    "import numpy as np\n",
    "import cv2\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.patches as patches\n",
    "import xml.etree.ElementTree as ET\n",
    "import imgaug as ia\n",
    "from imgaug import augmenters as iaa"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "SIIXXvcFeQRf"
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "from tensorflow import keras\n",
    "import tensorflow.keras.backend as K\n",
    "from tensorflow.keras.layers import Concatenate, concatenate, Dropout, LeakyReLU, Reshape, Activation, Conv2D, Input, MaxPooling2D, BatchNormalization, Flatten, Dense, Lambda"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 52
    },
    "id": "Mziy04f3eQRo",
    "outputId": "e157fa2e-d3f5-4746-a80b-da24865bdf0c"
   },
   "outputs": [],
   "source": [
    "tf.enable_eager_execution()\n",
    "\n",
    "print(\"TensorFlow version: {}\".format(tf.VERSION))\n",
    "print(\"Eager execution: {}\".format(tf.executing_eagerly()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "J_05RzLxeQRz"
   },
   "outputs": [],
   "source": [
    "# Parameters\n",
    "\n",
    "LABELS           = ('sugarbeet', 'weed')\n",
    "IMAGE_H, IMAGE_W = 512, 512\n",
    "GRID_H,  GRID_W  = 16, 16 # GRID size = IMAGE size / 32\n",
    "\n",
    "BOX              = 5\n",
    "CLASS            = len(LABELS)\n",
    "SCORE_THRESHOLD  = 0.5\n",
    "IOU_THRESHOLD    = 0.45\n",
    "ANCHORS          = [0.57273, 0.677385, 1.87446, 2.06253, 3.33843, 5.47434, 7.88282, 3.52778, 9.77052, 9.16828]\n",
    "\n",
    "TRAIN_BATCH_SIZE = 1\n",
    "VAL_BATCH_SIZE   = 1\n",
    "EPOCHS           = 100\n",
    "\n",
    "LAMBDA_NOOBJECT  = 1\n",
    "LAMBDA_OBJECT    = 5\n",
    "LAMBDA_CLASS     = 1\n",
    "LAMBDA_COORD     = 1\n",
    "\n",
    "max_annot        = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "-ahZ0qWBeQR8"
   },
   "outputs": [],
   "source": [
    "# Train and validation directory\n",
    "\n",
    "train_image_folder = 'data/train/image/'\n",
    "train_annot_folder = 'data/train/annotation/'\n",
    "val_image_folder = 'data/val/image/'\n",
    "val_annot_folder = 'data/val/annotation/'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "NM2Ny9vxeQSB"
   },
   "source": [
    "# 1. Define YOLO model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tyuMIK3VeQSC"
   },
   "outputs": [],
   "source": [
    "# Custom Keras layer\n",
    "\n",
    "class SpaceToDepth(keras.layers.Layer):\n",
    "\n",
    "    def __init__(self, block_size, **kwargs):\n",
    "        self.block_size = block_size\n",
    "        super(SpaceToDepth, self).__init__(**kwargs)\n",
    "\n",
    "    def call(self, inputs):\n",
    "        x = inputs\n",
    "        batch, height, width, depth = K.int_shape(x)\n",
    "        reduced_height = height // self.block_size\n",
    "        reduced_width = width // self.block_size\n",
    "        y = K.reshape(x, (batch, reduced_height, self.block_size,\n",
    "                             reduced_width, self.block_size, depth))\n",
    "        z = K.permute_dimensions(y, (0, 1, 3, 2, 4, 5))\n",
    "        t = K.reshape(z, (batch, reduced_height, reduced_width, -1))\n",
    "        return t\n",
    "\n",
    "    def compute_output_shape(self, input_shape):\n",
    "        shape =  (input_shape[0], input_shape[1] // self.block_size, input_shape[2] // self.block_size,\n",
    "                  input_shape[3] * self.block_size **2)\n",
    "        return tf.TensorShape(shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "z02gy62XeQSH"
   },
   "outputs": [],
   "source": [
    "# Yolo model (thanks to https://github.com/experiencor/keras-yolo2)\n",
    "\n",
    "input_image = tf.keras.layers.Input((IMAGE_H, IMAGE_W, 3), dtype='float32')\n",
    "\n",
    "# Layer 1\n",
    "x = Conv2D(32, (3,3), strides=(1,1), padding='same', name='conv_1', use_bias=False)(input_image)\n",
    "x = BatchNormalization(name='norm_1')(x)\n",
    "x = LeakyReLU(alpha=0.1)(x)\n",
    "x = MaxPooling2D(pool_size=(2, 2))(x)\n",
    "\n",
    "# Layer 2\n",
    "x = Conv2D(64, (3,3), strides=(1,1), padding='same', name='conv_2', use_bias=False)(x)\n",
    "x = BatchNormalization(name='norm_2')(x)\n",
    "x = LeakyReLU(alpha=0.1)(x)\n",
    "x = MaxPooling2D(pool_size=(2, 2))(x)\n",
    "\n",
    "# Layer 3\n",
    "x = Conv2D(128, (3,3), strides=(1,1), padding='same', name='conv_3', use_bias=False)(x)\n",
    "x = BatchNormalization(name='norm_3')(x)\n",
    "x = LeakyReLU(alpha=0.1)(x)\n",
    "\n",
    "# Layer 4\n",
    "x = Conv2D(64, (1,1), strides=(1,1), padding='same', name='conv_4', use_bias=False)(x)\n",
    "x = BatchNormalization(name='norm_4')(x)\n",
    "x = LeakyReLU(alpha=0.1)(x)\n",
    "\n",
    "# Layer 5\n",
    "x = Conv2D(128, (3,3), strides=(1,1), padding='same', name='conv_5', use_bias=False)(x)\n",
    "x = BatchNormalization(name='norm_5')(x)\n",
    "x = LeakyReLU(alpha=0.1)(x)\n",
    "x = MaxPooling2D(pool_size=(2, 2))(x)\n",
    "\n",
    "# Layer 6\n",
    "x = Conv2D(256, (3,3), strides=(1,1), padding='same', name='conv_6', use_bias=False)(x)\n",
    "x = BatchNormalization(name='norm_6')(x)\n",
    "x = LeakyReLU(alpha=0.1)(x)\n",
    "\n",
    "# Layer 7\n",
    "x = Conv2D(128, (1,1), strides=(1,1), padding='same', name='conv_7', use_bias=False)(x)\n",
    "x = BatchNormalization(name='norm_7')(x)\n",
    "x = LeakyReLU(alpha=0.1)(x)\n",
    "\n",
    "# Layer 8\n",
    "x = Conv2D(256, (3,3), strides=(1,1), padding='same', name='conv_8', use_bias=False)(x)\n",
    "x = BatchNormalization(name='norm_8')(x)\n",
    "x = LeakyReLU(alpha=0.1)(x)\n",
    "x = MaxPooling2D(pool_size=(2, 2))(x)\n",
    "\n",
    "# Layer 9\n",
    "x = Conv2D(512, (3,3), strides=(1,1), padding='same', name='conv_9', use_bias=False)(x)\n",
    "x = BatchNormalization(name='norm_9')(x)\n",
    "x = LeakyReLU(alpha=0.1)(x)\n",
    "\n",
    "# Layer 10\n",
    "x = Conv2D(256, (1,1), strides=(1,1), padding='same', name='conv_10', use_bias=False)(x)\n",
    "x = BatchNormalization(name='norm_10')(x)\n",
    "x = LeakyReLU(alpha=0.1)(x)\n",
    "\n",
    "# Layer 11\n",
    "x = Conv2D(512, (3,3), strides=(1,1), padding='same', name='conv_11', use_bias=False)(x)\n",
    "x = BatchNormalization(name='norm_11')(x)\n",
    "x = LeakyReLU(alpha=0.1)(x)\n",
    "\n",
    "# Layer 12\n",
    "x = Conv2D(256, (1,1), strides=(1,1), padding='same', name='conv_12', use_bias=False)(x)\n",
    "x = BatchNormalization(name='norm_12')(x)\n",
    "x = LeakyReLU(alpha=0.1)(x)\n",
    "\n",
    "# Layer 13\n",
    "x = Conv2D(512, (3,3), strides=(1,1), padding='same', name='conv_13', use_bias=False)(x)\n",
    "x = BatchNormalization(name='norm_13')(x)\n",
    "x = LeakyReLU(alpha=0.1)(x)\n",
    "\n",
    "skip_connection = x\n",
    "\n",
    "x = MaxPooling2D(pool_size=(2, 2))(x)\n",
    "\n",
    "# Layer 14\n",
    "x = Conv2D(1024, (3,3), strides=(1,1), padding='same', name='conv_14', use_bias=False)(x)\n",
    "x = BatchNormalization(name='norm_14')(x)\n",
    "x = LeakyReLU(alpha=0.1)(x)\n",
    "\n",
    "# Layer 15\n",
    "x = Conv2D(512, (1,1), strides=(1,1), padding='same', name='conv_15', use_bias=False)(x)\n",
    "x = BatchNormalization(name='norm_15')(x)\n",
    "x = LeakyReLU(alpha=0.1)(x)\n",
    "\n",
    "# Layer 16\n",
    "x = Conv2D(1024, (3,3), strides=(1,1), padding='same', name='conv_16', use_bias=False)(x)\n",
    "x = BatchNormalization(name='norm_16')(x)\n",
    "x = LeakyReLU(alpha=0.1)(x)\n",
    "\n",
    "# Layer 17\n",
    "x = Conv2D(512, (1,1), strides=(1,1), padding='same', name='conv_17', use_bias=False)(x)\n",
    "x = BatchNormalization(name='norm_17')(x)\n",
    "x = LeakyReLU(alpha=0.1)(x)\n",
    "\n",
    "# Layer 18\n",
    "x = Conv2D(1024, (3,3), strides=(1,1), padding='same', name='conv_18', use_bias=False)(x)\n",
    "x = BatchNormalization(name='norm_18')(x)\n",
    "x = LeakyReLU(alpha=0.1)(x)\n",
    "\n",
    "# Layer 19\n",
    "x = Conv2D(1024, (3,3), strides=(1,1), padding='same', name='conv_19', use_bias=False)(x)\n",
    "x = BatchNormalization(name='norm_19')(x)\n",
    "x = LeakyReLU(alpha=0.1)(x)\n",
    "\n",
    "# Layer 20\n",
    "x = Conv2D(1024, (3,3), strides=(1,1), padding='same', name='conv_20', use_bias=False)(x)\n",
    "x = BatchNormalization(name='norm_20')(x)\n",
    "x = LeakyReLU(alpha=0.1)(x)\n",
    "\n",
    "# Layer 21\n",
    "skip_connection = Conv2D(64, (1,1), strides=(1,1), padding='same', name='conv_21', use_bias=False)(skip_connection)\n",
    "skip_connection = BatchNormalization(name='norm_21')(skip_connection)\n",
    "skip_connection = LeakyReLU(alpha=0.1)(skip_connection)\n",
    "\n",
    "skip_connection = SpaceToDepth(block_size=2)(skip_connection)\n",
    "\n",
    "x = concatenate([skip_connection, x])\n",
    "\n",
    "# Layer 22\n",
    "x = Conv2D(1024, (3,3), strides=(1,1), padding='same', name='conv_22', use_bias=False)(x)\n",
    "x = BatchNormalization(name='norm_22')(x)\n",
    "x = LeakyReLU(alpha=0.1)(x)\n",
    "x = Dropout(0.7)(x) # add dropout\n",
    "\n",
    "# Layer 23\n",
    "x = Conv2D(BOX * (4 + 1 + CLASS), (1,1), strides=(1,1), padding='same', name='conv_23')(x)\n",
    "output = Reshape((GRID_H, GRID_W, BOX, 4 + 1 + CLASS))(x)\n",
    "\n",
    "model = keras.models.Model(input_image, output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "nF0cQvzseQSO",
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "model.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "M2q9Hs4XeQUh"
   },
   "source": [
    "# 2. Load YOLO pretrained weigts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "s9hVyf3seQVL"
   },
   "outputs": [],
   "source": [
    "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tABvrvEVeQVZ"
   },
   "outputs": [],
   "source": [
    "weight_reader = WeightReader('yolo.weights')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Bn8pfjvMeQVe"
   },
   "outputs": [],
   "source": [
    "weight_reader.reset()\n",
    "nb_conv = 23\n",
    "\n",
    "for i in range(1, nb_conv+1):\n",
    "    conv_layer = model.get_layer('conv_' + str(i))\n",
    "    conv_layer.trainable = True\n",
    "    \n",
    "    if i < nb_conv:\n",
    "        norm_layer = model.get_layer('norm_' + str(i))\n",
    "        norm_layer.trainable = True\n",
    "        \n",
    "        size = np.prod(norm_layer.get_weights()[0].shape)\n",
    "\n",
    "        beta  = weight_reader.read_bytes(size)\n",
    "        gamma = weight_reader.read_bytes(size)\n",
    "        mean  = weight_reader.read_bytes(size)\n",
    "        var   = weight_reader.read_bytes(size)\n",
    "\n",
    "        weights = norm_layer.set_weights([gamma, beta, mean, var])       \n",
    "        \n",
    "    if len(conv_layer.get_weights()) > 1:\n",
    "        bias   = weight_reader.read_bytes(np.prod(conv_layer.get_weights()[1].shape))\n",
    "        kernel = weight_reader.read_bytes(np.prod(conv_layer.get_weights()[0].shape))\n",
    "        kernel = kernel.reshape(list(reversed(conv_layer.get_weights()[0].shape)))\n",
    "        kernel = kernel.transpose([2,3,1,0])\n",
    "        conv_layer.set_weights([kernel, bias])\n",
    "    else:\n",
    "        kernel = weight_reader.read_bytes(np.prod(conv_layer.get_weights()[0].shape))\n",
    "        kernel = kernel.reshape(list(reversed(conv_layer.get_weights()[0].shape)))\n",
    "        kernel = kernel.transpose([2,3,1,0])\n",
    "        conv_layer.set_weights([kernel])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "V6aCwIGReQVn"
   },
   "outputs": [],
   "source": [
    "layer   = model.layers[-2] # last convolutional layer\n",
    "layer.trainable = True\n",
    "\n",
    "\n",
    "weights = layer.get_weights()\n",
    "\n",
    "new_kernel = np.random.normal(size=weights[0].shape)/(GRID_H*GRID_W)\n",
    "new_bias   = np.random.normal(size=weights[1].shape)/(GRID_H*GRID_W)\n",
    "\n",
    "layer.set_weights([new_kernel, new_bias])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "GsAii_vneQV4"
   },
   "source": [
    "# 3. Data generator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "sGgISdFOeQV7"
   },
   "outputs": [],
   "source": [
    "def parse_annotation(ann_dir, img_dir, labels):\n",
    "    '''\n",
    "    Parse XML files in PASCAL VOC format.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    - ann_dir : annotations files directory\n",
    "    - img_dir : images files directory\n",
    "    - labels : labels list\n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    - imgs_name : numpy array of images files path (shape : images count, 1)\n",
    "    - true_boxes : numpy array of annotations for each image (shape : image count, max annotation count, 5)\n",
    "        annotation format : xmin, ymin, xmax, ymax, class\n",
    "        xmin, ymin, xmax, ymax : image unit (pixel)\n",
    "        class = label index\n",
    "    '''\n",
    "    \n",
    "    max_annot = 0\n",
    "    imgs_name = []\n",
    "    annots = []\n",
    "    \n",
    "    # Parse file\n",
    "    for ann in sorted(os.listdir(ann_dir)):\n",
    "        annot_count = 0\n",
    "        boxes = []\n",
    "        try:\n",
    "            tree = ET.parse(ann_dir + ann)\n",
    "            for elem in tree.iter(): \n",
    "                if 'filename' in elem.tag:\n",
    "                    imgs_name.append(img_dir + elem.text)\n",
    "                if 'width' in elem.tag:\n",
    "                    w = int(elem.text)\n",
    "                if 'height' in elem.tag:\n",
    "                    h = int(elem.text)\n",
    "                if 'object' in elem.tag or 'part' in elem.tag:                  \n",
    "                    box = np.zeros((5))\n",
    "                    for attr in list(elem):\n",
    "                        if 'name' in attr.tag:\n",
    "                            box[4] = labels.index(attr.text) + 1 # 0:label for no bounding box\n",
    "                        if 'bndbox' in attr.tag:\n",
    "                            annot_count += 1\n",
    "                            for dim in list(attr):\n",
    "                                if 'xmin' in dim.tag:\n",
    "                                    box[0] = int(round(float(dim.text)))\n",
    "                                if 'ymin' in dim.tag:\n",
    "                                    box[1] = int(round(float(dim.text)))\n",
    "                                if 'xmax' in dim.tag:\n",
    "                                    box[2] = int(round(float(dim.text)))\n",
    "                                if 'ymax' in dim.tag:\n",
    "                                    box[3] = int(round(float(dim.text)))\n",
    "                    boxes.append(np.asarray(box))\n",
    "\n",
    "            if w != IMAGE_W or h != IMAGE_H :\n",
    "                print('Image size error')\n",
    "                break\n",
    "\n",
    "            annots.append(np.asarray(boxes))\n",
    "        except:\n",
    "            print(ann)\n",
    "\n",
    "        if annot_count > max_annot:\n",
    "            max_annot = annot_count\n",
    "           \n",
    "    # Rectify annotations boxes : len -> max_annot\n",
    "    imgs_name = np.array(imgs_name)  \n",
    "    true_boxes = np.zeros((imgs_name.shape[0], max_annot, 5))\n",
    "    for idx, boxes in enumerate(annots):\n",
    "        true_boxes[idx, :boxes.shape[0], :5] = boxes\n",
    "        \n",
    "    return imgs_name, true_boxes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Fv9lTCv2eQWQ"
   },
   "source": [
    "## 3.1. Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Nb7dHvVEeQWT"
   },
   "outputs": [],
   "source": [
    "def parse_function(img_obj, true_boxes):\n",
    "    x_img_string = tf.read_file(img_obj)\n",
    "    x_img = tf.image.decode_png(x_img_string, channels=3) # dtype=tf.uint8\n",
    "    x_img = tf.image.convert_image_dtype(x_img, tf.float32) # pixel value /255, dtype=tf.float32, channels : RGB\n",
    "    return x_img, true_boxes\n",
    "\n",
    "def get_dataset(img_dir, ann_dir, labels, batch_size):\n",
    "    '''\n",
    "    Create a YOLO dataset\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    - ann_dir : annotations files directory\n",
    "    - img_dir : images files directory\n",
    "    - labels : labels list\n",
    "    - batch_size : int\n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    - YOLO dataset : generate batch\n",
    "        batch : tupple(images, annotations)\n",
    "        batch[0] : images : tensor (shape : batch_size, IMAGE_W, IMAGE_H, 3)\n",
    "        batch[1] : annotations : tensor (shape : batch_size, max annot, 5)\n",
    "    Note : image pixel values = pixels value / 255. channels : RGB\n",
    "    '''\n",
    "    imgs_name, bbox = parse_annotation(ann_dir, img_dir, LABELS)\n",
    "    dataset = tf.data.Dataset.from_tensor_slices((imgs_name, bbox))    \n",
    "    dataset = dataset.shuffle(len(imgs_name))\n",
    "    dataset = dataset.repeat()\n",
    "    dataset = dataset.map(parse_function, num_parallel_calls=6)\n",
    "    dataset = dataset.batch(batch_size)\n",
    "    dataset = dataset.prefetch(10)\n",
    "    print('-------------------')\n",
    "    print('Dataset:')\n",
    "    print('Images count: {}'.format(len(imgs_name)))\n",
    "    print('Step per epoch: {}'.format(len(imgs_name) // batch_size))\n",
    "    print('Images per epoch: {}'.format(batch_size * (len(imgs_name) // batch_size)))\n",
    "    return dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "kFpvOoI5eQWd",
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "train_dataset = None\n",
    "train_dataset= get_dataset(train_image_folder, train_annot_folder, LABELS, TRAIN_BATCH_SIZE)\n",
    "\n",
    "val_dataset = None\n",
    "val_dataset= get_dataset(val_image_folder, val_annot_folder, LABELS, VAL_BATCH_SIZE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 361
    },
    "id": "KpCyhivLeQWt",
    "outputId": "2ac64459-1297-4854-d042-ac710dcab116"
   },
   "outputs": [],
   "source": [
    "# Test dataset\n",
    "\n",
    "def test_dataset(dataset):\n",
    "    for batch in dataset:\n",
    "        img = batch[0][0]\n",
    "        label = batch[1][0]\n",
    "        plt.figure(figsize=(2,2))\n",
    "        f, (ax1) = plt.subplots(1,1, figsize=(10, 10))\n",
    "        ax1.imshow(img)\n",
    "        ax1.set_title('Input image. Shape : {}'.format(img.shape))\n",
    "        for i in range(label.shape[0]):\n",
    "            box = label[i,:]\n",
    "            box = box.numpy()\n",
    "            x = box[0]\n",
    "            y = box[1]\n",
    "            w = box[2] - box[0]\n",
    "            h = box[3] - box[1]\n",
    "            if box[4] == 1:\n",
    "                color = (0, 1, 0)\n",
    "            else:\n",
    "                color = (1, 0, 0)\n",
    "            rect = patches.Rectangle((x, y), w, h, linewidth = 2, edgecolor=color,facecolor='none')\n",
    "            ax1.add_patch(rect)\n",
    "        break\n",
    "        \n",
    "test_dataset(train_dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_XU21szleQW6"
   },
   "source": [
    "## 3.2. Data augmentation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ukVO7ZcFeQXI"
   },
   "outputs": [],
   "source": [
    "def augmentation_generator(yolo_dataset):\n",
    "    '''\n",
    "    Augmented batch generator from a yolo dataset\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    - YOLO dataset\n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    - augmented batch : tensor (shape : batch_size, IMAGE_W, IMAGE_H, 3)\n",
    "        batch : tupple(images, annotations)\n",
    "        batch[0] : images : tensor (shape : batch_size, IMAGE_W, IMAGE_H, 3)\n",
    "        batch[1] : annotations : tensor (shape : batch_size, max annot, 5)\n",
    "    '''\n",
    "    for batch in yolo_dataset:\n",
    "        # conversion tensor->numpy\n",
    "        img = batch[0].numpy()\n",
    "        boxes = batch[1]. numpy()\n",
    "        # conversion bbox numpy->ia object\n",
    "        ia_boxes = []\n",
    "        for i in range(img.shape[0]):\n",
    "            ia_bbs = [ia.BoundingBox(x1=bb[0],\n",
    "                                       y1=bb[1],\n",
    "                                       x2=bb[2],\n",
    "                                       y2=bb[3]) for bb in boxes[i]\n",
    "                      if (bb[0] + bb[1] +bb[2] + bb[3] > 0)]\n",
    "            ia_boxes.append(ia.BoundingBoxesOnImage(ia_bbs, shape=(IMAGE_H, IMAGE_W, 3)))\n",
    "        # data augmentation\n",
    "        seq = iaa.Sequential([\n",
    "            iaa.Fliplr(0.5),\n",
    "            iaa.Flipud(0.5),\n",
    "            iaa.Multiply((0.4, 1.6)), # change brightness\n",
    "            #iaa.ContrastNormalization((0.5, 1.5)),\n",
    "            #iaa.Affine(translate_px={\"x\": (-100,100), \"y\": (-100,100)}, scale=(0.7, 1.30))\n",
    "            ])\n",
    "        #seq = iaa.Sequential([])\n",
    "        seq_det = seq.to_deterministic()\n",
    "        img_aug = seq_det.augment_images(img)\n",
    "        img_aug = np.clip(img_aug, 0, 1)\n",
    "        boxes_aug = seq_det.augment_bounding_boxes(ia_boxes)\n",
    "        # conversion ia object -> bbox numpy\n",
    "        for i in range(img.shape[0]):\n",
    "            boxes_aug[i] = boxes_aug[i].remove_out_of_image().cut_out_of_image()\n",
    "            for j, bb in enumerate(boxes_aug[i].bounding_boxes):\n",
    "                boxes[i,j,0] = bb.x1\n",
    "                boxes[i,j,1] = bb.y1\n",
    "                boxes[i,j,2] = bb.x2\n",
    "                boxes[i,j,3] = bb.y2\n",
    "        # conversion numpy->tensor\n",
    "        batch = (tf.convert_to_tensor(img_aug), tf.convert_to_tensor(boxes))\n",
    "        #batch = (img_aug, boxes)\n",
    "        yield batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "x5mZl3sleQXV"
   },
   "outputs": [],
   "source": [
    "aug_train_dataset = augmentation_generator(train_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "QmcrrBI0eQXZ",
    "outputId": "b5081b36-3a63-4fe1-9dcb-54b9d5a83e5e"
   },
   "outputs": [],
   "source": [
    "test_dataset(aug_train_dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hUhQfKFQeQXw"
   },
   "source": [
    "## 3.3. Process data to YOLO prediction format"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "8T_JPobpeQXz"
   },
   "outputs": [],
   "source": [
    "def process_true_boxes(true_boxes, anchors, image_width, image_height):\n",
    "    '''\n",
    "    Build image ground truth in YOLO format from image true_boxes and anchors.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    - true_boxes : tensor, shape (max_annot, 5), format : x1 y1 x2 y2 c, coords unit : image pixel\n",
    "    - anchors : list [anchor_1_width, anchor_1_height, anchor_2_width, anchor_2_height...]\n",
    "        anchors coords unit : grid cell\n",
    "    - image_width, image_height : int (pixels)\n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    - detector_mask : array, shape (GRID_W, GRID_H, anchors_count, 1)\n",
    "        1 if bounding box detected by grid cell, else 0\n",
    "    - matching_true_boxes : array, shape (GRID_W, GRID_H, anchors_count, 5)\n",
    "        Contains adjusted coords of bounding box in YOLO format\n",
    "    -true_boxes_grid : array, same shape than true_boxes (max_annot, 5),\n",
    "        format : x, y, w, h, c, coords unit : grid cell\n",
    "        \n",
    "    Note:\n",
    "    -----\n",
    "    Bounding box in YOLO Format : x, y, w, h, c\n",
    "    x, y : center of bounding box, unit : grid cell\n",
    "    w, h : width and height of bounding box, unit : grid cell\n",
    "    c : label index\n",
    "    ''' \n",
    "    \n",
    "    scale = IMAGE_W / GRID_W # scale = 32\n",
    "    \n",
    "    anchors_count = len(anchors) // 2\n",
    "    anchors = np.array(anchors)\n",
    "    anchors = anchors.reshape(len(anchors) // 2, 2)\n",
    "    \n",
    "    detector_mask = np.zeros((GRID_H, GRID_W, anchors_count, 1))\n",
    "    matching_true_boxes = np.zeros((GRID_H, GRID_W, anchors_count, 5))\n",
    "    \n",
    "    # convert true_boxes numpy array -> tensor\n",
    "    true_boxes = true_boxes.numpy()\n",
    "    \n",
    "    true_boxes_grid = np.zeros(true_boxes.shape)\n",
    "    \n",
    "    # convert bounding box coords and localize bounding box\n",
    "    for i, box in enumerate(true_boxes):\n",
    "        # convert box coords to x, y, w, h and convert to grids coord\n",
    "        w = (box[2] - box[0]) / scale\n",
    "        h = (box[3] - box[1]) / scale    \n",
    "        x = ((box[0] + box[2]) / 2) / scale\n",
    "        y = ((box[1] + box[3]) / 2) / scale\n",
    "        true_boxes_grid[i,...] = np.array([x, y, w, h, box[4]])\n",
    "        if w * h > 0: # box exists\n",
    "            # calculate iou between box and each anchors and find best anchors\n",
    "            best_iou = 0\n",
    "            best_anchor = 0\n",
    "            for i in range(anchors_count): \n",
    "                # iou (anchor and box are shifted to 0,0)\n",
    "                intersect = np.minimum(w, anchors[i,0]) * np.minimum(h, anchors[i,1])\n",
    "                union = (anchors[i,0] * anchors[i,1]) + (w * h) - intersect\n",
    "                iou = intersect / union\n",
    "                if iou > best_iou:\n",
    "                    best_iou = iou\n",
    "                    best_anchor = i\n",
    "            # localize box in detector_mask and matching true_boxes\n",
    "            if best_iou > 0:\n",
    "                x_coord = np.floor(x).astype('int')\n",
    "                y_coord = np.floor(y).astype('int')\n",
    "                detector_mask[y_coord, x_coord, best_anchor] = 1\n",
    "                yolo_box = np.array([x, y, w, h, box[4]])\n",
    "                matching_true_boxes[y_coord, x_coord, best_anchor] = yolo_box\n",
    "    return matching_true_boxes, detector_mask, true_boxes_grid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "21Jxs9OqeQYB"
   },
   "outputs": [],
   "source": [
    "def ground_truth_generator(dataset):\n",
    "    '''\n",
    "    Ground truth batch generator from a yolo dataset, ready to compare with YOLO prediction in loss function.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    - YOLO dataset. Generate batch:\n",
    "        batch : tupple(images, annotations)\n",
    "        batch[0] : images : tensor (shape : batch_size, IMAGE_W, IMAGE_H, 3)\n",
    "        batch[1] : annotations : tensor (shape : batch_size, max annot, 5)\n",
    "        \n",
    "    Returns\n",
    "    -------\n",
    "    - imgs : images to predict. tensor (shape : batch_size, IMAGE_H, IMAGE_W, 3)\n",
    "    - detector_mask : tensor, shape (batch, size, GRID_W, GRID_H, anchors_count, 1)\n",
    "        1 if bounding box detected by grid cell, else 0\n",
    "    - matching_true_boxes : tensor, shape (batch_size, GRID_W, GRID_H, anchors_count, 5)\n",
    "        Contains adjusted coords of bounding box in YOLO format\n",
    "    - class_one_hot : tensor, shape (batch_size, GRID_W, GRID_H, anchors_count, class_count)\n",
    "        One hot representation of bounding box label\n",
    "    - true_boxes_grid : annotations : tensor (shape : batch_size, max annot, 5)\n",
    "        true_boxes format : x, y, w, h, c, coords unit : grid cell\n",
    "    '''\n",
    "    for batch in dataset:\n",
    "        # imgs\n",
    "        imgs = batch[0]\n",
    "        \n",
    "        # true boxes\n",
    "        true_boxes = batch[1]\n",
    "        \n",
    "        # matching_true_boxes and detector_mask\n",
    "        batch_matching_true_boxes = []\n",
    "        batch_detector_mask = []\n",
    "        batch_true_boxes_grid = []\n",
    "        \n",
    "        for i in range(true_boxes.shape[0]):     \n",
    "            one_matching_true_boxes, one_detector_mask, true_boxes_grid = process_true_boxes(true_boxes[i],\n",
    "                                                                                           ANCHORS,\n",
    "                                                                                           IMAGE_H,\n",
    "                                                                                           IMAGE_W)\n",
    "            batch_matching_true_boxes.append(one_matching_true_boxes)\n",
    "            batch_detector_mask.append(one_detector_mask)\n",
    "            batch_true_boxes_grid.append(true_boxes_grid)\n",
    "                \n",
    "        detector_mask = tf.convert_to_tensor(np.array(batch_detector_mask), dtype='float32')\n",
    "        matching_true_boxes = tf.convert_to_tensor(np.array(batch_matching_true_boxes), dtype='float32')\n",
    "        true_boxes_grid = tf.convert_to_tensor(np.array(batch_true_boxes_grid), dtype='float32')\n",
    "        \n",
    "        # class one_hot\n",
    "        matching_classes = K.cast(matching_true_boxes[..., 4], 'int32') \n",
    "        class_one_hot = K.one_hot(matching_classes, CLASS + 1)[:,:,:,:,1:]\n",
    "        class_one_hot = tf.cast(class_one_hot, dtype='float32')\n",
    "        \n",
    "        batch = (imgs, detector_mask, matching_true_boxes, class_one_hot, true_boxes_grid)\n",
    "        yield batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "9px-bYDXeQYF"
   },
   "outputs": [],
   "source": [
    "# Ground true generator\n",
    "\n",
    "train_gen = ground_truth_generator(aug_train_dataset)\n",
    "val_gen = ground_truth_generator(val_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 530
    },
    "id": "Qrx4p1HReQYJ",
    "outputId": "c4a5bde4-229b-46fa-a532-336c9ea9bc36",
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# Test generator pipeline\n",
    "\n",
    "#model.load_weights('weights/training_do60_4_0.23566356.h5') # best weights, comment to start with YOLO weights\n",
    "\n",
    "# batch\n",
    "img, detector_mask, matching_true_boxes, class_one_hot, true_boxes = next(train_gen)\n",
    "\n",
    "# y\n",
    "matching_true_boxes = matching_true_boxes[0,...]\n",
    "detector_mask = detector_mask[0,...]\n",
    "class_one_hot = class_one_hot[0,...]\n",
    "y = K.concatenate((matching_true_boxes[...,0:4], detector_mask, class_one_hot), axis = -1)\n",
    "\n",
    "# y_hat\n",
    "y_hat = model.predict_on_batch(img)[0,...]\n",
    "\n",
    "# img\n",
    "img = img[0,...]\n",
    "\n",
    "# display prediction (Yolo Confidence value)\n",
    "plt.figure(figsize=(2,2))\n",
    "f, (ax1, ax2, ax3) = plt.subplots(1,3, figsize=(10, 10))\n",
    "ax1.imshow(img)\n",
    "ax1.set_title('Image')\n",
    "\n",
    "ax2.matshow((K.sum(y[:,:,:,4], axis=2))) # YOLO Confidence value\n",
    "ax2.set_title('Ground truth')\n",
    "ax2.xaxis.set_ticks_position('bottom')\n",
    "\n",
    "ax3.matshow(K.sum(y_hat[:,:,:,4], axis=2)) # YOLO Confidence value\n",
    "ax3.set_title('Prediction')\n",
    "ax3.xaxis.set_ticks_position('bottom')\n",
    "\n",
    "f.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7q_awL0EeQYO"
   },
   "source": [
    "# 4. Train\n",
    "\n",
    "## 4.1. Loss function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "sPne-C2veQYQ"
   },
   "outputs": [],
   "source": [
    "def iou(x1, y1, w1, h1, x2, y2, w2, h2):\n",
    "    '''\n",
    "    Calculate IOU between box1 and box2\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    - x, y : box center coords\n",
    "    - w : box width\n",
    "    - h : box height\n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    - IOU\n",
    "    '''   \n",
    "    xmin1 = x1 - 0.5*w1\n",
    "    xmax1 = x1 + 0.5*w1\n",
    "    ymin1 = y1 - 0.5*h1\n",
    "    ymax1 = y1 + 0.5*h1\n",
    "    xmin2 = x2 - 0.5*w2\n",
    "    xmax2 = x2 + 0.5*w2\n",
    "    ymin2 = y2 - 0.5*h2\n",
    "    ymax2 = y2 + 0.5*h2\n",
    "    interx = np.minimum(xmax1, xmax2) - np.maximum(xmin1, xmin2)\n",
    "    intery = np.minimum(ymax1, ymax2) - np.maximum(ymin1, ymin2)\n",
    "    inter = interx * intery\n",
    "    union = w1*h1 + w2*h2 - inter\n",
    "    iou = inter / (union + 1e-6)\n",
    "    return iou"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ROg1WywqeQYw"
   },
   "outputs": [],
   "source": [
    "# loss\n",
    "\n",
    "def yolov2_loss(detector_mask, matching_true_boxes, class_one_hot, true_boxes_grid, y_pred, info=False):\n",
    "    '''\n",
    "    Calculate YOLO V2 loss from prediction (y_pred) and ground truth tensors (detector_mask,\n",
    "    matching_true_boxes, class_one_hot, true_boxes_grid,)\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    - detector_mask : tensor, shape (batch, size, GRID_W, GRID_H, anchors_count, 1)\n",
    "        1 if bounding box detected by grid cell, else 0\n",
    "    - matching_true_boxes : tensor, shape (batch_size, GRID_W, GRID_H, anchors_count, 5)\n",
    "        Contains adjusted coords of bounding box in YOLO format\n",
    "    - class_one_hot : tensor, shape (batch_size, GRID_W, GRID_H, anchors_count, class_count)\n",
    "        One hot representation of bounding box label\n",
    "    - true_boxes_grid : annotations : tensor (shape : batch_size, max annot, 5)\n",
    "        true_boxes_grid format : x, y, w, h, c (coords unit : grid cell)\n",
    "    - y_pred : prediction from model. tensor (shape : batch_size, GRID_W, GRID_H, anchors count, (5 + labels count)\n",
    "    - info : boolean. True to get some info about loss value\n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    - loss : scalar\n",
    "    - sub_loss : sub loss list : coords loss, class loss and conf loss : scalar\n",
    "    '''\n",
    "    \n",
    "    # anchors tensor\n",
    "    anchors = np.array(ANCHORS)\n",
    "    anchors = anchors.reshape(len(anchors) // 2, 2)\n",
    "    \n",
    "    # grid coords tensor\n",
    "    coord_x = tf.to_float(tf.reshape(tf.tile(tf.range(GRID_W), [GRID_H]), (1, GRID_H, GRID_W, 1, 1)))\n",
    "    coord_y = tf.to_float(tf.reshape(tf.tile(tf.range(GRID_H), [GRID_W]), (1, GRID_W, GRID_H, 1, 1)))\n",
    "    coord_y = tf.transpose(coord_y, (0,2,1,3,4))\n",
    "    coords = tf.tile(tf.concat([coord_x,coord_y], -1), [TRAIN_BATCH_SIZE, 1, 1, 5, 1])\n",
    "    \n",
    "    # coordinate loss\n",
    "    pred_xy = K.sigmoid(y_pred[:,:,:,:,0:2]) # adjust coords between 0 and 1\n",
    "    pred_xy = (pred_xy + coords) # add cell coord for comparaison with ground truth. New coords in grid cell unit\n",
    "    pred_wh = K.exp(y_pred[:,:,:,:,2:4]) * anchors # adjust width and height for comparaison with ground truth. New coords in grid cell unit\n",
    "    nb_detector_mask = K.sum(tf.to_float(detector_mask > 0.0))\n",
    "    xy_loss = LAMBDA_COORD * K.sum(detector_mask * K.square(matching_true_boxes[...,:2] - pred_xy)) / (nb_detector_mask + 1e-6) # Non /2\n",
    "    wh_loss = LAMBDA_COORD * K.sum(detector_mask * K.square(K.sqrt(matching_true_boxes[...,2:4]) - \n",
    "                                                            K.sqrt(pred_wh))) / (nb_detector_mask + 1e-6)\n",
    "    coord_loss = xy_loss + wh_loss\n",
    "    \n",
    "    # class loss    \n",
    "    pred_box_class = y_pred[..., 5:]\n",
    "    true_box_class = tf.argmax(class_one_hot, -1)\n",
    "    class_loss = K.sparse_categorical_crossentropy(target=true_box_class, output=pred_box_class, from_logits=True)\n",
    "    class_loss = K.expand_dims(class_loss, -1) * detector_mask\n",
    "    class_loss = LAMBDA_CLASS * K.sum(class_loss) / (nb_detector_mask + 1e-6)\n",
    "    \n",
    "    # confidence loss\n",
    "    pred_conf = K.sigmoid(y_pred[...,4:5])\n",
    "    # for each detector : iou between prediction and ground truth\n",
    "    x1 = matching_true_boxes[...,0]\n",
    "    y1 = matching_true_boxes[...,1]\n",
    "    w1 = matching_true_boxes[...,2]\n",
    "    h1 = matching_true_boxes[...,3]\n",
    "    x2 = pred_xy[...,0]\n",
    "    y2 = pred_xy[...,1]\n",
    "    w2 = pred_wh[...,0]\n",
    "    h2 = pred_wh[...,1]\n",
    "    ious = iou(x1, y1, w1, h1, x2, y2, w2, h2)\n",
    "    ious = K.expand_dims(ious, -1)\n",
    "     \n",
    "    # for each detector : best ious between prediction and true_boxes (every bounding box of image)\n",
    "    pred_xy = K.expand_dims(pred_xy, 4) # shape : m, GRID_W, GRID_H, BOX, 1, 2 \n",
    "    pred_wh = K.expand_dims(pred_wh, 4)\n",
    "    pred_wh_half = pred_wh / 2.\n",
    "    pred_mins = pred_xy - pred_wh_half\n",
    "    pred_maxes = pred_xy + pred_wh_half\n",
    "    true_boxe_shape = K.int_shape(true_boxes_grid)\n",
    "    true_boxes_grid = K.reshape(true_boxes_grid, [true_boxe_shape[0], 1, 1, 1, true_boxe_shape[1], true_boxe_shape[2]])\n",
    "    true_xy = true_boxes_grid[...,0:2]\n",
    "    true_wh = true_boxes_grid[...,2:4]\n",
    "    true_wh_half = true_wh * 0.5\n",
    "    true_mins = true_xy - true_wh_half\n",
    "    true_maxes = true_xy + true_wh_half\n",
    "    intersect_mins = K.maximum(pred_mins, true_mins) # shape : m, GRID_W, GRID_H, BOX, max_annot, 2 \n",
    "    intersect_maxes = K.minimum(pred_maxes, true_maxes) # shape : m, GRID_W, GRID_H, BOX, max_annot, 2\n",
    "    intersect_wh = K.maximum(intersect_maxes - intersect_mins, 0.) # shape : m, GRID_W, GRID_H, BOX, max_annot, 1\n",
    "    intersect_areas = intersect_wh[..., 0] * intersect_wh[..., 1] # shape : m, GRID_W, GRID_H, BOX, max_annot, 1\n",
    "    pred_areas = pred_wh[..., 0] * pred_wh[..., 1] # shape : m, GRID_W, GRID_H, BOX, 1, 1\n",
    "    true_areas = true_wh[..., 0] * true_wh[..., 1] # shape : m, GRID_W, GRID_H, BOX, max_annot, 1\n",
    "    union_areas = pred_areas + true_areas - intersect_areas\n",
    "    iou_scores = intersect_areas / union_areas # shape : m, GRID_W, GRID_H, BOX, max_annot, 1\n",
    "    best_ious = K.max(iou_scores, axis=4)  # Best IOU scores.\n",
    "    best_ious = K.expand_dims(best_ious) # shape : m, GRID_W, GRID_H, BOX, 1\n",
    "    \n",
    "    # no object confidence loss\n",
    "    no_object_detection = K.cast(best_ious < 0.6, K.dtype(best_ious)) \n",
    "    noobj_mask = no_object_detection * (1 - detector_mask)\n",
    "    nb_noobj_mask  = K.sum(tf.to_float(noobj_mask  > 0.0))\n",
    "    \n",
    "    noobject_loss =  LAMBDA_NOOBJECT * K.sum(noobj_mask * K.square(-pred_conf)) / (nb_noobj_mask + 1e-6)\n",
    "    # object confidence loss\n",
    "    object_loss = LAMBDA_OBJECT * K.sum(detector_mask * K.square(ious - pred_conf)) / (nb_detector_mask + 1e-6)\n",
    "    # total confidence loss\n",
    "    conf_loss = noobject_loss + object_loss\n",
    "    \n",
    "    # total loss\n",
    "    loss = conf_loss + class_loss + coord_loss\n",
    "    sub_loss = [conf_loss, class_loss, coord_loss]  \n",
    "    \n",
    "#     # 'triple' mask\n",
    "#     true_box_conf_IOU = ious * detector_mask\n",
    "#     conf_mask = noobj_mask * LAMBDA_NOOBJECT\n",
    "#     conf_mask = conf_mask + detector_mask * LAMBDA_OBJECT\n",
    "#     nb_conf_box  = K.sum(tf.to_float(conf_mask  > 0.0))\n",
    "#     conf_loss = K.sum(K.square(true_box_conf_IOU - pred_conf) * conf_mask)  / (nb_conf_box  + 1e-6) \n",
    "    \n",
    "#     # total loss\n",
    "#     loss = conf_loss /2. + class_loss + coord_loss /2.\n",
    "#     sub_loss = [conf_loss /2., class_loss, coord_loss /2.]\n",
    "\n",
    "    if info:\n",
    "        print('conf_loss   : {:.4f}'.format(conf_loss))\n",
    "        print('class_loss  : {:.4f}'.format(class_loss))\n",
    "        print('coord_loss  : {:.4f}'.format(coord_loss))\n",
    "        print('    xy_loss : {:.4f}'.format(xy_loss))\n",
    "        print('    wh_loss : {:.4f}'.format(wh_loss))\n",
    "        print('--------------------')\n",
    "        print('total loss  : {:.4f}'.format(loss))\n",
    "        \n",
    "        # display masks for each anchors\n",
    "        for i in range(len(anchors)):\n",
    "            f, (ax1, ax2, ax3) = plt.subplots(1,3, figsize=(10, 5))\n",
    "            f.tight_layout()\n",
    "            f.suptitle('MASKS FOR ANCHOR {} :'.format(anchors[i,...]))\n",
    "            \n",
    "            ax1.matshow((K.sum(detector_mask[0,:,:,i], axis=2)), cmap='Greys', vmin=0, vmax=1)\n",
    "            ax1.set_title('detector_mask, count : {}'.format(K.sum(tf.to_int32(detector_mask[0,:,:,i]  > 0.))))\n",
    "            ax1.xaxis.set_ticks_position('bottom')\n",
    "            \n",
    "            ax2.matshow((K.sum(no_object_detection[0,:,:,i], axis=2)), cmap='Greys', vmin=0, vmax=1)\n",
    "            ax2.set_title('no_object_detection mask')\n",
    "            ax2.xaxis.set_ticks_position('bottom')\n",
    "            \n",
    "            ax3.matshow((K.sum(noobj_mask[0,:,:,i], axis=2)), cmap='Greys', vmin=0, vmax=1)\n",
    "            ax3.set_title('noobj_mask')\n",
    "            ax3.xaxis.set_ticks_position('bottom')\n",
    "              \n",
    "    return loss, sub_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 654
    },
    "id": "0EJqvrmheQY5",
    "outputId": "eed8d062-98f8-49cb-c5ee-f5aac3b41606",
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# test loss\n",
    "\n",
    "# get batch\n",
    "img, detector_mask, matching_true_boxes, class_one_hot, true_boxe_grid = next(train_gen)\n",
    "\n",
    "# first image in batch\n",
    "img = img[0:1]\n",
    "detector_mask = detector_mask[0:1]\n",
    "matching_true_boxes = matching_true_boxes[0:1]\n",
    "class_one_hot = class_one_hot[0:1]\n",
    "true_boxe_grid = true_boxe_grid[0:1]\n",
    "\n",
    "# predict\n",
    "y_pred = model.predict_on_batch(img)\n",
    "\n",
    "# plot img, ground truth and prediction\n",
    "f, (ax1, ax2, ax3) = plt.subplots(1,3, figsize=(10, 5))\n",
    "ax1.imshow(img[0,...])\n",
    "ax1.set_title('Image')\n",
    "ax2.matshow(K.sum(detector_mask[0,:,:,:,0], axis=2)) # YOLO Confidence value\n",
    "ax2.set_title('Ground truth, count : {}'.format(K.sum(tf.to_int32(detector_mask  > 0.))))\n",
    "ax2.xaxis.set_ticks_position('bottom')\n",
    "ax3.matshow(K.sum(y_pred[0,:,:,:,4], axis=2)) # YOLO Confidence value\n",
    "ax3.set_title('Prediction')\n",
    "ax3.xaxis.set_ticks_position('bottom')\n",
    "f.tight_layout()\n",
    "\n",
    "# loss info\n",
    "loss, sub_loss = yolov2_loss(detector_mask, matching_true_boxes, class_one_hot, true_boxe_grid, y_pred, info = True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "icqAjPgEeQZJ"
   },
   "outputs": [],
   "source": [
    "# gradients\n",
    "def grad(model, img, detector_mask, matching_true_boxes, class_one_hot, true_boxes, training=True):\n",
    "    with tf.GradientTape() as tape:\n",
    "        y_pred = model(img, training)\n",
    "        loss, sub_loss = yolov2_loss(detector_mask, matching_true_boxes, class_one_hot, true_boxes, y_pred)\n",
    "    return loss, sub_loss, tape.gradient(loss, model.trainable_variables)\n",
    "\n",
    "# save weights\n",
    "def save_best_weights(model, name, val_loss_avg):\n",
    "    # delete existing weights file\n",
    "    files = glob.glob(os.path.join('weights/', name + '*'))\n",
    "    for file in files:\n",
    "        os.remove(file)\n",
    "    # create new weights file\n",
    "    name = name + '_' + str(val_loss_avg) + '.h5'\n",
    "    path_name = os.path.join('weights/', name)\n",
    "    model.save_weights(path_name)\n",
    "\n",
    "# log (tensorboard)\n",
    "def log_loss(loss, val_loss):\n",
    "    with tf.contrib.summary.record_summaries_every_n_global_steps(1):\n",
    "        tf.contrib.summary.scalar('loss', loss)\n",
    "        tf.contrib.summary.scalar('val_loss', val_loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "fkUBWQQxeQZU"
   },
   "outputs": [],
   "source": [
    "# training\n",
    "def train(epochs, model, train_dataset, val_dataset, steps_per_epoch_train, steps_per_epoch_val, train_name = 'train'):\n",
    "    '''\n",
    "    Train YOLO model for n epochs.\n",
    "    Eval loss on training and validation dataset.\n",
    "    Log training loss and validation loss for tensorboard.\n",
    "    Save best weights during training (according to validation loss).\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    - epochs : integer, number of epochs to train the model.\n",
    "    - model : YOLO model.\n",
    "    - train_dataset : YOLO ground truth and image generator from training dataset.\n",
    "    - val_dataset : YOLO ground truth and image generator from validation dataset.\n",
    "    - steps_per_epoch_train : integer, number of batch to complete one epoch for train_dataset.\n",
    "    - steps_per_epoch_val : integer, number of batch to complete one epoch for val_dataset.\n",
    "    - train_name : string, training name used to log loss and save weights.\n",
    "    \n",
    "    Notes :\n",
    "    - train_dataset and val_dataset generate YOLO ground truth tensors : detector_mask,\n",
    "      matching_true_boxes, class_one_hot, true_boxes_grid. Shape of these tensors (batch size, tensor shape).\n",
    "    - steps per epoch = number of images in dataset // batch size of dataset\n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    - loss history : [train_loss_history, val_loss_history] : list of average loss for each epoch.\n",
    "    '''\n",
    "    num_epochs = epochs\n",
    "    steps_per_epoch_train = steps_per_epoch_train\n",
    "    steps_per_epoch_val = steps_per_epoch_val\n",
    "    train_loss_history = []\n",
    "    val_loss_history = []\n",
    "    best_val_loss = 1e6\n",
    "    \n",
    "    # optimizer\n",
    "    optimizer = tf.train.AdamOptimizer(learning_rate=1e-5, beta1=0.9, beta2=0.999, epsilon=1e-08)\n",
    "    global_step = tf.train.get_or_create_global_step()\n",
    "    \n",
    "    # log (tensorboard)\n",
    "    summary_writer = tf.contrib.summary.create_file_writer(os.path.join('logs/', train_name), flush_millis=20000)\n",
    "    summary_writer.set_as_default()\n",
    "    \n",
    "    # training\n",
    "    for epoch in range(num_epochs):\n",
    "        epoch_loss = []\n",
    "        epoch_val_loss = []\n",
    "        epoch_val_sub_loss = []\n",
    "        print('Epoch {} :'.format(epoch))\n",
    "        # train\n",
    "        for batch_idx in range(steps_per_epoch_train): \n",
    "            img, detector_mask, matching_true_boxes, class_one_hot, true_boxes =  next(train_dataset)\n",
    "            loss, _, grads = grad(model, img, detector_mask, matching_true_boxes, class_one_hot, true_boxes)\n",
    "            optimizer.apply_gradients(zip(grads, model.variables), global_step)\n",
    "            epoch_loss.append(loss)\n",
    "            print('-', end='')\n",
    "        print(' | ', end='')\n",
    "        # val\n",
    "        for batch_idx in range(steps_per_epoch_val): \n",
    "            img, detector_mask, matching_true_boxes, class_one_hot, true_boxes =  next(val_dataset)\n",
    "            loss, sub_loss, grads = grad(model, img, detector_mask, matching_true_boxes, class_one_hot, true_boxes, training=False)\n",
    "            epoch_val_loss.append(loss)\n",
    "            epoch_val_sub_loss.append(sub_loss)\n",
    "            print('-', end='')\n",
    "\n",
    "        loss_avg = np.mean(np.array(epoch_loss))\n",
    "        val_loss_avg = np.mean(np.array(epoch_val_loss))\n",
    "        sub_loss_avg = np.mean(np.array(epoch_val_sub_loss), axis=0)\n",
    "        train_loss_history.append(loss_avg)\n",
    "        val_loss_history.append(val_loss_avg)\n",
    "        \n",
    "        # log\n",
    "        log_loss(loss_avg, val_loss_avg)\n",
    "        \n",
    "        # save\n",
    "        if val_loss_avg < best_val_loss:\n",
    "            save_best_weights(model, train_name, val_loss_avg)\n",
    "            best_val_loss = val_loss_avg\n",
    "        \n",
    "        print(' loss = {:.4f}, val_loss = {:.4f} (conf={:.4f}, class={:.4f}, coords={:.4f})'.format(\n",
    "            loss_avg, val_loss_avg, sub_loss_avg[0], sub_loss_avg[1], sub_loss_avg[2]))\n",
    "        \n",
    "    return [train_loss_history, val_loss_history]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "qAiEBF2BeQZc",
    "outputId": "e93d5fd8-6de6-4f05-b05c-d8d1000d0b77",
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "results = train(EPOCHS, model, train_gen, val_gen, 10, 2, 'training_1')\n",
    "\n",
    "plt.plot(results[0])\n",
    "plt.plot(results[1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "FSltGdc6eQZl"
   },
   "source": [
    "# 5. Results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "nuKZ3pLReQZm",
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "def display_yolo(file, model, score_threshold, iou_threshold):\n",
    "    '''\n",
    "    Display predictions from YOLO model.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    - file : string list : list of images path.\n",
    "    - model : YOLO model.\n",
    "    - score_threshold : threshold used for filtering predicted bounding boxes.\n",
    "    - iou_threshold : threshold used for non max suppression.\n",
    "    '''\n",
    "    # load image\n",
    "    image = cv2.imread(file)\n",
    "\n",
    "    input_image = image[:,:,::-1]\n",
    "    input_image = image / 255.\n",
    "    input_image = np.expand_dims(input_image, 0)\n",
    "\n",
    "    # prediction\n",
    "    y_pred = model.predict_on_batch(input_image)\n",
    "\n",
    "    # post prediction process\n",
    "    # grid coords tensor\n",
    "    coord_x = tf.to_float(tf.reshape(tf.tile(tf.range(GRID_W), [GRID_H]), (1, GRID_H, GRID_W, 1, 1)))\n",
    "    coord_y = tf.to_float(tf.reshape(tf.tile(tf.range(GRID_H), [GRID_W]), (1, GRID_W, GRID_H, 1, 1)))\n",
    "    coord_y = tf.transpose(coord_y, (0,2,1,3,4))\n",
    "    coords = tf.tile(tf.concat([coord_x,coord_y], -1), [TRAIN_BATCH_SIZE, 1, 1, 5, 1])\n",
    "    dims = K.cast_to_floatx(K.int_shape(y_pred)[1:3])\n",
    "    dims = K.reshape(dims,(1,1,1,1,2))\n",
    "    # anchors tensor\n",
    "    anchors = np.array(ANCHORS)\n",
    "    anchors = anchors.reshape(len(anchors) // 2, 2)\n",
    "    # pred_xy and pred_wh shape (m, GRID_H, GRID_W, Anchors, 2)\n",
    "    pred_xy = K.sigmoid(y_pred[:,:,:,:,0:2])\n",
    "    pred_xy = (pred_xy + coords)\n",
    "    pred_xy = pred_xy / dims\n",
    "    pred_wh = K.exp(y_pred[:,:,:,:,2:4])\n",
    "    pred_wh = (pred_wh * anchors)\n",
    "    pred_wh = pred_wh / dims\n",
    "    # pred_confidence\n",
    "    box_conf = K.sigmoid(y_pred[:,:,:,:,4:5])  \n",
    "    # pred_class\n",
    "    box_class_prob = K.softmax(y_pred[:,:,:,:,5:])\n",
    "\n",
    "    # Reshape\n",
    "    pred_xy = pred_xy[0,...]\n",
    "    pred_wh = pred_wh[0,...]\n",
    "    box_conf = box_conf[0,...]\n",
    "    box_class_prob = box_class_prob[0,...]\n",
    "\n",
    "    # Convert box coords from x,y,w,h to x1,y1,x2,y2\n",
    "    box_xy1 = pred_xy - 0.5 * pred_wh\n",
    "    box_xy2 = pred_xy + 0.5 * pred_wh\n",
    "    boxes = K.concatenate((box_xy1, box_xy2), axis=-1)\n",
    "\n",
    "    # Filter boxes\n",
    "    box_scores = box_conf * box_class_prob\n",
    "    box_classes = K.argmax(box_scores, axis=-1) # best score index\n",
    "    box_class_scores = K.max(box_scores, axis=-1) # best score\n",
    "    prediction_mask = box_class_scores >= score_threshold\n",
    "    boxes = tf.boolean_mask(boxes, prediction_mask)\n",
    "    scores = tf.boolean_mask(box_class_scores, prediction_mask)\n",
    "    classes = tf.boolean_mask(box_classes, prediction_mask)\n",
    "\n",
    "    # Non Max Supression\n",
    "    selected_idx = tf.image.non_max_suppression(boxes, scores, 50, iou_threshold=iou_threshold)\n",
    "    boxes = K.gather(boxes, selected_idx)\n",
    "    scores = K.gather(scores, selected_idx)\n",
    "    classes = K.gather(classes, selected_idx)\n",
    "    \n",
    "    # Draw image\n",
    "    plt.figure(figsize=(2,2))\n",
    "    f, (ax1) = plt.subplots(1,1, figsize=(10, 10))\n",
    "    ax1.imshow(image[:,:,::-1])\n",
    "    count_detected = boxes.shape[0]\n",
    "    ax1.set_title('Detected objects count : {}'.format(count_detected))\n",
    "    for i in range(count_detected):\n",
    "        box = boxes[i,...]\n",
    "        x = box[0] * IMAGE_H\n",
    "        y = box[1] * IMAGE_W\n",
    "        w = (box[2] - box[0]) * IMAGE_H\n",
    "        h = (box[3] - box[1]) * IMAGE_W\n",
    "        classe = classes[i].numpy()\n",
    "        if classe == 0:\n",
    "            color = (0, 1, 0)\n",
    "        else:\n",
    "            color = (1, 0, 0)\n",
    "        rect = patches.Rectangle((x.numpy(), y.numpy()), w.numpy(), h.numpy(), linewidth = 3, edgecolor=color,facecolor='none')\n",
    "        ax1.add_patch(rect)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "iKbrGrdTeQZu",
    "outputId": "c8507363-3fb4-45b7-ac87-708a209ea0e5",
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "x_files =  glob.glob('data/val/image/*.png')\n",
    "\n",
    "score = SCORE_THRESHOLD\n",
    "iou_threshold = IOU_THRESHOLD\n",
    "\n",
    "score = 0.45\n",
    "iou_threshold = 0.2\n",
    "\n",
    "for file in x_files:\n",
    "    display_yolo(file, model, score, iou_threshold)"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [
    "_XU21szleQW6",
    "NkjrMqr_eQZI"
   ],
   "name": "Yolo_V2_tf_eager.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.7.6"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": false,
   "skip_h1_title": false,
   "title_cell": "Sommaire",
   "title_sidebar": "Contents",
   "toc_cell": true,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}