word_classifier_CTC.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training: Connectionist temporal classification - Word Classification\n",
    "Training of CTC model for words recognition\n",
    "## TODO\n",
    "```\n",
    "char2idx change? Indexs?\n",
    "Add propper accuracy\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import cv2\n",
    "import matplotlib as plt\n",
    "import tensorflow as tf\n",
    "import os\n",
    "from imgaug import augmenters as iaa\n",
    "\n",
    "sys.path.append('src')\n",
    "from ocr.datahelpers import load_words_data, char2idx, CHAR_SIZE\n",
    "from ocr.dataiterator import BucketDataIterator\n",
    "from ocr.helpers import img_extend, resize\n",
    "from ocr.mlhelpers import TrainingPlot\n",
    "from ocr.tfhelpers import create_cell"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib notebook\n",
    "plt.rcParams['figure.figsize'] = (9.0, 5.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Loading images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading words...\n",
      "data/sets/test.csv\n",
      " |████████████████████████████████████████| 100.0% \n",
      "-> Number of words: 19076\n",
      "Loading words...\n",
      "data/sets/dev.csv\n",
      " |████████████████████████████████████████| 100.0% \n",
      "-> Number of words: 19076\n"
     ]
    }
   ],
   "source": [
    "train_images, train_labels = load_words_data('data/sets/train.csv', is_csv=True)\n",
    "dev_images, dev_labels = load_words_data('data/sets/dev.csv', is_csv=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Settings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "slider_size = (64, 64)      # Second parameter can be edited\n",
    "\n",
    "layers = 4\n",
    "residual_layers = 0\n",
    "units = 512\n",
    "\n",
    "num_buckets = 10\n",
    "N_INPUT = slider_size[0]*slider_size[1]\n",
    "vocab_size = CHAR_SIZE + 2         # Number of different chars + <PAD> and <EOS>\n",
    "\n",
    "learning_rate = 1e-4               # 1e-4\n",
    "dropout = 0.4\n",
    "\n",
    "TRAIN_STEPS = 20000\n",
    "TEST_ITER = 150\n",
    "SAVE_ITER = 500\n",
    "EPOCH = 500\n",
    "BATCH_SIZE = 64\n",
    "\n",
    "model_name = 'Classifier1'\n",
    "save_loc = 'models/word-clas/CTC/'\n",
    "summaries_dir = 'logs/word-clas/CTC/' + model_name\n",
    "\n",
    "if not os.path.exists(save_loc):\n",
    "    os.makedirs(save_loc)\n",
    "save_loc += model_name"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training images: 19076\n",
      "Testing images: 19076\n"
     ]
    }
   ],
   "source": [
    "data = {\n",
    "    'train': (train_images, train_labels, np.empty(len(train_labels), dtype=object)),\n",
    "    'dev': (dev_images, dev_labels, np.empty(len(dev_labels), dtype=object))\n",
    "}\n",
    "\n",
    "for d in ['train', 'dev']:\n",
    "    for i in range(len(data[d][0])):\n",
    "        data[d][0][i] = resize(data[d][0][i], slider_size[1], True)\n",
    "        data[d][2][i] = [char2idx(c) for c in data[d][1][i]]\n",
    "\n",
    "print(\"Training images:\", len(train_images))\n",
    "print(\"Testing images:\", len(dev_images))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Data Augmentation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "seq = iaa.Sequential([\n",
    "    iaa.Sometimes(\n",
    "        0.3,\n",
    "        iaa.ElasticTransformation(alpha=(0.5, 10.0), sigma=5.0)),\n",
    "    iaa.OneOf([\n",
    "        iaa.GaussianBlur((0, 0.5)),\n",
    "        iaa.AverageBlur(k=(1, 3)),\n",
    "        iaa.MedianBlur(k=(1, 3)),\n",
    "    ]),\n",
    "    iaa.Sometimes(\n",
    "        0.3,\n",
    "        iaa.AdditiveGaussianNoise(scale=0.01*255)),\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DataIterator(BucketDataIterator):\n",
    "    def next_feed(self, batch_size):\n",
    "        \"\"\"Create feed directly for model training.\"\"\"\n",
    "        (inputs_,\n",
    "         targets_,\n",
    "         inputs_length_) = self.next_batch(batch_size)\n",
    "        return {\n",
    "            inputs: inputs_,\n",
    "            inputs_length: inputs_length_,\n",
    "            targets: targets_,\n",
    "            keep_prob: (1.0 - self.dropout) if self.train else 1.0\n",
    "        }       "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Iterator created.\n",
      "Iterator created.\n"
     ]
    }
   ],
   "source": [
    "train_iterator = DataIterator(\n",
    "    data['train'][0],\n",
    "    data['train'][2],\n",
    "    num_buckets,\n",
    "    slider_size,\n",
    "    augmentation=seq,\n",
    "    dropout=dropout,\n",
    "    train=True)\n",
    "test_iterator = DataIterator(\n",
    "    data['dev'][0],\n",
    "    data['dev'][2],\n",
    "    1,\n",
    "    slider_size,\n",
    "    train=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Placeholders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Input placehodlers\n",
    "# N_INPUT -> size of vector representing one image in sequence\n",
    "# Inputs - Time major: (max_seq_length, batch_size, vec_size)\n",
    "inputs = tf.placeholder(shape=(None, slider_size[0], None, 1),\n",
    "                        dtype=tf.float32,\n",
    "                        name='inputs')\n",
    "inputs_length = tf.placeholder(shape=(None,),\n",
    "                               dtype=tf.int32,\n",
    "                               name='inputs_length')\n",
    "targets = tf.sparse_placeholder(dtype=tf.int32,\n",
    "                                name='targets')\n",
    "keep_prob = tf.placeholder(tf.float32, name='keep_prob')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Graph"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### CNN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SCALE = 0.0003\n",
    "h_n_pool = 16 # = 2^(number of pooling layers on height)\n",
    "w_n_pool = 8  # = 2^(number of pooling layers on width)\n",
    "# 1. Convulation\n",
    "conv1 = tf.layers.conv2d(\n",
    "    inputs=inputs,\n",
    "    filters=64,\n",
    "    kernel_size=[7, 7],\n",
    "    strides=(2, 2),\n",
    "    padding=\"same\",\n",
    "    kernel_initializer=tf.contrib.layers.xavier_initializer(),\n",
    "    kernel_regularizer=tf.contrib.layers.l2_regularizer(scale=SCALE),\n",
    "    activation=tf.nn.relu)\n",
    "    \n",
    "conv12 = tf.layers.conv2d(\n",
    "    inputs=conv1,\n",
    "    filters=64,\n",
    "    kernel_size=[5, 5],\n",
    "    strides=(1, 1),\n",
    "    padding=\"same\",\n",
    "    kernel_initializer=tf.contrib.layers.xavier_initializer(),\n",
    "    kernel_regularizer=tf.contrib.layers.l2_regularizer(scale=SCALE),\n",
    "    activation=tf.nn.relu)\n",
    "# 2. Max Pool\n",
    "pool1 = tf.layers.max_pooling2d(conv12, pool_size=[2, 2], strides=2)\n",
    "# 3. Inception\n",
    "conv2 = tf.layers.conv2d(\n",
    "    inputs=pool1,\n",
    "    filters=64,\n",
    "    kernel_size=[5, 5],\n",
    "    strides=(1, 1),\n",
    "    padding=\"same\",\n",
    "    kernel_initializer=tf.contrib.layers.xavier_initializer(),\n",
    "    kernel_regularizer=tf.contrib.layers.l2_regularizer(scale=SCALE),\n",
    "    activation=tf.nn.relu)\n",
    "    \n",
    "conv22 = tf.layers.conv2d(\n",
    "    inputs=conv2,\n",
    "    filters=128,\n",
    "    kernel_size=[5, 5],\n",
    "    strides=(1, 1),\n",
    "    padding=\"same\",\n",
    "    kernel_initializer=tf.contrib.layers.xavier_initializer(),\n",
    "    kernel_regularizer=tf.contrib.layers.l2_regularizer(scale=SCALE),\n",
    "    activation=tf.nn.relu)\n",
    "# 4. Max Pool\n",
    "pool2 = tf.layers.max_pooling2d(conv22, pool_size=[2, 1], strides=[2, 1])\n",
    "# 5. Inception\n",
    "conv3 = tf.layers.conv2d(\n",
    "    inputs=pool2,\n",
    "    filters=128,\n",
    "    kernel_size=[5, 5],\n",
    "    strides=(1, 1),\n",
    "    padding=\"same\",\n",
    "    kernel_initializer=tf.contrib.layers.xavier_initializer(),\n",
    "    kernel_regularizer=tf.contrib.layers.l2_regularizer(scale=SCALE),\n",
    "    activation=tf.nn.relu)\n",
    "conv32 = tf.layers.conv2d(\n",
    "    inputs=conv3,\n",
    "    filters=256,\n",
    "    kernel_size=[5, 5],\n",
    "    strides=(1, 1),\n",
    "    padding=\"same\",\n",
    "    kernel_initializer=tf.contrib.layers.xavier_initializer(),\n",
    "    kernel_regularizer=tf.contrib.layers.l2_regularizer(scale=SCALE),\n",
    "    activation=tf.nn.relu)\n",
    "# 6. Max Pool\n",
    "pool3 = tf.layers.max_pooling2d(conv32, pool_size=[2, 2], strides=2)\n",
    "\n",
    "# Image patches for RNN\n",
    "image_patches1 = tf.layers.conv2d(\n",
    "    inputs=pool3,\n",
    "    filters=256,\n",
    "    kernel_size=[slider_size[0]//(h_n_pool), slider_size[1]//(w_n_pool)],\n",
    "    strides=(1, 1),\n",
    "    padding=\"same\",\n",
    "    kernel_initializer=tf.contrib.layers.xavier_initializer(),\n",
    "    kernel_regularizer=tf.contrib.layers.l2_regularizer(scale=SCALE),\n",
    "    activation=tf.nn.relu)\n",
    "\n",
    "image_patches = tf.layers.separable_conv2d(\n",
    "    inputs=image_patches1,\n",
    "    filters=1280,\n",
    "    kernel_size=[slider_size[0]//(h_n_pool), 1],\n",
    "    strides=(1, 1),\n",
    "    depth_multiplier=5,\n",
    "    name='image_patches')\n",
    "\n",
    "processed_inputs = tf.transpose(\n",
    "    tf.squeeze(image_patches, [1]),\n",
    "    [1, 0, 2],\n",
    "    name='squeeze_transpose')\n",
    "lengths = inputs_length//w_n_pool"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Bidirectional RNN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "enc_cell_fw = create_cell(units,\n",
    "                          layers,\n",
    "                          residual_layers,\n",
    "                          is_dropout=True,\n",
    "                          keep_prob=keep_prob)\n",
    "enc_cell_bw = create_cell(units,\n",
    "                          layers,\n",
    "                          residual_layers,\n",
    "                          is_dropout=True,\n",
    "                          keep_prob=keep_prob)\n",
    "bi_outputs, _ = tf.nn.bidirectional_dynamic_rnn(\n",
    "    cell_fw=enc_cell_fw,\n",
    "    cell_bw=enc_cell_bw,\n",
    "    inputs=processed_inputs,\n",
    "    sequence_length=lengths,\n",
    "    dtype=tf.float32,\n",
    "    time_major=True)\n",
    "\n",
    "con_outputs = tf.concat(bi_outputs, -1)\n",
    "logits = tf.layers.dense(inputs=con_outputs, units=vocab_size)\n",
    "\n",
    "# Final outputs \n",
    "decoded, log_prob = tf.nn.ctc_beam_search_decoder(\n",
    "    logits, lengths, merge_repeated=False)\n",
    "word_prediction = tf.sparse_tensor_to_dense(decoded[0], name='word_prediction')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Optimizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "ctc_loss = tf.nn.ctc_loss(\n",
    "        targets,\n",
    "        logits,\n",
    "        lengths,\n",
    "        time_major=True,\n",
    "        ctc_merge_repeated=True,\n",
    "        ignore_longer_outputs_than_inputs=True)\n",
    "regularization = tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES)\n",
    "loss = tf.identity(tf.reduce_mean(ctc_loss) + sum(regularization), name='loss')\n",
    "\n",
    "optimizer = tf.train.AdamOptimizer(learning_rate)\n",
    "train_step = optimizer.minimize(loss, name='train_step')\n",
    "\n",
    "# Label error rate\n",
    "label_err_rate = tf.reduce_mean(\n",
    "    tf.edit_distance(tf.cast(decoded[0], tf.int32), targets))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Accuracy + Padding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "# pred_length = decoded[0].get_shape().as_list()[1]\n",
    "# targets_length = targets.get_shape().as_list()[1]\n",
    "\n",
    "# pad_lenght = tf.maximum(pred_length, targets_length)\n",
    "\n",
    "# pred_pad = tf.pad(\n",
    "#     word_prediction,\n",
    "#     [[0, 0],\n",
    "#      [0, pad_lenght - pred_length]],\n",
    "#     constant_values=PAD,\n",
    "#     mode='CONSTANT')\n",
    "# targets_pad = tf.pad(\n",
    "#     test_targets,\n",
    "#     [[0, 0],\n",
    "#      [0, pad_lenght - targets_lenght]],\n",
    "#     constant_values=PAD,\n",
    "#     mode='CONSTANT')\n",
    "\n",
    "# acc_weights = tf.cast(tf.divide(pred_pad, pred_pad), tf.float32)\n",
    "\n",
    "# # acc_weights = tf.sequence_mask(\n",
    "# #     tf.subtract(final_seq_lengths, 1),    # word_inputs_length, try max(targets, inputs)\n",
    "# #     word_pad_lenght,\n",
    "# #     dtype=tf.float32)\n",
    "\n",
    "# correct_prediction = tf.equal(pred_pad, targets_pad)\n",
    "# accuracy = (tf.reduce_sum(tf.cast(correct_prediction, tf.float32) * acc_weights) \\\n",
    "#             / tf.reduce_sum(acc_weights))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "## Training\n",
    "sess = tf.Session()\n",
    "sess.run(tf.global_variables_initializer())\n",
    "saver = tf.train.Saver()\n",
    "\n",
    "# TesorBoard stats\n",
    "tf.summary.scalar('Loss', loss)\n",
    "tf.summary.scalar('Label_err_rate', label_err_rate)\n",
    "merged = tf.summary.merge_all()\n",
    "train_writer = tf.summary.FileWriter(summaries_dir + '/train',\n",
    "                                     sess.graph)\n",
    "test_writer = tf.summary.FileWriter(summaries_dir + '/test')\n",
    "\n",
    "# I recommend using tensorboard, comment this out if you want\n",
    "trainPlot = TrainingPlot(TRAIN_STEPS, TEST_ITER, LOSS_ITER)\n",
    "\n",
    "try:\n",
    "    for i_batch in range(TRAIN_STEPS):\n",
    "        fd = train_iterator.next_feed(BATCH_SIZE)\n",
    "        sess.run(train_step, feed_dict=fd)\n",
    "        \n",
    "        if i_batch % LOSS_ITER == 0:\n",
    "            # Plotting loss\n",
    "            tmpLoss = loss.eval(fd)\n",
    "            trainPlot.updateCost(tmpLoss, i_batch // LOSS_ITER)\n",
    "    \n",
    "        if i_batch % TEST_ITER == 0:\n",
    "            # Plotting accuracy\n",
    "            fd_test = test_iterator.next_feed(BATCH_SIZE)\n",
    "            \n",
    "            accTest = accuracy.eval(fd_test)\n",
    "            accTrain = accuracy.eval(fd)\n",
    "            trainPlot.updateAcc(accTest, accTrain, i_batch // TEST_ITER)\n",
    "            \n",
    "            testSummary = sess.run(merged, feed_dict=fd_test)\n",
    "            trainSummary = sess.run(merged, feed_dict=fd)\n",
    "            test_writer.add_summary(testSummary, i_batch)\n",
    "            train_writer.add_summary(trainSummary, i_batch)\n",
    "\n",
    "        if i_batch % SAVE_ITER == 0:\n",
    "            saver.save(sess, save_loc, global_step=i_batch)\n",
    "        \n",
    "        if i_batch % EPOCH == 0:\n",
    "            fd_test = test_iterator.next_feed(BATCH_SIZE) # BATCH_SIZE)\n",
    "            print('Batch %r - Loss: %r' % (i_batch, sess.run(loss, fd)))\n",
    "            print('    Train Label Err Rate: %r' % sess.run(label_err_rate,\n",
    "                                                            feed_dict=fd))\n",
    "            print('    Eval Label Err Rate: %r' % sess.run(label_err_rate,\n",
    "                                                           feed_dict=fd_test))\n",
    "            print()\n",
    "\n",
    "except KeyboardInterrupt:\n",
    "    print('Stopped on batch:', i_batch)\n",
    "    saver.save(sess, save_loc)\n",
    "    print('Training interrupted, model saved.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(5):\n",
    "    fd_test = test_iterator.next_feed(BATCH_SIZE)\n",
    "    predict_, target_ = sess.run([word_prediction, test_targets], fd_test)\n",
    "    for i, (inp, pred) in enumerate(zip(target_, predict_)):\n",
    "        print('    expected  > {}'.format(inp))\n",
    "        print('    predicted > {}'.format(pred))\n",
    "        if i >= 1:\n",
    "            break\n",
    "    print()"
   ]
  }
 ],
 "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}