SL_debug.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import cv2\n",
    "import keras\n",
    "import keras.backend as K\n",
    "from tensorflow.python.keras.backend import set_session\n",
    "from keras.callbacks import Callback\n",
    "import tensorflow as tf\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import random\n",
    "import os\n",
    "import time\n",
    "import glob\n",
    "\n",
    "from sl_model import SL512, DSODSL512\n",
    "from ssd_data import InputGenerator\n",
    "from ssd_data import preprocess\n",
    "from sl_utils import PriorUtil\n",
    "from sl_training import SegLinkLoss\n",
    "\n",
    "from sl_utils import plot_rbox, rbox_to_polygon\n",
    "\n",
    "from utils.model import load_weights, count_parameters, calc_memory_usage\n",
    "from utils.training import Logger, LearningRateDecay\n",
    "\n",
    "mean = lambda x: np.sum(x)/len(x)"
   ]
  },
  {
   "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": [
    "#model = SL512()\n",
    "model = DSODSL512()\n",
    "image_size = model.image_size\n",
    "batch_size = 24"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Ground Truth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "_, inputs, images, data = gt_util_val.sample_random_batch(batch_size=16, input_size=image_size)\n",
    "\n",
    "prior_util = PriorUtil(model)\n",
    "gen = InputGenerator(gt_util_train, prior_util, batch_size, image_size, augmentation=False)\n",
    "\n",
    "test_idx = 0\n",
    "test_input = inputs[test_idx]\n",
    "test_img = images[test_idx]\n",
    "test_gt = data[test_idx]\n",
    "\n",
    "plt.figure(figsize=[8]*2)\n",
    "plt.imshow(test_img)\n",
    "gt_util.plot_gt(test_gt, show_labels=False)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "#weights_path = './checkpoints/201711071436_sl512_synthtext/weights.001.h5'\n",
    "#weights_path = './checkpoints/201711132011_dsodsl512_synthtext/weights.001.h5'\n",
    "weights_path = './checkpoints/201806021007_dsodsl512_synthtext/weights.012.h5'\n",
    "\n",
    "segment_threshold = 0.55; link_threshold = 0.45\n",
    "\n",
    "load_weights(model, weights_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "preds = model.predict(inputs, batch_size=1, verbose=1)\n",
    "test_pred = preds[test_idx]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "test_encoded_gt = prior_util.encode(test_gt)\n",
    "#prior_util.print_gt_stats()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "prior_util.decode(test_pred, segment_threshold, link_threshold)\n",
    "\n",
    "plt.figure(figsize=[8]*2)\n",
    "plt.imshow(test_img)\n",
    "#gt_util.plot_gt(test_gt, show_labels=False)\n",
    "prior_util.plot_results(show_labels=False)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "#np.set_printoptions(precision=0, suppress=True, formatter={'all': None})\n",
    "np.set_printoptions(precision=0, suppress=True, formatter={'float': '{:.1f}'.format})\n",
    "\n",
    "mask = test_encoded_gt[:,0] == 0\n",
    "\n",
    "egt = test_encoded_gt[mask][:,7:23:2]\n",
    "mr = test_pred[mask][:,7:23:2]\n",
    "for i in range(len(mr)):\n",
    "    #print(egt[i])\n",
    "    #print(mr[i])\n",
    "    #print(np.abs(egt[i]-mr[i]))\n",
    "    #print(np.mean(np.abs(egt[i]-mr[i])))\n",
    "    pass\n",
    "#print(np.mean(np.abs(egt-mr)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "### Encoding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "#for m_idx in [0,1,2,3]:\n",
    "for m_idx in [2]:\n",
    "    plt.figure(figsize=[16]*2)\n",
    "    plt.axis('off')\n",
    "    plt.imshow(test_img)\n",
    "    #prior_util.prior_maps[m_idx-1].plot_locations()\n",
    "    prior_util.prior_maps[m_idx].plot_locations()\n",
    "    prior_util.prior_maps[m_idx].plot_boxes(range(0,200,40))\n",
    "    prior_util.plot_neighbors(m_idx, [0,20], cross_layer=False)\n",
    "    prior_util.plot_neighbors(m_idx, range(0,200,20), inter_layer=False)\n",
    "    prior_util.plot_assignment(m_idx)\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=[16]*2)\n",
    "plt.axis('off')\n",
    "plt.imshow(test_img)\n",
    "\n",
    "#prior_util.encode(test_gt, debug=False)\n",
    "    \n",
    "#for m_idx in [0,1,2,3]:\n",
    "for m_idx in range(len(prior_util.prior_maps)):\n",
    "#for m_idx in [2]:\n",
    "    #prior_util.prior_maps[m_idx-1].plot_locations()\n",
    "    #prior_util.prior_maps[m_idx].plot_locations()\n",
    "    #prior_util.prior_maps[m_idx].plot_boxes(range(0,200,40))\n",
    "    #prior_util.plot_neighbors(m_idx, [0,20], cross_layer=False)\n",
    "    #prior_util.plot_neighbors(m_idx, range(0,200,20), inter_layer=False)\n",
    "    prior_util.plot_assignment(m_idx)\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Decoding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "dummy_output = np.copy(test_encoded_gt)\n",
    "dummy_output[:,2:4] += np.random.randn(*dummy_output[:,2:4].shape)*0.5\n",
    "\n",
    "plt.figure(figsize=[16]*2)\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, 512); plt.ylim(512,0)\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": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "_, inputs, images, data = gt_util_val.sample_random_batch(batch_size=16, input_size=image_size)\n",
    "\n",
    "preds = model.predict(inputs, batch_size=1, verbose=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "experiment = 'sl512_synthtext'\n",
    "\n",
    "for fl in glob.glob('checkpoints/%s/result_*' % (experiment,)):\n",
    "    #os.remove(fl)\n",
    "    pass\n",
    "\n",
    "for i in range(len(preds)):\n",
    "    plt.figure(figsize=[8]*2)\n",
    "    plt.axis('off')\n",
    "    plt.imshow(images[i])\n",
    "    res = prior_util.decode(preds[i], segment_threshold=0.55, link_threshold=0.35)\n",
    "    prior_util.encode(data[i])\n",
    "    #prior_util.plot_gt()\n",
    "    prior_util.plot_results(res)\n",
    "    #plt.savefig('checkpoints/%s/result_%03d.jpg' % (experiment, i))\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "def decode(self, model_output,\n",
    "            segment_threshold=0.55, link_threshold=0.35, debug=False, debug_combining=False):\n",
    "    \"\"\"Decode local classification and regression results to combined bounding boxes.\n",
    "    \n",
    "    # Arguments\n",
    "        model_output: Array with SegLink model output of shape \n",
    "            (segments, 2 x segment_label + 5 x segment_offset + 16 x inter_layer_links_label \n",
    "            + 8 x cross_layer_links_label)\n",
    "        segment_threshold: Threshold for filtering segment confidence, float betwen 0 and 1.\n",
    "        link_threshold: Threshold for filtering link confidence, float betwen 0 and 1.\n",
    "\n",
    "    # Return\n",
    "        Array with rboxes of shape (results, x + y + w + h + theta + confidence).\n",
    "    \"\"\"\n",
    "    # TODO: handle the case when the line is vertical, tan(theta) == inf and x_proj[max_idx] == x_proj[max_idx]\n",
    "    \n",
    "    segment_labels = model_output[:,0:2]\n",
    "    segment_offsets = model_output[:,2:7]\n",
    "    inter_layer_links_labels = model_output[:,7:23]\n",
    "    cross_layer_links_labels = model_output[:,23:31]\n",
    "\n",
    "    priors_xy = self.priors_xy\n",
    "    priors_wh = self.priors_wh\n",
    "    priors_variances = self.priors_variances\n",
    "    inter_layer_neighbors_idxs = self.inter_layer_neighbors_idxs\n",
    "    cross_layer_neighbors_idxs = self.cross_layer_neighbors_idxs\n",
    "    map_offsets = self.map_offsets\n",
    "    first_map_offset = map_offsets[1] # 64*64\n",
    "    \n",
    "    # filter segments, only pos segments\n",
    "    segment_mask = segment_labels[:,1] > segment_threshold\n",
    "\n",
    "    # filter links, pos links connected with pos segments \n",
    "    inter_layer_link_mask = (inter_layer_links_labels[:,1::2] > link_threshold) & np.repeat(segment_mask[np.newaxis, :], 8, axis=0).T\n",
    "    cross_layer_link_mask = (cross_layer_links_labels[:,1::2] > link_threshold) & np.repeat(segment_mask[np.newaxis, :], 4, axis=0).T\n",
    "\n",
    "    # all pos segments\n",
    "    segment_idxs = np.ix_(segment_mask)[0]\n",
    "    # all segments with pos links\n",
    "    #inter_layer_link_idxs = np.ix_(np.logical_and.reduce(inter_layer_link_mask, axis=1))[0]\n",
    "    #cross_layer_link_idxs = np.ix_(np.logical_and.reduce(cross_layer_link_mask, axis=1))[0]\n",
    "    \n",
    "    # decode segments\n",
    "    offsets = segment_offsets[segment_idxs] # delta(x,y,w,h,theta)_s\n",
    "    offsets = np.copy(offsets)\n",
    "    offsets[:,:4] *= priors_variances[segment_idxs] # variances\n",
    "\n",
    "    rboxes_s = np.empty([len(offsets), 5]) # (x,y,w,h,theta)_s\n",
    "    rboxes_s[:,0:2] = priors_wh[segment_idxs] * offsets[:,0:2] + priors_xy[segment_idxs]\n",
    "    rboxes_s[:,2:4] = priors_wh[segment_idxs] * np.exp(offsets[:,2:4]) # priors_wh is filled with a_l by default\n",
    "    rboxes_s[:,4] = offsets[:,4]\n",
    "    rboxes_s_dict = {segment_idxs[i]: rboxes_s[i] for i in range(len(segment_idxs))}\n",
    "\n",
    "    nodes = list(segment_idxs)\n",
    "    adjacency = {n:set() for n in segment_idxs}\n",
    "    for s_idx in segment_idxs:\n",
    "        # collect inter layer links\n",
    "        for n in np.ix_(inter_layer_link_mask[s_idx])[0]:\n",
    "            n_idx = inter_layer_neighbors_idxs[s_idx, n]\n",
    "            if n_idx in nodes:\n",
    "                # since we add only links to pos segments, they are also valid\n",
    "                adjacency[s_idx].add(n_idx)\n",
    "                adjacency[n_idx].add(s_idx)\n",
    "        # collect cross layer links\n",
    "        if s_idx >= first_map_offset:\n",
    "            for n in np.ix_(cross_layer_link_mask[s_idx])[0]:\n",
    "                n_idx = cross_layer_neighbors_idxs[s_idx-first_map_offset, n]\n",
    "                if n_idx in nodes:\n",
    "                    adjacency[s_idx].add(n_idx)\n",
    "                    adjacency[n_idx].add(s_idx)\n",
    "    \n",
    "    # find connected components\n",
    "    ids = {n:None for n in segment_idxs}\n",
    "\n",
    "    def dfs(node, group_id):\n",
    "        if ids[node] == None:\n",
    "            ids[node] = group_id\n",
    "            for a in adjacency[node]:\n",
    "                dfs(a, group_id)\n",
    "\n",
    "    for i in range(len(nodes)):\n",
    "        dfs(nodes[i], i)\n",
    "    groups = {i:[] for i in set(ids.values())}\n",
    "    for k, v in ids.items():\n",
    "        groups[v].append(k)\n",
    "    \n",
    "    # combine segments\n",
    "    results = []\n",
    "    for f, k in enumerate(groups):\n",
    "        # decoded segment rboxes in group\n",
    "        idxs = np.array(groups[k])\n",
    "        rboxes_s = np.array([rboxes_s_dict[i] for i in idxs]) # (x,y,w,h,theta)_s\n",
    "        \n",
    "        # step 2, algorithm 1\n",
    "        #print('rboxes_s[:,4]', rboxes_s[:,4].shape)\n",
    "        theta_b = mean(rboxes_s[:,4])\n",
    "\n",
    "        # step 3, algorithm 1, find minimizing b in y = a*x + b\n",
    "        # minimize sum (a*x_i + b - y_i)^2 leads to b = mean(y_i - a*x_i)\n",
    "        a = np.tan(-theta_b)\n",
    "        b = mean(rboxes_s[:,1] - a * rboxes_s[:,0])\n",
    "\n",
    "        # step 4, algorithm 1, project centers on the line\n",
    "        # construct line y_p = a_p*x_p + b_p that contains the point and is orthognonal to y = a*x + b\n",
    "        # with a_p = -1/a and b_p = y_p - a_p * x_p we get th point of intersection\n",
    "        # x_s = (b_p - b) / (a - a_p) \n",
    "        # y_s = a * x_s + b\n",
    "        x_proj = (rboxes_s[:,1] + 1/a * rboxes_s[:,0] - b) / (a + 1/a)\n",
    "        y_proj = a * x_proj + b\n",
    "\n",
    "        # find the extreme points\n",
    "        max_idx = np.argmax(x_proj)\n",
    "        min_idx = np.argmin(x_proj)\n",
    "        x_p, y_p = x_proj[min_idx], y_proj[min_idx]\n",
    "        x_q, y_q = x_proj[max_idx], y_proj[max_idx]\n",
    "\n",
    "        # step 5 to 10, algorithm 1, compute the rbox values\n",
    "        w_p = rboxes_s[min_idx,2]\n",
    "        w_q = rboxes_s[max_idx,2]\n",
    "\n",
    "        x_b = (x_p + x_q) / 2\n",
    "        y_b = (y_p + y_q) / 2\n",
    "        w_b = ((x_q - x_p)**2 + (y_q - y_p)**2)**0.5 + (w_p + w_q) / 2\n",
    "        h_b = mean(rboxes_s[:,3])\n",
    "        \n",
    "        rbox_b = [x_b, y_b, w_b, h_b, theta_b]\n",
    "        \n",
    "        # confidence\n",
    "        confs_s = segment_labels[idxs,1]\n",
    "        #conf_b = mean(confs_s)\n",
    "        # weighted confidence by area of segments\n",
    "        boxes_s_area = rboxes_s[:, 2]*rboxes_s[:, 3]\n",
    "        conf_b = np.sum(confs_s * boxes_s_area) / np.sum(boxes_s_area)\n",
    "\n",
    "        results.append(rbox_b + [conf_b])\n",
    "        \n",
    "        # for debugging geometric construction\n",
    "        if debug_combining:\n",
    "            ax = plt.gca()\n",
    "            for rbox in rboxes_s:\n",
    "                c = 'grbck'\n",
    "                c = 'mkgcyb'\n",
    "                c = c[f%len(c)]\n",
    "                plot_rbox(rbox, color=c, linewidth=1)\n",
    "                # segment centers\n",
    "                plt.plot(rbox[0], rbox[1], 'o'+c, markersize=4)\n",
    "                # projected segment centers\n",
    "                plt.plot(x_proj, y_proj, 'oy', markersize=4)\n",
    "            # lines\n",
    "            x_l = np.array([0,512])\n",
    "            y_l = a * x_l + b\n",
    "            plt.plot(x_l, y_l, 'r')\n",
    "            # endpoints\n",
    "            plt.plot(x_p, y_p, 'or', markersize=6)\n",
    "            plt.plot(x_q, y_q, 'or', markersize=6)\n",
    "            # combined box\n",
    "            plot_rbox(rbox_b, color='r', linewidth=2)\n",
    "\n",
    "    if len(results) > 0:\n",
    "        results = np.asarray(results)\n",
    "    else:\n",
    "        results = np.empty((0,6))\n",
    "    self.results = results\n",
    "\n",
    "    # debug\n",
    "    if debug:\n",
    "        ax = plt.gca()\n",
    "        \n",
    "        # plot positive links\n",
    "        inter_layer_link_mask = inter_layer_links_labels[:,1::2] > link_threshold\n",
    "        for idx in range(len(inter_layer_link_mask)):\n",
    "            p1 = priors_xy[idx]\n",
    "            for n_idx in inter_layer_neighbors_idxs[idx][inter_layer_link_mask[idx]]:\n",
    "                p2 = priors_xy[n_idx]\n",
    "                plt.plot([p1[0], p2[0]], [p1[1], p2[1]], 'y-', linewidth=2)\n",
    "        \n",
    "        cross_layer_link_mask = cross_layer_links_labels[:,1::2] > link_threshold\n",
    "        for idx in range(len(cross_layer_neighbors_idxs)):\n",
    "            p1 = priors_xy[idx+first_map_offset]\n",
    "            for n_idx in cross_layer_neighbors_idxs[idx][cross_layer_link_mask[idx+first_map_offset]]:\n",
    "                p2 = priors_xy[n_idx]\n",
    "                plt.plot([p1[0], p2[0]], [p1[1], p2[1]], '-', color='orange', linewidth=2)\n",
    "                \n",
    "        # plot segments\n",
    "        keys = list(rboxes_s_dict.keys())\n",
    "        for k in keys:\n",
    "            plot_rbox(rboxes_s_dict[k], color='k', linewidth=2)\n",
    "\n",
    "        # plot links\n",
    "        for k in keys:\n",
    "            p1 = rboxes_s_dict[k][:2]\n",
    "            for m in adjacency[k]:\n",
    "                p2 = rboxes_s_dict[m][:2]\n",
    "                plt.plot([p1[0], p2[0]], [p1[1], p2[1]], 'm-', linewidth=1)\n",
    "\n",
    "        # plot priors\n",
    "        for k in keys:\n",
    "            p1 = rboxes_s_dict[k][:2]\n",
    "            p2 = priors_xy[k]\n",
    "            plt.plot([p1[0]], [p1[1]], 'mo', markersize=4)\n",
    "            plt.plot([p2[0]], [p2[1]], 'go', markersize=4)\n",
    "            plt.plot([p1[0], p2[0]], [p1[1], p2[1]], 'g-', linewidth=1)\n",
    "    \n",
    "    return results\n",
    "\n",
    "\n",
    "plt.figure(figsize=[14]*2)\n",
    "ax = plt.gca()\n",
    "plt.imshow(test_img)\n",
    "#res = prior_util.decode(dummy_output, debug=False)\n",
    "res = decode(prior_util, dummy_output, debug=False)\n",
    "#res = decode(prior_util, test_pred, debug=False)\n",
    "#prior_util.plot_gt()\n",
    "prior_util.plot_results(res)\n",
    "plt.axis('off'); plt.xlim(0, 512); plt.ylim(512,0)\n",
    "plt.show()\n",
    "\n",
    "#print(res)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "np.set_printoptions(precision=0, suppress=True, formatter={'float': '{:7.3f}'.format})\n",
    "\n",
    "i = 6\n",
    "plt.figure(figsize=[14]*2)\n",
    "plt.axis('off')\n",
    "plt.imshow(images[i])\n",
    "res = decode(prior_util, preds[i], segment_threshold=0.55, link_threshold=0.35, debug=False, debug_combining=False)\n",
    "prior_util.encode(data[i])\n",
    "prior_util.plot_gt()\n",
    "prior_util.plot_results(res)\n",
    "#plt.savefig('checkpoints/%s/result_%03d.jpg' % (experiment, i))\n",
    "for i in range(len(prior_util.prior_maps)):\n",
    "    #prior_util.plot_neighbors(i, list(range(10))*5)\n",
    "    pass\n",
    "print(res)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "%%timeit\n",
    "res = decode(prior_util, preds[i], segment_threshold=0.55, link_threshold=0.35, debug=False, debug_combining=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=[16]*2)\n",
    "\n",
    "decode(prior_util, test_pred, segment_threshold=0.55, link_threshold=0.35, debug=False, debug_combining=True)\n",
    "\n",
    "#plt.show()\n",
    "\n",
    "#plt.figure(figsize=[24,24])\n",
    "plt.imshow(test_img)\n",
    "#gt_util.plot_gt(test_gt, show_labels=False)\n",
    "#prior_util.plot_results(show_labels=False)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def encode(self, gt_data, debug=False):\n",
    "    \"\"\"Encode ground truth polygones to segments and links for local classification and regression.\n",
    "\n",
    "    # Arguments\n",
    "        gt_data: shape (boxes, 4 xy + classes)\n",
    "\n",
    "    # Return\n",
    "        shape (priors, 2 segment_labels + 5 segment_offsets + 2*8 inter_layer_links_labels + 2*4 cross_layer_links_labels)\n",
    "    \"\"\"\n",
    "\n",
    "    rboxes = []\n",
    "    polygons = []\n",
    "    for word in gt_data:\n",
    "        xy = np.reshape(word[:8], (-1, 2))\n",
    "        xy = np.copy(xy) * self.image_size\n",
    "        polygons.append(xy)\n",
    "        rbox = polygon_to_rbox(xy)\n",
    "        rboxes.append(rbox)\n",
    "    rboxes = self.gt_rboxes = np.array(rboxes)\n",
    "    polygnos = self.gt_polygons = np.array(polygons)\n",
    "\n",
    "    # compute segments\n",
    "    for i in range(len(self.prior_maps)):\n",
    "        m = self.prior_maps[i]\n",
    "\n",
    "        # compute priors\n",
    "        #m.compute_priors()\n",
    "\n",
    "        num_priors = len(m.priors)\n",
    "\n",
    "        # assigne gt to priors\n",
    "        a_l = m.minmax_size[0]\n",
    "        match_indices = np.full(num_priors, -1, dtype=np.int32)\n",
    "        min_lhs_eq_11 = np.full(num_priors, 1e6, dtype=np.float32)\n",
    "        for j in range(len(rboxes)): # ~12.9 ms\n",
    "            cx, cy, w, h, theta = rboxes[j]\n",
    "            c = rboxes[j,:2]\n",
    "            # constraint on ratio between box size and word height, equation (11)\n",
    "            lhs_eq_11 = max(a_l/h, h/a_l)\n",
    "            if lhs_eq_11 <= 1.5:\n",
    "                R = rot_matrix(theta)\n",
    "                for k in range(num_priors): # hurts\n",
    "                    # is center of prior is in gt rbox\n",
    "                    d = np.abs(np.dot(m.priors_xy[k]-c, R.T))\n",
    "                    if d[0] < w/2. and d[1] < h/2.:\n",
    "                        # is lhs of equation (11) minimal for prior\n",
    "                        if lhs_eq_11 < min_lhs_eq_11[k]:\n",
    "                            min_lhs_eq_11[k] = lhs_eq_11\n",
    "                            match_indices[k] = j   \n",
    "        m.match_indices = match_indices\n",
    "\n",
    "        segment_mask = match_indices != -1\n",
    "\n",
    "        # segment labels\n",
    "        m.segment_labels = np.empty((num_priors, 2), dtype=np.int8)\n",
    "        m.segment_labels[:, 0] = np.logical_not(segment_mask)\n",
    "        m.segment_labels[:, 1] = segment_mask\n",
    "\n",
    "        # compute offsets only for assigned boxes\n",
    "        m.segment_offsets = np.zeros((num_priors, 5))\n",
    "        pos_segment_idxs = np.nonzero(segment_mask)[0]\n",
    "        for j in pos_segment_idxs: # box_idx # ~4 ms\n",
    "            gt_idx = match_indices[j]\n",
    "            rbox = rboxes[gt_idx]\n",
    "            polygon = polygons[gt_idx]\n",
    "            cx, cy, w, h, theta = rbox\n",
    "            R = rot_matrix(theta)\n",
    "            prior_x, prior_y = m.priors_xy[j]\n",
    "            prior_w, prior_h = m.priors_wh[j]\n",
    "\n",
    "            # step 2 figuer 5, rotate word anticlockwise around the center of prior\n",
    "            d = rbox[:2] - m.priors_xy[j]\n",
    "            poly_loc = rbox_to_polygon([*d, w, h, theta])\n",
    "            poly_loc_easy = polygon - m.priors_xy[j]\n",
    "\n",
    "            poly_loc_rot = np.dot(poly_loc, R.T)\n",
    "\n",
    "            # step 3 figure 5, crop word to left and right of prior\n",
    "            poly_loc_coped = np.copy(poly_loc_rot)\n",
    "            poly_loc_coped[:,0] = np.clip(poly_loc_coped[:,0], -prior_w/2., prior_w/2.)\n",
    "\n",
    "            # step 4 figure 5, rotate croped word box clockwisely\n",
    "            poly_loc_rot_back = np.dot(poly_loc_coped, R)\n",
    "            rbox_loc_rot_back = polygon_to_rbox(poly_loc_rot_back)\n",
    "\n",
    "            # encode, solve (3) to (7) to get local offsets\n",
    "            offset = np.array([*(rbox_loc_rot_back[:2]/a_l), \n",
    "                               *(np.log(rbox_loc_rot_back[2:4]/a_l)), \n",
    "                               rbox_loc_rot_back[4]])\n",
    "            offset[:4] /= m.priors[j,-4:] # variances\n",
    "            m.segment_offsets[j] = offset\n",
    "\n",
    "            # for debugging local geometry\n",
    "            if debug:\n",
    "                prior_poly_loc = np.array([[-prior_w, +prior_h],\n",
    "                                           [+prior_w, +prior_h],\n",
    "                                           [+prior_w, -prior_h],\n",
    "                                           [-prior_w, -prior_h]])/2.\n",
    "                plt.figure(figsize=[10]*2)\n",
    "                ax = plt.gca()\n",
    "                ax.add_patch(plt.Polygon(prior_poly_loc, fill=False, edgecolor='r', linewidth=1))\n",
    "                ax.add_patch(plt.Polygon(poly_loc, fill=False, edgecolor='b', linewidth=1))\n",
    "                ax.add_patch(plt.Polygon(np.dot(poly_loc, R.T), fill=False, edgecolor='k', linewidth=1))\n",
    "                #ax.add_patch(plt.Polygon(poly_loc_easy, fill=False, edgecolor='r', linewidth=1))\n",
    "                #ax.add_patch(plt.Polygon(np.dot(poly_loc_easy, R.T), fill=False, edgecolor='y', linewidth=1))\n",
    "                ax.add_patch(plt.Polygon(poly_loc_coped, fill=False, edgecolor='c', linewidth=1))\n",
    "                ax.add_patch(plt.Polygon(poly_loc_rot_back, fill=False, edgecolor='y', linewidth=1))\n",
    "                lim = 50; plt.xlim(-lim,lim); plt.ylim(-lim,lim); plt.grid()\n",
    "                plt.show()\n",
    "                break\n",
    "\n",
    "        # compute link labels\n",
    "        m.inter_layer_links_labels = np.zeros((num_priors,16), dtype=np.int8)\n",
    "        m.cross_layer_links_labels = np.zeros((num_priors,8), dtype=np.int8)\n",
    "        if i > 0:\n",
    "            previous_map = self.prior_maps[i-1]\n",
    "        # we only have to check neighbors if we are positive\n",
    "        for idx in pos_segment_idxs:\n",
    "            neighbor_idxs = m.inter_layer_neighbors_idxs[idx]\n",
    "            for n, neighbor_idx in enumerate(neighbor_idxs):\n",
    "                # valide neighbors\n",
    "                if m.inter_layer_neighbors_valide[idx,n]:\n",
    "                    # neighbor matched to the same word\n",
    "                    if match_indices[idx] == match_indices[neighbor_idx]:\n",
    "                        # since we are positive and match to the same word, neighbor has to be positive\n",
    "                        m.inter_layer_links_labels[idx, n*2+1] = 1\n",
    "            # would be nice, but we refere to invalide neighbors\n",
    "            #label = m.inter_layer_neighbors_valide[idx] & (match_indices[neighbor_idxs] == match_indices[idx])\n",
    "            #m.inter_layer_links_labels[idx, 1::2] = label\n",
    "\n",
    "            if i > 0:\n",
    "                neighbor_idxs = m.cross_layer_neighbors_idxs[idx]\n",
    "                for n, neighbor_idx in enumerate(neighbor_idxs):\n",
    "                    # cross layer neighbors are always valide\n",
    "                    if match_indices[idx] == previous_map.match_indices[neighbor_idx]:\n",
    "                        m.cross_layer_links_labels[idx, n*2+1] = 1\n",
    "\n",
    "        m.inter_layer_links_labels[:,::2] = np.logical_not(m.inter_layer_links_labels[:,1::2])\n",
    "        m.cross_layer_links_labels[:,::2] = np.logical_not(m.cross_layer_links_labels[:,1::2])\n",
    "\n",
    "    # collect encoded ground truth\n",
    "    maps = self.prior_maps\n",
    "    segment_labels = np.concatenate([m.segment_labels for m in maps])\n",
    "    segment_offsets = np.concatenate([m.segment_offsets for m in maps])\n",
    "    inter_layer_links_labels = np.concatenate([m.inter_layer_links_labels for m in maps])\n",
    "    cross_layer_links_labels = np.concatenate([m.cross_layer_links_labels for m in maps])\n",
    "    return np.concatenate([segment_labels, segment_offsets, inter_layer_links_labels, cross_layer_links_labels], axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "def plot_assignment(self, map_idx=None):\n",
    "    \"\"\"Draw the combined bounding boxes in the current figure.\n",
    "\n",
    "    # Arguments\n",
    "        map_idx: The index of the considered ProroMap.\n",
    "            If None, all maps are considered.\n",
    "    \"\"\"\n",
    "    ax = plt.gca()\n",
    "    # groud truth polygones\n",
    "    for p in self.gt_polygons:\n",
    "        ax.add_patch(plt.Polygon(p, fill=False, edgecolor='y', linewidth=4))\n",
    "    # groud truth rboxes\n",
    "    rboxes = self.gt_rboxes\n",
    "    for rbox in rboxes:\n",
    "        box = rbox_to_polygon(rbox)\n",
    "        ax.add_patch(plt.Polygon(box, fill=False, edgecolor='b', linewidth=2))\n",
    "    plt.plot(rboxes[:,0], rboxes[:,1], 'go',  markersize=4)\n",
    "    m = self.prior_maps[map_idx]\n",
    "    # assigned boxes\n",
    "    for idx in np.nonzero(m.segment_labels[:, 1])[0]:\n",
    "        p_prior = m.priors_xy[idx]\n",
    "        p_word = rboxes[m.match_indices[idx]][:2]\n",
    "        plt.plot([p_prior[0], p_word[0]], [p_prior[1], p_word[1]], 'm-', linewidth=1)\n",
    "        #plt.plot([p_word[0]], [p_word[1]], 'ro',  markersize=8)\n",
    "    # links\n",
    "    labels = m.inter_layer_links_labels[:,1::2]\n",
    "    idxs = np.nonzero(np.any(labels, axis=1))[0]\n",
    "    for idx in idxs:\n",
    "        for n_idx in m.inter_layer_neighbors_idxs[idx, np.nonzero(labels[idx])[0]]:\n",
    "            x, y = m.priors_xy[idx]\n",
    "            n_x, n_y = m.priors_xy[n_idx]\n",
    "            plt.plot([x, n_x], [y, n_y], '-c', linewidth=1)\n",
    "    if map_idx > 0:\n",
    "        n_m = self.prior_maps[map_idx-1]\n",
    "        labels = m.cross_layer_links_labels[:,1::2]\n",
    "        idxs = np.nonzero(np.any(labels, axis=1))[0]\n",
    "        for idx in idxs:\n",
    "            x, y = m.priors_xy[idx]\n",
    "            for n_idx in m.cross_layer_neighbors_idxs[idx, np.nonzero(labels[idx])[0]]:\n",
    "                n_x, n_y = n_m.priors_xy[n_idx]\n",
    "                plt.plot([x, n_x], [y, n_y], '-c', linewidth=1)\n",
    "\n",
    "prior_util.encode(test_gt, debug=True)\n",
    "plt.figure(figsize=[8]*2)\n",
    "\n",
    "plt.imshow(test_img)\n",
    "\n",
    "plot_assignment(prior_util, 1)\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": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# backup\n",
    "\n",
    "def decode(self, model_output,\n",
    "           #segment_labels, segment_offsets, inter_layer_links_labels, cross_layer_links_labels,\n",
    "            segment_threshold=0.55, link_threshold=0.35, debug=False):\n",
    "    \"\"\"Decode local classification and regression results to combined bounding boxes.\n",
    "\n",
    "    # Return\n",
    "        Array with rboxes of shape (results, x + y + w + h + theta).\n",
    "    \"\"\"\n",
    "    # TODO: handle the case when the line is vertical, tan(theta) == inf and x_proj[max_idx] == x_proj[max_idx]\n",
    "\n",
    "    segment_labels = model_output[:,0:2]\n",
    "    segment_offsets = model_output[:,2:7]\n",
    "    inter_layer_links_labels = model_output[:,7:23]\n",
    "    cross_layer_links_labels = model_output[:,23:31]\n",
    "\n",
    "    map_offsets = self.map_offsets\n",
    "\n",
    "    # filter segments, only pos segments\n",
    "    segment_mask = segment_labels[:,1] > segment_threshold\n",
    "\n",
    "    # filter links, pos links connected with pos segments \n",
    "    #inter_layer_link_mask = (inter_layer_links_labels[:,1::2] > link_threshold) & np.tile(segment_mask[:,np.newaxis], (1, 8))\n",
    "    #cross_layer_link_mask = (cross_layer_links_labels[:,1::2] > link_threshold) & np.tile(segment_mask[map_offsets[1]:,np.newaxis], (1, 4))\n",
    "    #cross_layer_link_mask = (cross_layer_links_labels[:,1::2] > link_threshold) & np.tile(segment_mask[:,np.newaxis], (1, 4))\n",
    "    inter_layer_link_mask = (inter_layer_links_labels[:,1::2] > link_threshold) & np.repeat(segment_mask[np.newaxis, :], 8, axis=0).T\n",
    "    cross_layer_link_mask = (cross_layer_links_labels[:,1::2] > link_threshold) & np.repeat(segment_mask[np.newaxis, :], 4, axis=0).T\n",
    "    \n",
    "    # all pos segments\n",
    "    segment_idxs = np.ix_(segment_mask)[0]\n",
    "    # all segments with pos links\n",
    "    #inter_layer_link_idxs = np.argwhere(np.any(inter_layer_link_mask, axis=1))[:,0]\n",
    "    #cross_layer_link_idxs = np.argwhere(np.any(cross_layer_link_mask, axis=1))[:,0]\n",
    "    #inter_layer_link_idxs = np.ix_(np.logical_and.reduce(inter_layer_link_mask, axis=1))[0]\n",
    "    #cross_layer_link_idxs = np.ix_(np.logical_and.reduce(cross_layer_link_mask, axis=1))[0]\n",
    "    \n",
    "    nodes = list(segment_idxs)\n",
    "    adjacency = {n:set() for n in segment_idxs}\n",
    "    rboxes_s_dict = {}\n",
    "    \n",
    "    for i in range(len(self.prior_maps)):\n",
    "        m = self.prior_maps[i]\n",
    "        sl = slice(map_offsets[i], map_offsets[i+1])\n",
    "\n",
    "        # decode segments\n",
    "        mask = segment_mask[sl]\n",
    "        in_map_idxs = np.where(mask)[0]\n",
    "        offsets = segment_offsets[sl][mask] # delta(x,y,w,h,theta)_s\n",
    "        offsets = np.copy(offsets)\n",
    "        offsets[:,:4] *= m.priors_variances[mask] # variances\n",
    "\n",
    "\n",
    "        rboxes_s = np.empty([len(offsets), 5]) # (x,y,w,h,theta)_s\n",
    "        rboxes_s[:,0:2] = m.priors_wh[mask] * offsets[:,0:2] + m.priors_xy[mask]\n",
    "        rboxes_s[:,2:4] = m.priors_wh[mask] * np.exp(offsets[:,2:4]) # priors_wh is filled with a_l by default\n",
    "        rboxes_s[:,4] = offsets[:,4]\n",
    "        for idx, rbox in zip(in_map_idxs+map_offsets[i], rboxes_s):\n",
    "            rboxes_s_dict[idx] = rbox\n",
    "\n",
    "        # collect inter layer links\n",
    "        for s_in_map_idx, n in np.array(np.where(inter_layer_link_mask[sl])).T:\n",
    "            n_in_map_idx = m.inter_layer_neighbors_idxs[s_in_map_idx, n]\n",
    "            s_idx = s_in_map_idx + map_offsets[i]\n",
    "            n_idx = n_in_map_idx + map_offsets[i]\n",
    "            if n_idx in nodes:\n",
    "                # since we add only links to pos segments, they are also valide\n",
    "                adjacency[s_idx].add(n_idx)\n",
    "                adjacency[n_idx].add(s_idx)\n",
    "        \n",
    "        # collect cross layer links\n",
    "        if i > 0:\n",
    "            #sl = slice(map_offsets[i] - map_offsets[1], map_offsets[i+1] - map_offsets[1])\n",
    "            sl = slice(map_offsets[i], map_offsets[i+1])\n",
    "            for s_in_map_idx, n in np.array(np.where(cross_layer_link_mask[sl])).T:\n",
    "                n_in_map_idx = m.cross_layer_neighbors_idxs[s_in_map_idx, n]\n",
    "                s_idx = s_in_map_idx + map_offsets[i]\n",
    "                n_idx = n_in_map_idx + map_offsets[i-1]\n",
    "                if n_idx in nodes:\n",
    "                    adjacency[s_idx].add(n_idx)\n",
    "                    adjacency[n_idx].add(s_idx)\n",
    "\n",
    "    # find connected components\n",
    "    ids = {n:None for n in segment_idxs}\n",
    "\n",
    "    def dfs(node, group_id):\n",
    "        if ids[node] == None:\n",
    "            ids[node] = group_id\n",
    "            for a in adjacency[node]:\n",
    "                dfs(a, group_id)\n",
    "\n",
    "    for i in range(len(nodes)):\n",
    "        dfs(nodes[i], i)\n",
    "    groups = {i:[] for i in set(ids.values())}\n",
    "    for k, v in ids.items():\n",
    "        groups[v].append(k)\n",
    "    \n",
    "    # combine segments\n",
    "    rboxes_b = []\n",
    "    for f, k in enumerate(groups):\n",
    "        # decoded segment rboxes in group\n",
    "        idxs = np.array(groups[k])\n",
    "        rboxes_s = np.array([rboxes_s_dict[i] for i in idxs]) # (x,y,w,h,theta)_s\n",
    "        n = len(rboxes_s)\n",
    "\n",
    "        # step 2, algorithm 1\n",
    "        theta_b = mean(rboxes_s[:,4])\n",
    "\n",
    "        # step 3, algorithm 1, find minimizing b in y = a*x + b\n",
    "        # minimize sum (a*x_i + b - y_i)^2 leads to b = mean(y_i - a*x_i)\n",
    "        a = np.tan(-theta_b)\n",
    "        b = mean(rboxes_s[:,1] - a * rboxes_s[:,0])\n",
    "\n",
    "        # step 4, algorithm 1, project centers on the line\n",
    "        # construct line y_p = a_p*x_p + b_p that contains the point and is orthognonal to y = a*x + b\n",
    "        # with a_p = -1/a and b_p = y_p - a_p * x_p we get th point of intersection\n",
    "        # x_s = (b_p - b) / (a - a_p) \n",
    "        # y_s = a * x_s + b\n",
    "        x_proj = (rboxes_s[:,1] + 1/a * rboxes_s[:,0] - b) / (a + 1/a)\n",
    "        y_proj = a * x_proj + b\n",
    "\n",
    "        # find the extreme points\n",
    "        max_idx = np.argmax(x_proj)\n",
    "        min_idx = np.argmin(x_proj)\n",
    "        x_p, y_p = x_proj[min_idx], y_proj[min_idx]\n",
    "        x_q, y_q = x_proj[max_idx], y_proj[max_idx]\n",
    "\n",
    "        # step 5 to 10, algorithm 1, compute the rbox values\n",
    "        w_p = rboxes_s[min_idx,2]\n",
    "        w_q = rboxes_s[max_idx,2]\n",
    "\n",
    "        x_b = (x_p + x_q) / 2\n",
    "        y_b = (y_p + y_q) / 2\n",
    "        w_b = ((x_q - x_p)**2 + (y_q - y_p)**2)**0.5 + (w_p + w_q) / 2\n",
    "        h_b = mean(rboxes_s[:,3])\n",
    "\n",
    "        rbox_b = np.array([x_b, y_b, w_b, h_b, theta_b])\n",
    "        rboxes_b.append(rbox_b)\n",
    "\n",
    "        # for debugging geometric construction\n",
    "        if debug:\n",
    "            ax = plt.gca()\n",
    "            for rbox in rboxes_s:\n",
    "                box = rbox_to_polygon(rbox)\n",
    "                c = 'grbck'\n",
    "                c = 'mkgcyb'\n",
    "                c = c[f%len(c)]\n",
    "                ax.add_patch(plt.Polygon(box, fill=False, edgecolor=c, linewidth=1))\n",
    "                # segment centers\n",
    "                plt.plot(rbox[0], rbox[1], 'o'+c,  markersize=4)\n",
    "                # projected segment centers\n",
    "                plt.plot(x_proj, y_proj, 'oy',  markersize=4)\n",
    "            # lines\n",
    "            x_l = np.array([0,512])\n",
    "            y_l = a * x_l + b\n",
    "            plt.plot(x_l, y_l, 'r')\n",
    "            # endpoints\n",
    "            plt.plot(x_p, y_p, 'or', markersize=8)\n",
    "            plt.plot(x_q, y_q, 'or', markersize=8)\n",
    "            # combined box\n",
    "            box = rbox_to_polygon(rbox_b)\n",
    "            ax.add_patch(plt.Polygon(box, fill=False, edgecolor='r', linewidth=2))\n",
    "\n",
    "    rboxes_b = np.array(rboxes_b)\n",
    "    self.rboxes_b = rboxes_b\n",
    "    \n",
    "    # debug\n",
    "    if False:\n",
    "    # plot positive links\n",
    "        xy = np.concatenate([m.priors_xy for m in self.prior_maps])\n",
    "\n",
    "        inter_layer_neighbors_idxs = np.concatenate([ \n",
    "            self.prior_maps[i].inter_layer_neighbors_idxs + map_offsets[i] for i in range(len(self.prior_maps)) ])\n",
    "        inter_layer_link_mask = inter_layer_links_labels[:,1::2] > link_threshold\n",
    "        for idx in range(len(inter_layer_link_mask)):\n",
    "            n_idxs = inter_layer_neighbors_idxs[idx][inter_layer_link_mask[idx]]\n",
    "            p1 = xy[idx]\n",
    "            for n_idx in n_idxs:\n",
    "                p2 = xy[n_idx]\n",
    "                plt.plot([p1[0], p2[0]], [p1[1], p2[1]], 'y-', linewidth=2)\n",
    "\n",
    "        #cross_layer_neighbors_idxs = np.concatenate([\n",
    "        #    self.prior_maps[i].cross_layer_neighbors_idxs + map_offsets[i] for i in range(1, len(self.prior_maps)) ])\n",
    "        cross_layer_neighbors_idxs = np.concatenate([\n",
    "            self.prior_maps[i].cross_layer_neighbors_idxs + map_offsets[i-1] for i in range(1, len(self.prior_maps)) ])\n",
    "        cross_layer_link_mask = cross_layer_links_labels[:,1::2] > link_threshold\n",
    "        print('cross_layer_link_mask', len(cross_layer_link_mask))\n",
    "        print('cross_layer_neighbors_idxs', len(cross_layer_neighbors_idxs))\n",
    "        for idx in range(len(cross_layer_neighbors_idxs)):\n",
    "            n_idxs = cross_layer_neighbors_idxs[idx][cross_layer_link_mask[idx+map_offsets[1]]]\n",
    "            p1 = xy[idx+map_offsets[1]]\n",
    "            for n_idx in n_idxs:\n",
    "                try:\n",
    "                    p2 = xy[n_idx]\n",
    "                    plt.plot([p1[0], p2[0]], [p1[1], p2[1]], '-', color='orange', linewidth=2)\n",
    "                except Exception as e:\n",
    "                    print(e)\n",
    "\n",
    "        # plot segments\n",
    "        keys = list(rboxes_s_dict.keys())\n",
    "        ax = plt.gca()\n",
    "        for k in keys:\n",
    "            box = rbox_to_polygon(rboxes_s_dict[k])\n",
    "            ax.add_patch(plt.Polygon(box, fill=False, edgecolor='k', linewidth=2))\n",
    "\n",
    "        # plot links\n",
    "        for k in keys:\n",
    "            p1 = rboxes_s_dict[k][:2]\n",
    "            for m in adjacency[k]:\n",
    "                p2 = rboxes_s_dict[m][:2]\n",
    "                plt.plot([p1[0], p2[0]], [p1[1], p2[1]], 'm-', linewidth=1)\n",
    "\n",
    "        # plot priors\n",
    "        xy = np.concatenate([m.priors_xy for m in self.prior_maps])\n",
    "        for k in keys:\n",
    "            p1 = rboxes_s_dict[k][:2]\n",
    "            p2 = xy[k]\n",
    "            plt.plot([p1[0]], [p1[1]], 'mo', markersize=4)\n",
    "            plt.plot([p2[0]], [p2[1]], 'go', markersize=4)\n",
    "            plt.plot([p1[0], p2[0]], [p1[1], p2[1]], 'g-', linewidth=1)\n",
    "            #print(k, type(k))\n",
    "    \n",
    "    return rboxes_b\n",
    "\n",
    "def plot_results(self, results=None):\n",
    "    if results is None:\n",
    "        results = self.rboxes_b\n",
    "    ax = plt.gca()\n",
    "    for rbox_b in results:\n",
    "        box = rbox_to_polygon(rbox_b)\n",
    "        ax.add_patch(plt.Polygon(box, fill=False, edgecolor='r', linewidth=2))\n",
    "\n",
    "        \n",
    "def plot_linking(self):\n",
    "    \n",
    "    for i in range(len(self.prior_maps)):\n",
    "        m = self.prior_maps[i]\n",
    "        \n",
    "#i = 6\n",
    "#plt.figure(figsize=[14]*2)\n",
    "#plt.imshow(images[i])\n",
    "#res = decode(prior_util, preds[i], segment_threshold=0.55, link_threshold=0.35, debug=False)\n",
    "#prior_util.encode(data[i])\n",
    "##prior_util.plot_gt()\n",
    "#plot_results(prior_util, res)\n",
    "#plot_linking(prior_util)\n",
    "##plt.savefig('checkpoints/%s/result_%03d.jpg' % (experiment, i))\n",
    "#for i in range(len(prior_util.prior_maps)):\n",
    "#    #prior_util.plot_neighbors(i, list(range(10))*5)\n",
    "#    pass\n",
    "#plt.show()\n",
    "\n",
    "plt.figure(figsize=[14]*2)\n",
    "ax = plt.gca()\n",
    "plt.imshow(test_img)\n",
    "#res = prior_util.decode(dummy_output, debug=False)\n",
    "#res = decode(prior_util, dummy_output, debug=False)\n",
    "res = decode(prior_util, test_pred, debug=False)\n",
    "#prior_util.plot_gt()\n",
    "prior_util.plot_results(res)\n",
    "plt.axis('off'); plt.xlim(0, 512); plt.ylim(512,0)\n",
    "plt.show()\n",
    "\n",
    "#print(res)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "%%timeit\n",
    "res = decode(prior_util, preds[i], segment_threshold=0.55, link_threshold=0.35, debug=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "dummy_output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "test_pred"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from tensorflow.python.keras.utils.vis_utils import plot_model\n",
    "plot_model(model, to_file='model.png', show_shapes=False, show_layer_names=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "dummy = np.array([dummy_output, dummy_output, dummy_output, dummy_output])\n",
    "\n",
    "dummy_batch_size = dummy.shape[0]\n",
    "\n",
    "n_seg = 64*64+32*32+16*16+8*8+4*4+2*2+1*1\n",
    "n_inter = 64*64+32*32+16*16+8*8+4*4+2*2+1*1\n",
    "n_cross = 32*32+16*16+8*8+4*4+2*2+1*1\n",
    "print('seg ', n_seg * dummy_batch_size, 21844)\n",
    "print('link', (n_inter*8 + n_cross*4 ) * dummy_batch_size, 196592)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "for name in prior_util.source_layers_names:\n",
    "    print('%-10s %s' % (name, model.get_layer(name).output_shape[1:3]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "%%timeit\n",
    "x, y = next(gen.generate())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "%%timeit\n",
    "result = prior_util.decode(dummy_output, segment_threshold=0.9, link_threshold=0.7, debug=False);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# profile encoding batch\n",
    "from cProfile import Profile\n",
    "pr = Profile()\n",
    "pr.enable()\n",
    "x, y = next(gen.generate())\n",
    "pr.disable()\n",
    "pr.print_stats(sort='cumulative')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# profile decoding\n",
    "from cProfile import Profile\n",
    "pr = Profile()\n",
    "pr.enable()\n",
    "for i in range(1000):\n",
    "    #results = prior_util.decode(test_encoded_gt)\n",
    "    result = prior_util.decode(dummy_output, segment_threshold=0.55, link_threshold=0.35, debug=False)\n",
    "pr.disable()\n",
    "pr.print_stats(sort='cumulative')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "j = 0\n",
    "for img_batch, gt_batch in gen.generate():\n",
    "    #print(np.max(gt_batch[:,:,2:7], axis=(0,1)))\n",
    "    if np.any(np.isnan(gt_batch[:,:,2:7])):\n",
    "        print(gt_batch[:,:,2:7])\n",
    "    \n",
    "    if np.max(gt_batch[:,:,1]) > 1.0:\n",
    "        print(gt_batch.shape, np.max(gt_batch[:,:,1]), np.sum(gt_batch[:,:,1]))\n",
    "    \n",
    "    j += 1\n",
    "    if j > 10: break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from sl_utils import polygon_to_rbox, rbox_to_polygon\n",
    "\n",
    "xy = test_gt[1][:8]\n",
    "#print(xy)\n",
    "xy = np.reshape(xy, (-1, 2))\n",
    "print(xy)\n",
    "rbox = polygon_to_rbox(xy)\n",
    "print(rbox)\n",
    "box = rbox_to_polygon(rbox)\n",
    "print(box)\n",
    "\n",
    "plt.figure()\n",
    "ax = plt.gca()\n",
    "ax.add_patch(plt.Polygon(xy, fill=False, edgecolor='r'))\n",
    "ax.add_patch(plt.Polygon(box, fill=False, edgecolor='b'))\n",
    "#plt.xlim(-0,1.2); plt.ylim(-0,1.2)\n",
    "#plt.xlim(0.4,0.6); plt.ylim(0.5,0.8)\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": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# plot line segments\n",
    "from matplotlib.collections import LineCollection\n",
    "n_idxs = inter_layer_neighbors_idxs[idx][inter_layer_link_mask[idx]]\n",
    "p2s = priors_xy[n_idxs]\n",
    "lines = np.array([np.tile(p1,(len(p2s),1)),p2s]).transpose((1,0,2))\n",
    "ax.add_collection(LineCollection(lines, colors='y', linewidths=2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# memory usage\n",
    "s = ''\n",
    "fs = '%-20s %5d kb\\n'\n",
    "s += fs % ('priors', prior_util.priors.nbytes/1024)\n",
    "s += fs % ('input image', test_img.nbytes/1024)\n",
    "s += fs % ('encoded gt / output', test_encoded_gt.nbytes/1024)\n",
    "print(s)\n",
    "calc_memory_usage(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "### Generate some samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class PriorUtilDummy(object):\n",
    "    def encode(self, gt_data, overlap_threshold=0.5, debug=False):\n",
    "        return np.zeros((10,10))\n",
    "\n",
    "gen = InputGenerator(gt_util_train, PriorUtilDummy(), batch_size=8, input_size=(512, 512), \n",
    "        saturation_var=0.5,\n",
    "        brightness_var=0.5,\n",
    "        contrast_var=0.5,\n",
    "        lighting_std=0.5,\n",
    "        hflip_prob=0.0,\n",
    "        vflip_prob=0.0,\n",
    "        do_crop=True,\n",
    "        add_noise=True,\n",
    "        crop_area_range=[0.9, 1.0],\n",
    "        aspect_ratio_range=[3.9/3, 4.1/3])\n",
    "\n",
    "g = gen.generate()\n",
    "inputs, targets = next(g)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "batch_size = 16\n",
    "print('queue size %0.3f gb' % (inputs[0].nbytes*batch_size/(1024**3),))\n",
    "print('num_batches', gen.num_batches)"
   ]
  },
  {
   "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
}