TBPP_end2end_predict_GPUonly.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "import cv2\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = '0'\n",
    "\n",
    "np.set_printoptions(linewidth=120, precision=5, suppress=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "from data_synthtext import GTUtility\n",
    "\n",
    "#file_name = 'gt_util_synthtext_seglink.pkl'\n",
    "file_name = 'gt_util_synthtext_seglink_xxs.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": {},
   "outputs": [],
   "source": [
    "batch_size = 16\n",
    "image_size = (512, 512)\n",
    "\n",
    "idxs, inputs, images, data = gt_util_val.sample_random_batch(batch_size=batch_size, input_size=image_size)\n",
    "\n",
    "images = np.array(images, dtype='float32')\n",
    "\n",
    "images_orig = [cv2.imread(os.path.join(gt_util_val.image_path, gt_util_val.image_names[idx])) for idx in idxs]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Detection model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tbpp_model import TBPP512, TBPP512_dense\n",
    "from tbpp_utils import PriorUtil\n",
    "\n",
    "det_model = TBPP512_dense()\n",
    "prior_util = PriorUtil(det_model)\n",
    "\n",
    "weights_path = './checkpoints/201807091503_dsodtbpp512fl_synthtext/weights.018.h5'\n",
    "confidence_threshold = 0.25\n",
    "\n",
    "det_model.load_weights(weights_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Recognition model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from crnn_model import CRNN\n",
    "from crnn_utils import alphabet87 as alphabet\n",
    "\n",
    "input_width = 1024*2\n",
    "input_height = 32\n",
    "\n",
    "weights_path = './checkpoints/202002030820_crnn_cnn_synthtext_concat_continued/weights.200000.h5'\n",
    "\n",
    "rec_model = CRNN((input_width, input_height, 1), len(alphabet), gru=False, cnn=True, prediction_only=True)\n",
    "rec_model.load_weights(weights_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Detection + Recognition model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from keras.models import Model\n",
    "from keras.layers import Input, Activation, Lambda\n",
    "\n",
    "from tbpp_layers import TBPPDecodeAndCrop\n",
    "\n",
    "x_in = Input(shape=det_model.input_shape[1:])\n",
    "x = Lambda(lambda x: x*255)(x_in)\n",
    "x_in_det = Lambda(lambda x: x-[104,117,123])(x)\n",
    "x_det = det_model(x_in_det)\n",
    "x_in_rec, x_det_dec = TBPPDecodeAndCrop(prior_util,\n",
    "                                        output_size=(input_height, input_width),\n",
    "                                        confidence_threshold=0.7,\n",
    "                                        iou_threshold=0.45,\n",
    "                                        top_k=200)([x, x_det])\n",
    "x_rec = rec_model(x_in_rec)\n",
    "\n",
    "model = Model([x_in,], [x_det_dec, x_in_rec, x_rec])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "from crnn_utils import decode\n",
    "\n",
    "res = model.predict([images])\n",
    "print([r.shape for r in res])\n",
    "\n",
    "for i in range(batch_size):\n",
    "    im = res[1][i]\n",
    "    im = np.transpose(im, (1,0,2))\n",
    "    plt.figure(figsize=(16,2))\n",
    "    plt.imshow(im[:,:1000,0]/255, vmin=0, vmax=1, cmap='gray')\n",
    "    plt.show()\n",
    "    \n",
    "    chars = [alphabet[c] for c in np.argmax(res[2][i], axis=1)]\n",
    "    print(''.join(chars))\n",
    "    \n",
    "    offsets = [0] + [int(e//4) for e in res[0][i][:,19] if e > 0]\n",
    "    print(offsets)\n",
    "    res_strs = [ decode(chars[offsets[i]:offsets[i+1]]) for i in range(len(offsets)-1) ]\n",
    "    res_strs = [ s.strip() for s in res_strs ]\n",
    "    print(res_strs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit -n 10 -r 10\n",
    "res = model.predict([images[:1]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "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.8.10"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}