SL_384x512.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import tensorflow as tf\n",
    "import keras\n",
    "import cv2\n",
    "import os, time, pickle\n",
    "from tqdm.notebook import tqdm\n",
    "\n",
    "from ssd_data import InputGenerator\n",
    "from sl_utils import PriorUtil\n",
    "from sl_training import SegLinkLoss, SegLinkFocalLoss\n",
    "from sl_metric import evaluate_results\n",
    "from ssd_metric import fscore\n",
    "from utils.model import load_weights, count_parameters, calc_memory_usage\n",
    "from utils.training import Logger, LearningRateDecay, MetricUtility"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "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)\n",
    "gt_util_train, _ = gt_util.split(0.25)\n",
    "gt_util_val, _ = gt_util.split(0.25)\n",
    "\n",
    "print(gt_util)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sl_model import SL384x512_dense\n",
    "model = SL384x512_dense()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(len(model.source_layers)):\n",
    "    l = model.source_layers[i]\n",
    "    print('%-2s %s' %(i, l.get_shape().as_list()))\n",
    "print()\n",
    "count_parameters(model)\n",
    "calc_memory_usage(model)\n",
    "\n",
    "c = 0\n",
    "for l in model.layers:\n",
    "    if l.__class__.__name__ == \"Conv2D\":\n",
    "        c += 1\n",
    "print(c)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ssd_data import preprocess\n",
    "\n",
    "inputs = []\n",
    "images = []\n",
    "data = []\n",
    "\n",
    "gtu = gt_util_val\n",
    "image_size = model.image_size\n",
    "\n",
    "np.random.seed(1337)\n",
    "\n",
    "for i in [0]: #np.random.randint(0, gtu.num_samples, 16):\n",
    "\n",
    "    img_path = os.path.join(gtu.image_path, gtu.image_names[i])\n",
    "    img = cv2.imread(img_path)\n",
    "    print('img_shape', img.shape)\n",
    "    \n",
    "    image_size_cv = image_size[::-1]\n",
    "    print('image_size_ssd', image_size)\n",
    "    print('image_size_cv ', image_size_cv)\n",
    "    inpt = preprocess(img, image_size_cv)\n",
    "    inputs.append(inpt)\n",
    "    print('image_size_inp', inpt.shape)\n",
    "    \n",
    "    img = cv2.resize(img, image_size_cv, cv2.INTER_LINEAR).astype('float32') # should we do resizing\n",
    "    print('image_size_img', img.shape)\n",
    "    img = img[:, :, (2,1,0)] # BGR to RGB\n",
    "    img /= 255\n",
    "    images.append(img)\n",
    "    \n",
    "    boxes = gtu.data[i]\n",
    "    data.append(boxes)\n",
    "    \n",
    "    print()\n",
    "\n",
    "inputs = np.asarray(inputs)\n",
    "\n",
    "test_idx = 0\n",
    "test_input = inputs[test_idx]\n",
    "test_img = images[test_idx]\n",
    "test_gt = data[test_idx]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "### Encoding/Decoding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prior_util = PriorUtil(model)\n",
    "\n",
    "plt.figure(figsize=[12,9])\n",
    "plt.axis('off')\n",
    "plt.axis('equal')\n",
    "plt.imshow(test_img)\n",
    "\n",
    "test_encoded_gt = prior_util.encode(test_gt, debug=False)\n",
    "\n",
    "loc_idxs = list(range(1000))\n",
    "\n",
    "for m_idx in [5]:\n",
    "#for m_idx in [0,1,2,3,4,5]:\n",
    "    #prior_util.prior_maps[m_idx-1].plot_locations()\n",
    "    m = prior_util.prior_maps[m_idx]\n",
    "    m.plot_locations()\n",
    "    #m.plot_boxes(loc_idxs)\n",
    "    #prior_util.plot_neighbors(m_idx, loc_idxs, cross_layer=False)\n",
    "    prior_util.plot_neighbors(m_idx, loc_idxs, inter_layer=False)\n",
    "    prior_util.plot_assignment(m_idx)\n",
    "    \n",
    "plt.show()\n",
    "\n",
    "dummy_output = np.copy(test_encoded_gt)\n",
    "#dummy_output[:,2:4] += np.random.randn(*dummy_output[:,2:4].shape)*0.05\n",
    "\n",
    "plt.figure(figsize=[12,9])\n",
    "ax = plt.gca()\n",
    "plt.imshow(test_img)\n",
    "res = prior_util.decode(dummy_output, debug=False, debug_combining=True)\n",
    "#res = decode(prior_util, dummy_output, debug=False)\n",
    "prior_util.plot_gt()\n",
    "prior_util.plot_results(res)\n",
    "plt.axis('off'); plt.xlim(0, image_size[1]); plt.ylim(image_size[0],0)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "epochs = 100\n",
    "initial_epoch = 0\n",
    "batch_size = 6\n",
    "freeze = []\n",
    "experiment = 'sl384x512_synthtext'\n",
    "\n",
    "prior_util = PriorUtil(model)\n",
    "\n",
    "#optimizer = tf.optimizers.SGD(learning_rate=1e-3, momentum=0.9, decay=0, nesterov=True)\n",
    "optimizer = tf.optimizers.Adam(learning_rate=1e-3, beta_1=0.9, beta_2=0.999, epsilon=0.001, decay=0.0)\n",
    "\n",
    "#loss = SegLinkLoss(lambda_offsets=1.0, lambda_links=1.0, neg_pos_ratio=3.0)\n",
    "loss = SegLinkFocalLoss(lambda_segments=100.0, lambda_offsets=1.0, lambda_links=100.0,\n",
    "                        gamma_segments=2, gamma_links=2, first_map_size=(96,128))\n",
    "\n",
    "#regularizer = None\n",
    "regularizer = keras.regularizers.l2(5e-4) # None if disabled\n",
    "\n",
    "gen_train = InputGenerator(gt_util_train, prior_util, batch_size, model.image_size, augmentation=False)\n",
    "gen_val = InputGenerator(gt_util_val, prior_util, batch_size, model.image_size, augmentation=False)\n",
    "\n",
    "\n",
    "dataset_train, dataset_val = gen_train.get_dataset(), gen_val.get_dataset()\n",
    "iterator_train, iterator_val = iter(dataset_train), iter(dataset_val)\n",
    "\n",
    "checkdir = './checkpoints/' + time.strftime('%Y%m%d%H%M') + '_' + experiment\n",
    "\n",
    "if not os.path.exists(checkdir):\n",
    "    os.makedirs(checkdir)\n",
    "\n",
    "with open(checkdir+'/source.py','wb') as f:\n",
    "    source = ''.join(['# In[%i]\\n%s\\n\\n' % (i, In[i]) for i in range(len(In))])\n",
    "    f.write(source.encode())\n",
    "\n",
    "print(checkdir)\n",
    "\n",
    "for l in model.layers:\n",
    "    l.trainable = not l.name in freeze\n",
    "    if regularizer and l.__class__.__name__.startswith('Conv'):\n",
    "        model.add_loss(lambda l=l: regularizer(l.kernel))\n",
    "\n",
    "metric_util = MetricUtility(loss.metric_names, logdir=checkdir)\n",
    "\n",
    "@tf.function\n",
    "def step(x, y_true, training=False):\n",
    "    if training:\n",
    "        with tf.GradientTape() as tape:\n",
    "            y_pred = model(x, training=True)\n",
    "            metric_values = loss.compute(y_true, y_pred)\n",
    "            total_loss = metric_values['loss']\n",
    "            if len(model.losses):\n",
    "                total_loss += tf.add_n(model.losses)\n",
    "        gradients = tape.gradient(total_loss, model.trainable_variables)\n",
    "        optimizer.apply_gradients(zip(gradients, model.trainable_variables))\n",
    "    else:\n",
    "        y_pred = model(x, training=True)\n",
    "        metric_values = loss.compute(y_true, y_pred)\n",
    "    return metric_values\n",
    "\n",
    "#tf.profiler.experimental.start('./tblog')\n",
    "\n",
    "for k in tqdm(range(initial_epoch, epochs), 'total', leave=False):\n",
    "    print('\\nepoch %i/%i' % (k+1, epochs))\n",
    "    metric_util.on_epoch_begin()\n",
    "\n",
    "    for i in tqdm(range(gen_train.num_batches//4), 'training', leave=False):\n",
    "        x, y_true = next(iterator_train)\n",
    "        metric_values = step(x, y_true, training=True)\n",
    "        metric_util.update(metric_values, training=True)\n",
    "        #if i == 100: break\n",
    "    \n",
    "    model.save_weights(checkdir+'/weights.%03i.h5' % (k+1,))\n",
    "\n",
    "    for i in tqdm(range(gen_val.num_batches), 'validation', leave=False):\n",
    "        x, y_true = next(iterator_val)\n",
    "        metric_values = step(x, y_true, training=False)\n",
    "        metric_util.update(metric_values, training=False)\n",
    "        #if i == 10: break\n",
    "\n",
    "    metric_util.on_epoch_end(verbose=1)\n",
    "    #if k == 1: break\n",
    "\n",
    "#tf.profiler.experimental.stop()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Predict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "weights_path = './checkpoints/201809251754_sl384x512_synthtext/weights.020.h5'\n",
    "segment_threshold = 0.50; link_threshold = 0.45\n",
    "#load_weights(model, weights_path)\n",
    "model.load_weights(weights_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_, inputs, images, data = gt_util_val.sample_random_batch(batch_size=1024, input_size=model.image_size)\n",
    "\n",
    "preds = model.predict(inputs, batch_size=1, verbose=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "results = [prior_util.decode(p, segment_threshold, link_threshold) for p in preds]\n",
    "\n",
    "for i in range(8):\n",
    "    plt.figure(figsize=[8]*2)\n",
    "    plt.imshow(images[i])\n",
    "    prior_util.encode(data[i])\n",
    "    prior_util.plot_gt()\n",
    "    prior_util.plot_results(results[i])\n",
    "    plt.axis('off')\n",
    "    #plt.savefig('plots/%s_test_%03i.pgf' % (plot_name, i), bbox_inches='tight')\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Grid Search"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "#steps_seg, steps_lnk = np.arange(0.1, 1, 0.1), np.arange(0.1, 1, 0.1)\n",
    "steps_seg, steps_lnk = np.arange(0.05, 1, 0.05), np.arange(0.05, 1, 0.05)\n",
    "\n",
    "fmes_grid = np.zeros((len(steps_seg),len(steps_lnk)))\n",
    "\n",
    "for i, st in enumerate(steps_seg):\n",
    "    for j, lt in enumerate(steps_lnk):\n",
    "        results = [prior_util.decode(p, st, lt) for p in preds]\n",
    "        TP, FP, FN = evaluate_results(data, results, image_size=image_size)\n",
    "        recall = TP / (TP+FN)\n",
    "        precision = TP / (TP+FP)\n",
    "        fmes = fscore(precision, recall)\n",
    "        fmes_grid[i,j] = fmes\n",
    "        print('segment_threshold %.2f link_threshold %.2f f-measure %.2f' % (st, lt, fmes))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_idx = np.argmax(fmes_grid)\n",
    "max_idx1 = max_idx//fmes_grid.shape[0]\n",
    "max_idx2 = max_idx%fmes_grid.shape[0]\n",
    "print(steps_seg[max_idx1], steps_seg[max_idx2], fmes_grid[max_idx1,max_idx2])\n",
    "plt.figure(figsize=[8]*2)\n",
    "plt.imshow(fmes_grid, cmap='jet', origin='lower', interpolation='bicubic') # nearest, bilinear, bicubic\n",
    "plt.title('f-measure')\n",
    "plt.xticks(range(len(steps_lnk)), steps_lnk.astype('float32'))\n",
    "plt.yticks(range(len(steps_seg)), steps_seg.astype('float32'))\n",
    "plt.plot(max_idx2, max_idx1, 'or')\n",
    "plt.xlabel('link_threshold')\n",
    "plt.ylabel('segment_threshold')\n",
    "plt.grid()\n",
    "#plt.savefig('plots/%s_gridsearch.pgf' % (plot_name), bbox_inches='tight')\n",
    "plt.show()"
   ]
  },
  {
   "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
}