SL_end2end_predict.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import glob\n",
    "import os\n",
    "import cv2\n",
    "import pickle\n",
    "import tensorflow as tf"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### SegLink Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from sl_model import SL512, DSODSL512\n",
    "from sl_utils import PriorUtil\n",
    "from ssd_data import InputGenerator\n",
    "from ssd_data import preprocess\n",
    "\n",
    "from utils.model import load_weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "Model = SL512\n",
    "weights_path = './checkpoints/201809231008_sl512_synthtext/weights.002.h5'\n",
    "segment_threshold = 0.6; link_threshold = 0.25\n",
    "plot_name = 'sl512_crnn_sythtext'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "Model = DSODSL512\n",
    "weights_path = './checkpoints/201806021007_dsodsl512_synthtext/weights.012.h5'\n",
    "segment_threshold = 0.55; link_threshold = 0.45\n",
    "plot_name = 'dsodsl512_crnn_sythtext'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "sl_graph = tf.Graph()\n",
    "with sl_graph.as_default():\n",
    "    sl_session = tf.Session()\n",
    "    with sl_session.as_default():\n",
    "        model = Model()\n",
    "        prior_util = PriorUtil(model)\n",
    "        load_weights(model, weights_path)\n",
    "\n",
    "image_size = model.image_size"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### CRNN Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from crnn_model import CRNN\n",
    "from crnn_utils import alphabet87 as alphabet\n",
    "\n",
    "input_width = 256\n",
    "input_height = 32\n",
    "\n",
    "weights_path = './checkpoints/201806190711_crnn_gru_synthtext/weights.400000.h5'\n",
    "\n",
    "crnn_graph = tf.Graph()\n",
    "with crnn_graph.as_default():\n",
    "    crnn_session = tf.Session()\n",
    "    with crnn_session.as_default():\n",
    "        crnn_model = CRNN((input_width, input_height, 1), len(alphabet), prediction_only=True)\n",
    "        load_weights(crnn_model, weights_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Detection real world images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "inputs = []\n",
    "images = []\n",
    "images_orig = []\n",
    "data = []\n",
    "\n",
    "for img_path in glob.glob('data/images/test_images_seglink/*'):\n",
    "    img = cv2.imread(img_path)\n",
    "    images_orig.append(np.copy(img))\n",
    "    inputs.append(preprocess(img, image_size))\n",
    "    h, w = image_size\n",
    "    img = cv2.resize(img, (w,h), cv2.INTER_LINEAR).astype('float32') # should we do resizing\n",
    "    img = img[:, :, (2,1,0)] # BGR to RGB\n",
    "    img /= 255\n",
    "    images.append(img)\n",
    "    \n",
    "inputs = np.asarray(inputs)\n",
    "\n",
    "with sl_graph.as_default():\n",
    "    with sl_session.as_default():\n",
    "        preds = model.predict(inputs, batch_size=1, verbose=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "%%timeit\n",
    "for i in range(len(inputs)):\n",
    "    with sl_graph.as_default():\n",
    "        with sl_session.as_default():\n",
    "            preds = model.predict(inputs[i:i+1], batch_size=1, verbose=0)\n",
    "    #res = prior_util.decode(preds[0], segment_threshold, link_threshold)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Detection SynthText"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import pickle\n",
    "from data_synthtext import GTUtility\n",
    "\n",
    "file_name = 'gt_util_synthtext_seglink.pkl'\n",
    "with open(file_name, 'rb') as f:\n",
    "    gt_util = pickle.load(f)\n",
    "gt_util_train, gt_util_val = gt_util.split(0.9)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "idxs, inputs, images, data = gt_util_val.sample_random_batch(batch_size=32, input_size=image_size)\n",
    "\n",
    "images_orig = [cv2.imread(os.path.join(gt_util_val.image_path, gt_util_val.image_names[idx])) for idx in idxs]\n",
    "\n",
    "with sl_graph.as_default():\n",
    "    with sl_session.as_default():\n",
    "        preds = model.predict(inputs, batch_size=1, verbose=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Recognition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "from crnn_data import crop_words\n",
    "from crnn_utils import decode\n",
    "from sl_utils import rbox_to_polygon, polygon_to_rbox\n",
    "from utils.vis import plot_box, escape_latex\n",
    "\n",
    "#for k in range(len(preds)):\n",
    "for k in [0,2,3,9,10,11,24]:\n",
    "    plt.figure(figsize=[8]*2)\n",
    "    plt.imshow(images[k])\n",
    "    res = prior_util.decode(preds[k], segment_threshold, link_threshold)\n",
    "    \n",
    "    #print(res.shape)\n",
    "    \n",
    "    img = images_orig[k]\n",
    "    #mean = np.array([104,117,123])\n",
    "    #img -= mean[np.newaxis, np.newaxis, :]\n",
    "    rboxes = res[:,:5]\n",
    "    if len(rboxes) == 0:\n",
    "        plt.axis('off')\n",
    "        plt.show()\n",
    "        continue\n",
    "        \n",
    "    bh = rboxes[:,3]\n",
    "    rboxes[:,2] += bh * 0.1\n",
    "    rboxes[:,3] += bh * 0.2\n",
    "    \n",
    "    boxes = np.asarray([rbox_to_polygon(r) for r in rboxes])\n",
    "    boxes = np.flip(boxes, axis=1) # TODO: fix order of points, why?\n",
    "    boxes = np.reshape(boxes, (-1, 8))\n",
    "    \n",
    "    boxes_mask = np.array([not (np.any(b < 0-10) or np.any(b > 512+10)) for b in boxes]) # box inside image\n",
    "    #boxes_mask = np.logical_and(boxes_mask, [b[2] > 0.8*b[3] for b in rboxes]) # width > height, in square world\n",
    "    \n",
    "    boxes = boxes[boxes_mask]\n",
    "    rboxes = rboxes[boxes_mask]\n",
    "    if len(boxes) == 0:\n",
    "        boxes = np.empty((0,8))\n",
    "    \n",
    "    # plot boxes\n",
    "    for box in boxes:\n",
    "        c = 'rgby'\n",
    "        for i in range(4):\n",
    "            x, y = box[i*2:i*2+2]\n",
    "            plt.plot(x,y, c[i], marker='o', markersize=4)\n",
    "        plot_box(box, 'polygon')\n",
    "    \n",
    "    words = crop_words(img, np.clip(boxes/512,0,1), input_height, width=input_width, grayscale=True)\n",
    "    words = np.asarray([w.transpose(1,0,2) for w in words])\n",
    "    \n",
    "    if len(words) > 0:\n",
    "        with crnn_graph.as_default():\n",
    "            with crnn_session.as_default():\n",
    "                    res_crnn = crnn_model.predict(words)\n",
    "\n",
    "    #print('rboxes', len(rboxes), 'words', len(words), 'res_crnn', len(res_crnn))\n",
    "    for i in range(len(words)):\n",
    "        chars = [alphabet[c] for c in np.argmax(res_crnn[i], axis=1)]\n",
    "        \n",
    "        #gt_str = texts[i]\n",
    "        res_str = decode(chars)\n",
    "        \n",
    "        #ed = editdistance.eval(gt_str, res_str)\n",
    "        #ed = levenshtein(gt_str, res_str)\n",
    "        #ed_norm = ed / len(gt_str)\n",
    "        #mean_ed += ed\n",
    "        #mean_ed_norm += ed_norm\n",
    "        \n",
    "        #print('%-20s %s' % (res_str, ''.join(chars)))\n",
    "        #print('%s %-20s %0.2f' % (''.join(chars), res_str, res[i,5]))\n",
    "        \n",
    "        #print('%-20s %-20s %s %0.2f' % (\n",
    "        #    gt_str,\n",
    "        #    res_str,\n",
    "        #    ''.join(chars),\n",
    "        #    ed_norm))\n",
    "        x, y, w, h, theta = rboxes[i]\n",
    "        \n",
    "        #res_str = re.sub(r\"([#$%&_{}])\", r\"\\\\\\1\" , res_str)\n",
    "        #print(res_str, '   ', escape_latex(res_str))\n",
    "        \n",
    "        \n",
    "        #plt.text(x+h*np.sin(theta)/2, y+h*np.cos(theta)/2, escape_latex(res_str), rotation=theta/np.pi*180, \n",
    "        #         horizontalalignment='center', size='x-large' , color='cyan') # magenta, lime\n",
    "        plt.text(x+h*np.sin(theta)/2, y+h*np.cos(theta)/2, escape_latex(res_str), rotation=theta/np.pi*180, \n",
    "                 horizontalalignment='center', size='xx-large' , color='lime') # magenta, lime\n",
    "    \n",
    "    plt.axis('off')\n",
    "    \n",
    "    file_name = 'plots/%s_endtoend_realworld_%03i.pgf' % (plot_name, k)\n",
    "    #plt.savefig(file_name, bbox_inches='tight')\n",
    "    #print(file_name)\n",
    "    \n",
    "    plt.show()\n",
    "    \n",
    "    if False:\n",
    "        for i in range(len(words)):\n",
    "            plt.figure(figsize=[30,0.5])\n",
    "            plt.imshow(words[i][:,:,0].T, cmap='gray')\n",
    "            plt.axis('off')\n",
    "            plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "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.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}