2_line_word_segmentation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Line / Word segmentation\n",
    "This notebook demonstates a deep convolutional network method to provide bounding boxes surrounding handwritten texts.\n",
    "\n",
    "Input: An image containing handwritten text\n",
    "Output: A list of bounding boxes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import multiprocessing\n",
    "import os\n",
    "import random\n",
    "import time\n",
    "\n",
    "import cv2\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.patches as patches\n",
    "\n",
    "import mxnet as mx\n",
    "from mxnet import nd, autograd, gluon\n",
    "from mxnet.image import resize_short\n",
    "from mxboard import SummaryWriter\n",
    "from mxnet.gluon.model_zoo.vision import resnet34_v1\n",
    "from mxnet.contrib.ndarray import MultiBoxPrior, MultiBoxTarget, MultiBoxDetection, box_nms\n",
    "import numpy as np\n",
    "from skimage.draw import line_aa\n",
    "from skimage import transform as skimage_tf\n",
    "\n",
    "np.seterr(all='raise')\n",
    "\n",
    "mx.random.seed(42)\n",
    "\n",
    "from ocr.utils.iam_dataset import IAMDataset\n",
    "from ocr.utils.draw_box_on_image import draw_boxes_on_image"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Network definition\n",
    "We utilised an SSD network with horizonal anchor boxes to identify each line."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SSD(gluon.Block):\n",
    "    def __init__(self, num_classes, ctx, **kwargs):\n",
    "        super(SSD, self).__init__(**kwargs)\n",
    "\n",
    "        # Seven sets of anchor boxes are defined. For each set, n=2 sizes and m=3 ratios are defined.\n",
    "        # Four anchor boxes (n + m - 1) are generated: 2 square anchor boxes based on the n=2 sizes and 2 rectanges based on\n",
    "        # the sizes and the ratios. See https://discuss.mxnet.io/t/question-regarding-ssd-algorithm/1307 for more information.\n",
    "        \n",
    "        #self.anchor_sizes = [[.1, .2], [.2, .3], [.2, .4], [.4, .6], [.5, .7], [.6, .8], [.7, .9]]\n",
    "        #self.anchor_ratios = [[1, 3, 5], [1, 3, 5], [1, 6, 8], [1, 5, 7], [1, 6, 8], [1, 7, 9], [1, 7, 10]]\n",
    "\n",
    "        self.anchor_sizes = [[.1, .2], [.2, .3], [.2, .4], [.3, .4], [.3, .5], [.4, .6]]\n",
    "        self.anchor_ratios = [[1, 3, 5], [1, 3, 5], [1, 6, 8], [1, 4, 7], [1, 6, 8], [1, 5, 7]]\n",
    "\n",
    "        self.num_anchors = len(self.anchor_sizes)\n",
    "        self.num_classes = num_classes\n",
    "        self.ctx = ctx\n",
    "        with self.name_scope():\n",
    "            self.body, self.downsamples, self.class_preds, self.box_preds = self.get_ssd_model()\n",
    "            self.downsamples.initialize(mx.init.Normal(), ctx=self.ctx)\n",
    "            self.class_preds.initialize(mx.init.Normal(), ctx=self.ctx)\n",
    "            self.box_preds.initialize(mx.init.Normal(), ctx=self.ctx)\n",
    "\n",
    "    def get_body(self):\n",
    "        '''\n",
    "        Create the feature extraction network of the SSD based on resnet34.\n",
    "        The first layer of the res-net is converted into grayscale by averaging the weights of the 3 channels\n",
    "        of the original resnet.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        network: gluon.nn.HybridSequential\n",
    "            The body network for feature extraction based on resnet\n",
    "        \n",
    "        '''\n",
    "        pretrained = resnet34_v1(pretrained=True, ctx=self.ctx)\n",
    "        pretrained_2 = resnet34_v1(pretrained=True, ctx=mx.cpu(0))\n",
    "        first_weights = pretrained_2.features[0].weight.data().mean(axis=1).expand_dims(axis=1)\n",
    "        # First weights could be replaced with individual channels.\n",
    "        \n",
    "        body = gluon.nn.HybridSequential()\n",
    "        with body.name_scope():\n",
    "            first_layer = gluon.nn.Conv2D(channels=64, kernel_size=(7, 7), padding=(3, 3), strides=(2, 2), in_channels=1, use_bias=False)\n",
    "            first_layer.initialize(mx.init.Normal(), ctx=self.ctx)\n",
    "            first_layer.weight.set_data(first_weights)\n",
    "            body.add(first_layer)\n",
    "            body.add(*pretrained.features[1:-3])\n",
    "        return body\n",
    "\n",
    "    def get_class_predictor(self, num_anchors_predicted):\n",
    "        '''\n",
    "        Creates the category prediction network (takes input from each downsampled feature)\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        \n",
    "        num_anchors_predicted: int\n",
    "            Given n sizes and m ratios, the number of boxes predicted is n+m-1.\n",
    "            e.g., sizes=[.1, .2], ratios=[1, 3, 5] the number of anchors predicted is 4.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "\n",
    "        network: gluon.nn.HybridSequential\n",
    "            The class predictor network\n",
    "        '''\n",
    "        return gluon.nn.Conv2D(num_anchors_predicted*(self.num_classes + 1), kernel_size=3, padding=1)\n",
    "\n",
    "    def get_box_predictor(self, num_anchors_predicted):\n",
    "        '''\n",
    "        Creates the bounding box prediction network (takes input from each downsampled feature)\n",
    "        \n",
    "        Parameters\n",
    "        ----------\n",
    "        \n",
    "        num_anchors_predicted: int\n",
    "            Given n sizes and m ratios, the number of boxes predicted is n+m-1.\n",
    "            e.g., sizes=[.1, .2], ratios=[1, 3, 5] the number of anchors predicted is 4.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "\n",
    "        pred: gluon.nn.HybridSequential\n",
    "            The box predictor network\n",
    "        '''\n",
    "        pred = gluon.nn.HybridSequential()\n",
    "        with pred.name_scope():\n",
    "            pred.add(gluon.nn.Conv2D(channels=num_anchors_predicted*4, kernel_size=3, padding=1))\n",
    "        return pred\n",
    "\n",
    "    def get_down_sampler(self, num_filters):\n",
    "        '''\n",
    "        Creates a two-stacked Conv-BatchNorm-Relu and then a pooling layer to\n",
    "        downsample the image features by half.\n",
    "        '''\n",
    "        out = gluon.nn.HybridSequential()\n",
    "        for _ in range(2):\n",
    "            out.add(gluon.nn.Conv2D(num_filters, 3, strides=1, padding=1))\n",
    "            out.add(gluon.nn.BatchNorm(in_channels=num_filters))\n",
    "            out.add(gluon.nn.Activation('relu'))\n",
    "        out.add(gluon.nn.MaxPool2D(2))\n",
    "        out.hybridize()\n",
    "        return out\n",
    "\n",
    "    def get_ssd_model(self):\n",
    "        '''\n",
    "        Creates the SSD model that includes the image feature, downsample, category\n",
    "        and bounding boxes prediction networks.\n",
    "        '''\n",
    "        body = self.get_body()\n",
    "        downsamples = gluon.nn.HybridSequential()\n",
    "        class_preds = gluon.nn.HybridSequential()\n",
    "        box_preds = gluon.nn.HybridSequential()\n",
    "\n",
    "        downsamples.add(self.get_down_sampler(32))\n",
    "        downsamples.add(self.get_down_sampler(32))\n",
    "        downsamples.add(self.get_down_sampler(32))\n",
    "\n",
    "        for scale in range(self.num_anchors):\n",
    "            num_anchors_predicted = len(self.anchor_sizes[0]) + len(self.anchor_ratios[0]) - 1\n",
    "            class_preds.add(self.get_class_predictor(num_anchors_predicted))\n",
    "            box_preds.add(self.get_box_predictor(num_anchors_predicted))\n",
    "\n",
    "        return body, downsamples, class_preds, box_preds\n",
    "\n",
    "    def ssd_forward(self, x):\n",
    "        '''\n",
    "        Helper function of the forward pass of the sdd\n",
    "        '''\n",
    "        x = self.body(x)\n",
    "\n",
    "        default_anchors = []\n",
    "        predicted_boxes = []\n",
    "        predicted_classes = []\n",
    "\n",
    "        for i in range(self.num_anchors):\n",
    "            default_anchors.append(MultiBoxPrior(x, sizes=self.anchor_sizes[i], ratios=self.anchor_ratios[i]))\n",
    "            predicted_boxes.append(self._flatten_prediction(self.box_preds[i](x)))\n",
    "            predicted_classes.append(self._flatten_prediction(self.class_preds[i](x)))\n",
    "            if i < len(self.downsamples):\n",
    "                x = self.downsamples[i](x)\n",
    "            elif i == 3:\n",
    "                x = nd.Pooling(x, global_pool=True, pool_type='max', kernel=(4, 4))\n",
    "        return default_anchors, predicted_classes, predicted_boxes\n",
    "\n",
    "    def forward(self, x):\n",
    "        default_anchors, predicted_classes, predicted_boxes = self.ssd_forward(x)\n",
    "        # we want to concatenate anchors, class predictions, box predictions from different layers\n",
    "        anchors = nd.concat(*default_anchors, dim=1)\n",
    "        box_preds = nd.concat(*predicted_boxes, dim=1)\n",
    "        class_preds = nd.concat(*predicted_classes, dim=1)\n",
    "        class_preds = nd.reshape(class_preds, shape=(0, -1, self.num_classes + 1))\n",
    "        return anchors, class_preds, box_preds\n",
    "\n",
    "    def _flatten_prediction(self, pred):\n",
    "        '''\n",
    "        Helper function to flatten the predicted bounding boxes and categories\n",
    "        '''\n",
    "        return nd.flatten(nd.transpose(pred, axes=(0, 2, 3, 1)))\n",
    "\n",
    "    def training_targets(self, default_anchors, class_predicts, labels):\n",
    "        '''\n",
    "        Helper function to obtain the bounding boxes from the anchors.\n",
    "        '''\n",
    "        class_predicts = nd.transpose(class_predicts, axes=(0, 2, 1))\n",
    "        box_target, box_mask, cls_target = MultiBoxTarget(default_anchors, labels, class_predicts)\n",
    "        return box_target, box_mask, cls_target"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loss function\n",
    "The SSD was trained to minimize the classification error (box is text/box is not text) and the smoothed L1 loss between the predicted and actual bounding box. The smooth L1 loss is defined below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SmoothL1Loss(gluon.loss.Loss):\n",
    "    '''\n",
    "    A SmoothL1loss function defined in https://gluon.mxnet.io/chapter08_computer-vision/object-detection.html\n",
    "    '''\n",
    "    def __init__(self, batch_axis=0, **kwargs):\n",
    "        super(SmoothL1Loss, self).__init__(None, batch_axis, **kwargs)\n",
    "\n",
    "    def hybrid_forward(self, F, output, label, mask):\n",
    "        loss = F.smooth_l1((output - label) * mask, scalar=1.0)\n",
    "        return F.mean(loss, self._batch_axis, exclude=True)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data transform and augmentation\n",
    "\n",
    "The data is transformed so that the images and labels could be fed into the network. The training data is augmented so that the images are randomly translated and lines are randomly removed. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def augment_transform(image, label):\n",
    "    '''\n",
    "    1) Function that randomly translates the input image by +-width_range and +-height_range.\n",
    "    The labels (bounding boxes) are also translated by the same amount.\n",
    "    2) Each line can also be randomly removed for augmentation. Labels are also reduced to correspond to this\n",
    "    data and label are converted into tensors by calling the \"transform\" function.\n",
    "    '''\n",
    "    ty = random.uniform(-random_y_translation, random_y_translation)\n",
    "    tx = random.uniform(-random_x_translation, random_x_translation)\n",
    "\n",
    "    st = skimage_tf.SimilarityTransform(translation=(tx*image.shape[1], ty*image.shape[0]))\n",
    "    image = skimage_tf.warp(image, st, cval=1.0)\n",
    "\n",
    "    label[:, 0] = label[:, 0] - tx/2 #NOTE: Check why it has to be halfed (found experimentally)\n",
    "    label[:, 1] = label[:, 1] - ty/2\n",
    "    \n",
    "    index = np.random.uniform(0, 1.0, size=label.shape[0]) > random_remove_box\n",
    "    for i, should_output_bb in enumerate(index):\n",
    "        if should_output_bb == False:\n",
    "            (x, y, w, h) = label[i]\n",
    "            (x1, y1, x2, y2) = (x, y, x + w, y + h)\n",
    "            (x1, y1, x2, y2) = (x1 * image.shape[1], y1 * image.shape[0],\n",
    "                                x2 * image.shape[1], y2 * image.shape[0])\n",
    "            (x1, y1, x2, y2) = (int(x1), int(y1), int(x2), int(y2))\n",
    "            x1 = 0 if x1 < 0 else x1\n",
    "            y1 = 0 if y1 < 0 else y1\n",
    "            x2 = 0 if x2 < 0 else x2\n",
    "            y2 = 0 if y2 < 0 else y2\n",
    "            image_h, image_w = image.shape\n",
    "            x1 = image_w-1 if x1 >= image_w else x1\n",
    "            y1 = image_h-1 if y1 >= image_h else y1\n",
    "            x2 = image_w-1 if x2 >= image_w else x2\n",
    "            y2 = image_h-1 if y2 >= image_h else y2\n",
    "            image[y1:y2, x1:x2] = image[y1, x1]\n",
    "    \n",
    "    augmented_labels = label[index, :]\n",
    "    return transform(image*255., augmented_labels)\n",
    "\n",
    "def transform(image, label):\n",
    "    '''\n",
    "    Function that converts resizes image into the input image tensor for a CNN.\n",
    "    The labels (bounding boxes) are expanded, converted into (x, y, x+w, y+h), and\n",
    "    zero padded to the maximum number of labels. Finally, it is converted into a float\n",
    "    tensor.\n",
    "    '''\n",
    "    max_label_n = 128 if detection_box == \"word\" else 13\n",
    "\n",
    "    # Resize the image\n",
    "    image = np.expand_dims(image, axis=2)\n",
    "    image = mx.nd.array(image)\n",
    "    image = resize_short(image, image_size)\n",
    "    image = image.transpose([2, 0, 1])/255.\n",
    "\n",
    "    # Expand the bounding box by expand_bb_scale\n",
    "    bb = label.copy()\n",
    "    new_w = (1 + expand_bb_scale) * bb[:, 2]\n",
    "    new_h = (1 + expand_bb_scale) * bb[:, 3]\n",
    "    \n",
    "    bb[:, 0] = bb[:, 0] - (new_w - bb[:, 2])/2\n",
    "    bb[:, 1] = bb[:, 1] - (new_h - bb[:, 3])/2\n",
    "    bb[:, 2] = new_w\n",
    "    bb[:, 3] = new_h\n",
    "    label = bb \n",
    "\n",
    "    # Convert the predicted bounding box from (x, y, w, h to (x, y, x + w, y + h)\n",
    "    label = label.astype(np.float32)\n",
    "    label[:, 2] = label[:, 0] + label[:, 2]\n",
    "    label[:, 3] = label[:, 1] + label[:, 3]\n",
    "\n",
    "    # Zero pad the data\n",
    "    label_n = label.shape[0]\n",
    "    label_padded = np.zeros(shape=(max_label_n, 5))\n",
    "    label_padded[:label_n, 1:] = label\n",
    "    label_padded[:label_n, 0] = np.ones(shape=(1, label_n))\n",
    "    label_padded = mx.nd.array(label_padded)\n",
    "    return image, label_padded\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Helper function to generate output images for MXBoard"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_output_image(box_predictions, default_anchors, cls_probs, box_target, box_mask, cls_target, x, y):\n",
    "    '''\n",
    "    Generate the image with the predicted and actual bounding boxes.\n",
    "    Parameters\n",
    "    ----------\n",
    "    box_predictions: nd.array\n",
    "        Bounding box predictions relative to the anchor boxes, output of the network\n",
    "\n",
    "    default_anchors: nd.array\n",
    "        Anchors used, output of the network\n",
    "    \n",
    "    cls_probs: nd.array\n",
    "        Output of nd.SoftmaxActivation(nd.transpose(class_predictions, (0, 2, 1)), mode='channel')\n",
    "        where class_predictions is the output of the network.\n",
    "\n",
    "    box_target: nd.array\n",
    "        Output classification probabilities from network.training_targets(default_anchors, class_predictions, y)\n",
    "\n",
    "    box_mask: nd.array\n",
    "        Output bounding box predictions from network.training_targets(default_anchors, class_predictions, y) \n",
    "\n",
    "    cls_target: nd.array\n",
    "        Output targets from network.training_targets(default_anchors, class_predictions, y)\n",
    "    \n",
    "    x: nd.array\n",
    "       The input images\n",
    "\n",
    "    y: nd.array\n",
    "        The actual labels\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    output_image: np.array\n",
    "        The images with the predicted and actual bounding boxes drawn on\n",
    "\n",
    "    number_of_bbs: int\n",
    "        The number of predicting bounding boxes\n",
    "    '''\n",
    "    output = MultiBoxDetection(*[cls_probs, box_predictions, default_anchors], force_suppress=True, clip=False)\n",
    "    output = box_nms(output, overlap_thresh=overlap_thres, valid_thresh=min_c, topk=topk)\n",
    "    output = output.asnumpy()\n",
    "\n",
    "    number_of_bbs = 0\n",
    "    predicted_bb = []\n",
    "    for b in range(output.shape[0]):\n",
    "        predicted_bb_ = output[b, output[b, :, 0] != -1]\n",
    "        predicted_bb_ = predicted_bb_[:, 2:]\n",
    "        number_of_bbs += predicted_bb_.shape[0]\n",
    "        predicted_bb_[:, 2] = predicted_bb_[:, 2] - predicted_bb_[:, 0]\n",
    "        predicted_bb_[:, 3] = predicted_bb_[:, 3] - predicted_bb_[:, 1]\n",
    "        predicted_bb.append(predicted_bb_)\n",
    "        \n",
    "    labels = y[:, :, 1:].asnumpy()\n",
    "    labels[:, :, 2] = labels[:, :, 2] - labels[:, :, 0]\n",
    "    labels[:, :, 3] = labels[:, :, 3] - labels[:, :, 1]\n",
    "\n",
    "    output_image = draw_boxes_on_image(predicted_bb, labels, x.asnumpy())\n",
    "    output_image[output_image<0] = 0\n",
    "    output_image[output_image>1] = 1\n",
    "\n",
    "    return output_image, number_of_bbs\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_epoch(e, network, dataloader, trainer, log_dir, print_name, is_train, update_metric):\n",
    "    '''\n",
    "    Run one epoch to train or test the SSD network\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "        \n",
    "    e: int\n",
    "        The epoch number\n",
    "\n",
    "    network: nn.Gluon.HybridSequential\n",
    "        The SSD network\n",
    "\n",
    "    dataloader: gluon.data.DataLoader\n",
    "        The train or testing dataloader that is wrapped around the iam_dataset\n",
    "    \n",
    "    log_dir: Str\n",
    "        The directory to store the log files for mxboard\n",
    "\n",
    "    print_name: Str\n",
    "        Name to print for associating with the data. usually this will be \"train\" and \"test\"\n",
    "    \n",
    "    is_train: bool\n",
    "        Boolean to indicate whether or not the CNN should be updated. is_train should only be set to true for the training data\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "\n",
    "    network: gluon.nn.HybridSequential\n",
    "        The class predictor network\n",
    "    '''\n",
    "\n",
    "    total_losses = [0 for ctx_i in ctx]\n",
    "    for i, (X, Y) in enumerate(dataloader):\n",
    "        X = gluon.utils.split_and_load(X, ctx)\n",
    "        Y = gluon.utils.split_and_load(Y, ctx)\n",
    "        \n",
    "        with autograd.record(train_mode=is_train):\n",
    "            losses = []\n",
    "            for x, y in zip(X, Y):\n",
    "                default_anchors, class_predictions, box_predictions = network(x)\n",
    "                box_target, box_mask, cls_target = network.training_targets(default_anchors, class_predictions, y)\n",
    "                # losses\n",
    "                loss_class = cls_loss(class_predictions, cls_target)\n",
    "                loss_box = box_loss(box_predictions, box_target, box_mask)\n",
    "                # sum all losses\n",
    "                loss = loss_class + loss_box\n",
    "                losses.append(loss)\n",
    "            \n",
    "        if is_train:\n",
    "            for loss in losses:\n",
    "                loss.backward()\n",
    "            step_size = 0\n",
    "            for x in X:\n",
    "                step_size += x.shape[0]\n",
    "            trainer.step(step_size)\n",
    "\n",
    "        for index, loss in enumerate(losses):\n",
    "            total_losses[index] += loss.mean().asscalar()\n",
    "            \n",
    "        if update_metric:\n",
    "            cls_metric.update([cls_target], [nd.transpose(class_predictions, (0, 2, 1))])\n",
    "            box_metric.update([box_target], [box_predictions * box_mask])\n",
    "            \n",
    "        if i == 0 and e % send_image_every_n == 0 and e > 0:\n",
    "            cls_probs = nd.SoftmaxActivation(nd.transpose(class_predictions, (0, 2, 1)), mode='channel')\n",
    "            output_image, number_of_bbs = generate_output_image(box_predictions, default_anchors,\n",
    "                                                                cls_probs, box_target, box_mask,\n",
    "                                                                cls_target, x, y)\n",
    "            print(\"Number of predicted {} BBs = {}\".format(print_name, number_of_bbs))\n",
    "            with SummaryWriter(logdir=log_dir, verbose=False, flush_secs=5) as sw:\n",
    "                sw.add_image('bb_{}_image'.format(print_name), output_image, global_step=e)\n",
    "        \n",
    "\n",
    "    total_loss = 0\n",
    "    for loss in total_losses:\n",
    "        total_loss += loss / (len(dataloader)*len(total_losses))\n",
    "\n",
    "    with SummaryWriter(logdir=log_dir, verbose=False, flush_secs=5) as sw:\n",
    "        if update_metric:\n",
    "            name1, val1 = cls_metric.get()\n",
    "            name2, val2 = box_metric.get()\n",
    "            sw.add_scalar(name1, {\"test\": val1}, global_step=e)\n",
    "            sw.add_scalar(name2, {\"test\": val2}, global_step=e)\n",
    "        sw.add_scalar('loss', {print_name: total_loss}, global_step=e)\n",
    "            \n",
    "    return total_loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define default parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "detection_box = \"word\" #\"word\" or \"line\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "gpu_count = 4\n",
    "expand_bb_scale = 0.05\n",
    "min_c = 0.01\n",
    "overlap_thres = 0.1 if detection_box == \"line\" else 0.001\n",
    "topk = 150 if detection_box == \"line\" else 200\n",
    "\n",
    "epochs = 20\n",
    "learning_rate = 0.00005\n",
    "batch_size = 32\n",
    "image_size = 350\n",
    "\n",
    "random_x_translation, random_y_translation = (0.03, 0.03) if detection_box == \"line\" else (0.005, 0.005)\n",
    "random_remove_box = 0.1\n",
    "\n",
    "\n",
    "log_dir = \"./logs/line_word_segmentation\"\n",
    "checkpoint_dir, checkpoint_name = \"model_checkpoint\", \"ssd_\"+detection_box+\".params\"\n",
    "\n",
    "print_every_n = 5\n",
    "send_image_every_n = 20\n",
    "save_every_n = 50"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of training samples: 967\n",
      "Number of testing samples: 232\n"
     ]
    }
   ],
   "source": [
    "ctx = [mx.gpu(i) for i in range(gpu_count)]\n",
    "\n",
    "train_ds = IAMDataset(\"form_bb\", output_data=\"bb\", output_parse_method=detection_box, train=True)\n",
    "print(\"Number of training samples: {}\".format(len(train_ds)))\n",
    "\n",
    "test_ds = IAMDataset(\"form_bb\", output_data=\"bb\", output_parse_method=detection_box, train=False)\n",
    "print(\"Number of testing samples: {}\".format(len(test_ds)))\n",
    "\n",
    "train_data = gluon.data.DataLoader(train_ds.transform(augment_transform), batch_size, shuffle=True, last_batch=\"rollover\", num_workers=8)\n",
    "test_data = gluon.data.DataLoader(test_ds.transform(transform), batch_size, shuffle=False, last_batch=\"keep\", num_workers=8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "net = SSD(num_classes=2, ctx=ctx)\n",
    "net.hybridize()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "cls_loss = gluon.loss.SoftmaxCrossEntropyLoss()\n",
    "box_loss = SmoothL1Loss()\n",
    "cls_loss.hybridize()\n",
    "box_loss.hybridize()\n",
    "\n",
    "best_test_loss = 10e5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "if os.path.isfile(os.path.join(checkpoint_dir, checkpoint_name)):\n",
    "    net.load_parameters(os.path.join(checkpoint_dir, checkpoint_name), ctx=ctx)\n",
    "    print(\"Parameters loaded\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saving network, previous best test loss 1000000.000000, current test loss 0.741341\n",
      "Epoch 0, train_loss 0.862036, test_loss 0.741341, test accuracy=0.961873, mae=0.029940\n",
      "Saving network, previous best test loss 0.741341, current test loss 0.349544\n",
      "Saving network, previous best test loss 0.349544, current test loss 0.187104\n",
      "Saving network, previous best test loss 0.187104, current test loss 0.136065\n",
      "Saving network, previous best test loss 0.136065, current test loss 0.112224\n",
      "Saving network, previous best test loss 0.112224, current test loss 0.098453\n",
      "Epoch 5, train_loss 0.095902, test_loss 0.098453, test accuracy=0.986375, mae=0.022022\n",
      "Saving network, previous best test loss 0.098453, current test loss 0.090050\n",
      "Saving network, previous best test loss 0.090050, current test loss 0.083589\n",
      "Saving network, previous best test loss 0.083589, current test loss 0.078738\n",
      "Saving network, previous best test loss 0.078738, current test loss 0.074685\n",
      "Saving network, previous best test loss 0.074685, current test loss 0.072138\n",
      "Epoch 10, train_loss 0.065310, test_loss 0.072138, test accuracy=0.986463, mae=0.019575\n",
      "Saving network, previous best test loss 0.072138, current test loss 0.069216\n",
      "Saving network, previous best test loss 0.069216, current test loss 0.066652\n",
      "Saving network, previous best test loss 0.066652, current test loss 0.064299\n",
      "Saving network, previous best test loss 0.064299, current test loss 0.063023\n",
      "Saving network, previous best test loss 0.063023, current test loss 0.060886\n",
      "Epoch 15, train_loss 0.054017, test_loss 0.060886, test accuracy=0.987364, mae=0.017628\n",
      "Saving network, previous best test loss 0.060886, current test loss 0.059685\n",
      "Saving network, previous best test loss 0.059685, current test loss 0.058271\n",
      "Saving network, previous best test loss 0.058271, current test loss 0.057139\n",
      "Saving network, previous best test loss 0.057139, current test loss 0.056071\n"
     ]
    }
   ],
   "source": [
    "trainer = gluon.Trainer(net.collect_params(), 'adam', {'learning_rate': learning_rate})\n",
    "for e in range(epochs):\n",
    "    cls_metric = mx.metric.Accuracy()\n",
    "    box_metric = mx.metric.MAE()\n",
    "    train_loss = run_epoch(e, net, train_data, trainer, log_dir, print_name=\"train\", is_train=True, update_metric=False)\n",
    "    test_loss = run_epoch(e, net, test_data, trainer, log_dir, print_name=\"test\", is_train=False, update_metric=True)    \n",
    "    if test_loss < best_test_loss:\n",
    "        print(\"Saving network, previous best test loss {:.6f}, current test loss {:.6f}\".format(best_test_loss, test_loss))\n",
    "        net.save_parameters(os.path.join(checkpoint_dir, checkpoint_name))\n",
    "        best_test_loss = test_loss\n",
    "        \n",
    "    if e % print_every_n == 0:\n",
    "        name1, val1 = cls_metric.get()\n",
    "        name2, val2 = box_metric.get()\n",
    "        print(\"Epoch {0}, train_loss {1:.6f}, test_loss {2:.6f}, test {3}={4:.6f}, {5}={6:.6f}\".format(e, train_loss, test_loss, name1, val1, name2, val2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "if os.path.isfile(os.path.join(checkpoint_dir, checkpoint_name)):\n",
    "    net.load_parameters(os.path.join(checkpoint_dir, checkpoint_name), ctx=ctx)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Results"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Helper function to predict the bounding box"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict_bounding_boxes(net, image, bb):\n",
    "    '''\n",
    "    Given the outputs of the dataset (image and bounding box) and the network, \n",
    "    the predicted bounding boxes are provided.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    net: SSD\n",
    "    The trained SSD network.\n",
    "    \n",
    "    image: np.array\n",
    "    A grayscale image of the handwriting passages.\n",
    "    \n",
    "    bb: [(x1, y1, x2, y2)]\n",
    "    A tuple that contains the bounding box.\n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    predicted_bb: [(x, y, w, h)]\n",
    "    The predicted bounding boxes.\n",
    "    \n",
    "    actual_bb: [(x, y, w, h)]\n",
    "    The actual bounding bounding boxes.\n",
    "    '''\n",
    "    image, bb = transform(image, bb)\n",
    "\n",
    "    image = image.as_in_context(ctx[0])\n",
    "    image = image.expand_dims(axis=0)\n",
    "\n",
    "    bb = bb.as_in_context(ctx[0])\n",
    "    bb = bb.expand_dims(axis=0)\n",
    "\n",
    "    default_anchors, class_predictions, box_predictions = net(image)\n",
    "    box_target, box_mask, cls_target = net.training_targets(default_anchors, class_predictions, bb)\n",
    "    cls_probs = nd.SoftmaxActivation(nd.transpose(class_predictions, (0, 2, 1)), mode='channel')\n",
    "\n",
    "    predicted_bb = MultiBoxDetection(*[cls_probs, box_predictions, default_anchors], force_suppress=True, clip=False)\n",
    "    predicted_bb = box_nms(predicted_bb, overlap_thresh=overlap_thres, valid_thresh=min_c, topk=topk)\n",
    "    predicted_bb = predicted_bb.asnumpy()\n",
    "    predicted_bb = predicted_bb[0, predicted_bb[0, :, 0] != -1]\n",
    "    predicted_bb = predicted_bb[:, 2:]\n",
    "    predicted_bb[:, 2] = predicted_bb[:, 2] - predicted_bb[:, 0]\n",
    "    predicted_bb[:, 3] = predicted_bb[:, 3] - predicted_bb[:, 1]\n",
    "\n",
    "    labeled_bb = bb[:, :, 1:].asnumpy()\n",
    "    labeled_bb[:, :, 2] = labeled_bb[:, :, 2] - labeled_bb[:, :, 0]\n",
    "    labeled_bb[:, :, 3] = labeled_bb[:, :, 3] - labeled_bb[:, :, 1]\n",
    "    labeled_bb = labeled_bb[0]\n",
    "    return predicted_bb, labeled_bb"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Quantitative analysis\n",
    "The mean IOU was calculated for the train and test set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train iou 0.4409386990807991 test iou 0.4430850615806017\n"
     ]
    }
   ],
   "source": [
    "def get_iou(box1, box2):\n",
    "    '''\n",
    "    Calculate the IOU between two bounding boxes (x, y, w, h)\n",
    "    '''\n",
    "    # source: https://www.pyimagesearch.com/2016/11/07/intersection-over-union-iou-for-object-detection/\n",
    "    x1 = max(box1[0], box2[0])\n",
    "    y1 = max(box1[1], box2[1])\n",
    "    x2 = min(box1[2], box2[2])\n",
    "    y2 = min(box1[3], box2[3])\n",
    "\n",
    "    inter_area = max(0, x2 - x1 + 1) * max(0, y2 - y1 + 1)\n",
    "\n",
    "    box1_area = (box1[2] - box1[0] + 1) * (box1[3] - box1[1] + 1)\n",
    "    box2_area = (box2[2] - box2[0] + 1) * (box2[3] - box2[1] + 1)\n",
    "\n",
    "    iou = inter_area / float(box1_area + box2_area - inter_area)\n",
    "    return iou\n",
    "\n",
    "def calculate_iou(dataset):\n",
    "    '''\n",
    "    Iterate through the dataset and calculate the mean IOU between the actual and \n",
    "    predicted bounding boxes.\n",
    "    '''\n",
    "    ious = []\n",
    "    for i in range(len(dataset)):\n",
    "        iou_i = 0.0\n",
    "        count_i = 0\n",
    "        image, bb = dataset[i]\n",
    "        predicted_bb, actual_bb = predict_bounding_boxes(net, image, bb)\n",
    "        # A naive 1-1 bouding box matching algorithm was used. This algorithm\n",
    "        # doesn't account for insertions or deletes and may result in lower IOU values.\n",
    "        for predicted_bb_i, actual_bb_i in zip(predicted_bb, actual_bb):\n",
    "            iou = get_iou(predicted_bb_i, actual_bb_i)\n",
    "            iou_i += iou\n",
    "            count_i += 1\n",
    "        ious.append(iou_i/count_i)\n",
    "    return np.mean(ious)\n",
    "        \n",
    "train_iou = calculate_iou(train_ds)\n",
    "test_iou = calculate_iou(test_ds)\n",
    "print(\"Train iou {} test iou {}\".format(train_iou, test_iou))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Qualitative analysis"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Visualise the predicted and actual bounding boxes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "figs_to_plot = 2\n",
    "fig, axs = plt.subplots(figs_to_plot, 2, figsize=(15, 10 * figs_to_plot))\n",
    "ds = test_ds\n",
    "for i in range(figs_to_plot):\n",
    "    n = int(random.random()*len(ds))\n",
    "    image, bb = ds[n]\n",
    "    predicted_bb, actual_bb = predict_bounding_boxes(net, image, bb)\n",
    "    \n",
    "    for j in range(actual_bb.shape[0]):\n",
    "        (x, y, w, h) = actual_bb[j]\n",
    "        axs[i][0].imshow(image, cmap='Greys_r')\n",
    "        image_h, image_w = image.shape[-2:]\n",
    "        (x, y, w, h) = (x * image_w, y * image_h, w * image_w, h * image_h)\n",
    "        rect = patches.Rectangle((x, y), w, h, fill=False, color=\"r\")\n",
    "        axs[i][0].add_patch(rect)\n",
    "        axs[i][0].set_title(\"BB actual\")\n",
    "\n",
    "    for j in range(predicted_bb.shape[0]):\n",
    "        axs[i][1].imshow(image, cmap='Greys_r')\n",
    "        (x, y, w, h) = predicted_bb[j]\n",
    "        image_h, image_w = image.shape[-2:]\n",
    "        (x, y, w, h) = (x * image_w, y * image_h, w * image_w, h * image_h)\n",
    "        rect = patches.Rectangle((x, y), w, h, fill=False, color=\"r\")\n",
    "        axs[i][1].add_patch(rect)\n",
    "        axs[i][1].set_title(\"BB predicted\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}