generate_dataset.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import os\n",
    "from tqdm import tqdm_notebook as tqdm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('train/receipt_00329.json') as json_file:\n",
    "    data = json.load(json_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dict_keys(['dontcare', 'valid_line', 'meta', 'roi', 'repeating_symbol'])"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'version': 'v0.1',\n",
       " 'split': 'train',\n",
       " 'image_id': 329,\n",
       " 'image_size': {'width': 1836, 'height': 3264}}"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data['meta']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{}"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data['roi']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[]"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data['repeating_symbol']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "9"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(data['valid_line'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'words': [{'quad': {'x2': 958,\n",
       "    'y3': 2094,\n",
       "    'x3': 958,\n",
       "    'y4': 2094,\n",
       "    'x1': 819,\n",
       "    'y1': 2033,\n",
       "    'x4': 819,\n",
       "    'y2': 2033},\n",
       "   'is_key': 1,\n",
       "   'row_id': 2163225,\n",
       "   'text': 'TOTAL'},\n",
       "  {'quad': {'x2': 1373,\n",
       "    'y3': 2101,\n",
       "    'x3': 1373,\n",
       "    'y4': 2101,\n",
       "    'x1': 1173,\n",
       "    'y1': 1989,\n",
       "    'x4': 1173,\n",
       "    'y2': 1989},\n",
       "   'is_key': 0,\n",
       "   'row_id': 2163225,\n",
       "   'text': '30.000'}],\n",
       " 'category': 'total.total_price',\n",
       " 'group_id': 26}"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data['valid_line'][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'words': [{'quad': {'x2': 430,\n",
       "    'y3': 2283,\n",
       "    'x3': 437,\n",
       "    'y4': 2296,\n",
       "    'x1': 300,\n",
       "    'y1': 2226,\n",
       "    'x4': 307,\n",
       "    'y2': 2213},\n",
       "   'is_key': 1,\n",
       "   'row_id': 2163226,\n",
       "   'text': 'CASH'},\n",
       "  {'quad': {'x2': 1370,\n",
       "    'y3': 2247,\n",
       "    'x3': 1370,\n",
       "    'y4': 2247,\n",
       "    'x1': 1183,\n",
       "    'y1': 2182,\n",
       "    'x4': 1183,\n",
       "    'y2': 2182},\n",
       "   'is_key': 0,\n",
       "   'row_id': 2163226,\n",
       "   'text': '50.000'}],\n",
       " 'category': 'total.cashprice',\n",
       " 'group_id': 26}"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data['valid_line'][1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'words': [{'quad': {'x2': 512,\n",
       "    'y3': 2340,\n",
       "    'x3': 522,\n",
       "    'y4': 2372,\n",
       "    'x1': 299,\n",
       "    'y1': 2303,\n",
       "    'x4': 309,\n",
       "    'y2': 2272},\n",
       "   'is_key': 1,\n",
       "   'row_id': 2163227,\n",
       "   'text': 'CHANGED'},\n",
       "  {'quad': {'x2': 1371,\n",
       "    'y3': 2319,\n",
       "    'x3': 1369,\n",
       "    'y4': 2314,\n",
       "    'x1': 1186,\n",
       "    'y1': 2249,\n",
       "    'x4': 1185,\n",
       "    'y2': 2254},\n",
       "   'is_key': 0,\n",
       "   'row_id': 2163227,\n",
       "   'text': '20,000'}],\n",
       " 'category': 'total.changeprice',\n",
       " 'group_id': 26}"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data['valid_line'][2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'words': [{'quad': {'x2': 329,\n",
       "    'y3': 1948,\n",
       "    'x3': 329,\n",
       "    'y4': 1948,\n",
       "    'x1': 292,\n",
       "    'y1': 1887,\n",
       "    'x4': 292,\n",
       "    'y2': 1887},\n",
       "   'is_key': 0,\n",
       "   'row_id': 2163224,\n",
       "   'text': '1'}],\n",
       " 'category': 'menu.cnt',\n",
       " 'group_id': 28}"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data['valid_line'][6]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generating the dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def scale(box, size):\n",
    "    return (1000 * box) / size"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "<ipython-input-20-9e2222b121cf>:6: TqdmDeprecationWarning: This function will be removed in tqdm==5.0.0\n",
      "Please use `tqdm.notebook.tqdm` instead of `tqdm.tqdm_notebook`\n",
      "  for js in tqdm(os.listdir(path)):\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0e40d9f3350743648399310b56bfd46e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(HTML(value=''), FloatProgress(value=0.0, max=800.0), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "tokens_train = []\n",
    "boxes_train = []\n",
    "labels_train = []\n",
    "json_train = []\n",
    "path = 'train/'\n",
    "for js in tqdm(os.listdir(path)):\n",
    "    with open(path + js) as f:\n",
    "        json_train.append(json.load(f))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "<ipython-input-21-6327ac5b5e9e>:1: TqdmDeprecationWarning: This function will be removed in tqdm==5.0.0\n",
      "Please use `tqdm.notebook.tqdm` instead of `tqdm.tqdm_notebook`\n",
      "  for js in tqdm(json_train):\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "045d36bd989e460ea6b8bc19e284c01f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(HTML(value=''), FloatProgress(value=0.0, max=800.0), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "for js in tqdm(json_train):\n",
    "    labels = []\n",
    "    tokens = []\n",
    "    boxes = []\n",
    "    width, height = js['meta']['image_size']['width'], js['meta']['image_size']['height']\n",
    "    for elem in js['valid_line']:\n",
    "        for word in elem['words']:\n",
    "            labels.append(elem['category'])\n",
    "            box = word['quad']\n",
    "            y1, y3 = int(scale(box['y1'], height)), int(scale(box['y3'], height))\n",
    "            x1, x3 = int(scale(box['x1'], width)), int(scale(box['x3'], width))\n",
    "            txt = word['text']\n",
    "            tokens.append(txt)\n",
    "            boxes.append([x1, y3, x3, y1])\n",
    "    tokens_train.append(tokens)\n",
    "    labels_train.append(labels)\n",
    "    boxes_train.append(boxes)\n",
    "            \n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['menu.nm',\n",
       " 'menu.nm',\n",
       " 'menu.cnt',\n",
       " 'menu.price',\n",
       " 'total.total_price',\n",
       " 'total.total_price',\n",
       " 'total.cashprice',\n",
       " 'total.cashprice',\n",
       " 'total.menuqty_cnt']"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "labels_train[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['Tebu', 'Lemon', '1', '22.000', 'Total', '22.000', 'CASH', '22.000', '1']"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokens_train[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[[225, 587, 293, 563],\n",
       " [306, 587, 390, 562],\n",
       " [195, 587, 205, 564],\n",
       " [759, 584, 861, 557],\n",
       " [224, 669, 310, 623],\n",
       " [761, 670, 865, 622],\n",
       " [156, 751, 294, 707],\n",
       " [661, 755, 866, 704],\n",
       " [191, 671, 204, 625]]"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "boxes_train[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Validation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "<ipython-input-26-e1e3046c6d63>:6: TqdmDeprecationWarning: This function will be removed in tqdm==5.0.0\n",
      "Please use `tqdm.notebook.tqdm` instead of `tqdm.tqdm_notebook`\n",
      "  for js in tqdm(os.listdir(path)):\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "20d452d6b375481bb312b1f05a056d61",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(HTML(value=''), FloatProgress(value=0.0), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "tokens_val = []\n",
    "labels_val = []\n",
    "boxes_val = []\n",
    "json_val = []\n",
    "path = 'valid/'\n",
    "for js in tqdm(os.listdir(path)):\n",
    "    with open(path + js) as f:\n",
    "        json_val.append(json.load(f))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "<ipython-input-27-6886740b3123>:1: TqdmDeprecationWarning: This function will be removed in tqdm==5.0.0\n",
      "Please use `tqdm.notebook.tqdm` instead of `tqdm.tqdm_notebook`\n",
      "  for js in tqdm(json_val):\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3dbaf20c3de6436ca51735d60bad8790",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(HTML(value=''), FloatProgress(value=0.0), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "for js in tqdm(json_val):\n",
    "    labels = []\n",
    "    tokens = []\n",
    "    boxes = []\n",
    "    width, height = js['meta']['image_size']['width'], js['meta']['image_size']['height']\n",
    "    for elem in js['valid_line']:\n",
    "        for word in elem['words']:\n",
    "            labels.append(elem['category'])\n",
    "            box = word['quad']\n",
    "            y1, y3 = int(scale(box['y1'], height)), int(scale(box['y3'], height))\n",
    "            x1, x3 = int(scale(box['x1'], width)), int(scale(box['x3'], width))\n",
    "            txt = word['text']\n",
    "            tokens.append(txt)\n",
    "            boxes.append([x1, y3, x3, y1])\n",
    "    tokens_val.append(tokens)\n",
    "    labels_val.append(labels)\n",
    "    boxes_val.append(boxes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['menu.cnt',\n",
       " 'menu.nm',\n",
       " 'menu.price',\n",
       " 'menu.cnt',\n",
       " 'menu.nm',\n",
       " 'menu.nm',\n",
       " 'menu.price',\n",
       " 'sub_total.subtotal_price',\n",
       " 'sub_total.subtotal_price',\n",
       " 'sub_total.tax_price',\n",
       " 'sub_total.tax_price',\n",
       " 'sub_total.tax_price',\n",
       " 'sub_total.tax_price',\n",
       " 'total.total_price',\n",
       " 'total.total_price',\n",
       " 'total.menuqty_cnt',\n",
       " 'total.cashprice',\n",
       " 'total.cashprice',\n",
       " 'total.changeprice',\n",
       " 'total.changeprice']"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "labels_val[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['1',\n",
       " 'SAYAP',\n",
       " '13,636',\n",
       " '1',\n",
       " 'PAHA',\n",
       " 'BAWAH',\n",
       " '13,636',\n",
       " 'Subtotal',\n",
       " '27,272',\n",
       " 'P.',\n",
       " 'Resto',\n",
       " '10%',\n",
       " '2,700',\n",
       " 'Total',\n",
       " '30,000',\n",
       " '2',\n",
       " 'Tendered',\n",
       " '50,000',\n",
       " 'Kembali',\n",
       " '20,000']"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokens_val[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[[300, 422, 311, 411],\n",
       " [330, 423, 409, 410],\n",
       " [590, 428, 684, 413],\n",
       " [298, 435, 311, 424],\n",
       " [327, 437, 395, 424],\n",
       " [411, 439, 494, 425],\n",
       " [590, 440, 681, 428],\n",
       " [297, 461, 429, 446],\n",
       " [586, 464, 684, 450],\n",
       " [294, 475, 325, 460],\n",
       " [329, 475, 411, 461],\n",
       " [427, 475, 477, 463],\n",
       " [600, 477, 681, 464],\n",
       " [330, 500, 411, 477],\n",
       " [588, 504, 686, 478],\n",
       " [297, 500, 311, 475],\n",
       " [297, 531, 430, 515],\n",
       " [587, 530, 686, 517],\n",
       " [291, 571, 405, 557],\n",
       " [588, 571, 686, 557]]"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "boxes_val[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "<ipython-input-32-957d45014907>:6: TqdmDeprecationWarning: This function will be removed in tqdm==5.0.0\n",
      "Please use `tqdm.notebook.tqdm` instead of `tqdm.tqdm_notebook`\n",
      "  for js in tqdm(os.listdir(path)):\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e7bc09c7a7f34394995d9a821e33bd4b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(HTML(value=''), FloatProgress(value=0.0), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "tokens_test = []\n",
    "labels_test = []\n",
    "boxes_test = []\n",
    "json_test = []\n",
    "path = 'test/'\n",
    "for js in tqdm(os.listdir(path)):\n",
    "    with open(path + js) as f:\n",
    "        json_test.append(json.load(f))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "<ipython-input-33-e75a102702f0>:1: TqdmDeprecationWarning: This function will be removed in tqdm==5.0.0\n",
      "Please use `tqdm.notebook.tqdm` instead of `tqdm.tqdm_notebook`\n",
      "  for js in tqdm(json_test):\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "85cf6048126942ad990da66eaec09ff3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(HTML(value=''), FloatProgress(value=0.0), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "for js in tqdm(json_test):\n",
    "    labels = []\n",
    "    tokens = []\n",
    "    boxes = []\n",
    "    width, height = js['meta']['image_size']['width'], js['meta']['image_size']['height']\n",
    "    for elem in js['valid_line']:\n",
    "        for word in elem['words']:\n",
    "            labels.append(elem['category'])\n",
    "            box = word['quad']\n",
    "            y1, y3 = int(scale(box['y1'], height)), int(scale(box['y3'], height))\n",
    "            x1, x3 = int(scale(box['x1'], width)), int(scale(box['x3'], width))\n",
    "            txt = word['text']\n",
    "            tokens.append(txt)\n",
    "            boxes.append([x1, y3, x3, y1])\n",
    "    tokens_test.append(tokens)\n",
    "    labels_test.append(labels)\n",
    "    boxes_test.append(boxes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['menu.nm',\n",
       " 'menu.nm',\n",
       " 'menu.cnt',\n",
       " 'menu.unitprice',\n",
       " 'menu.price',\n",
       " 'menu.nm',\n",
       " 'menu.nm',\n",
       " 'menu.cnt',\n",
       " 'menu.unitprice',\n",
       " 'menu.price',\n",
       " 'sub_total.subtotal_price',\n",
       " 'sub_total.subtotal_price',\n",
       " 'sub_total.tax_price',\n",
       " 'sub_total.tax_price',\n",
       " 'sub_total.tax_price',\n",
       " 'sub_total.tax_price',\n",
       " 'total.total_price',\n",
       " 'total.total_price',\n",
       " 'total.total_price',\n",
       " 'total.cashprice',\n",
       " 'total.cashprice',\n",
       " 'total.cashprice',\n",
       " 'total.cashprice']"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "labels_test[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['BASO',\n",
       " 'TAHU',\n",
       " '1',\n",
       " '43,181',\n",
       " '43,181',\n",
       " 'ES',\n",
       " 'JERUK',\n",
       " '1',\n",
       " '13,000',\n",
       " '13,000',\n",
       " 'TOTAL',\n",
       " '56,181',\n",
       " 'TAX',\n",
       " '10.00',\n",
       " '%',\n",
       " '5,618',\n",
       " 'GRAND',\n",
       " 'TOTAL',\n",
       " '61,799',\n",
       " 'TUNAI',\n",
       " '62,000',\n",
       " 'KEMBALI',\n",
       " '201']"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokens_test[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[[145, 510, 233, 483],\n",
       " [243, 505, 333, 483],\n",
       " [528, 496, 554, 472],\n",
       " [564, 496, 686, 471],\n",
       " [771, 492, 888, 466],\n",
       " [145, 547, 195, 514],\n",
       " [206, 543, 312, 517],\n",
       " [533, 527, 560, 503],\n",
       " [570, 529, 695, 503],\n",
       " [777, 529, 902, 497],\n",
       " [138, 580, 254, 550],\n",
       " [792, 564, 917, 529],\n",
       " [133, 655, 206, 625],\n",
       " [217, 652, 332, 624],\n",
       " [342, 648, 370, 618],\n",
       " [834, 637, 935, 604],\n",
       " [119, 691, 239, 665],\n",
       " [253, 688, 370, 659],\n",
       " [820, 671, 952, 640],\n",
       " [119, 733, 239, 700],\n",
       " [831, 712, 978, 678],\n",
       " [109, 775, 278, 743],\n",
       " [902, 750, 981, 718]]"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "boxes_test[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Saving"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "with open('train.pkl', 'wb') as t:\n",
    "    pickle.dump([tokens_train, labels_train, boxes_train], t)\n",
    "with open('val.pkl', 'wb') as t:\n",
    "    pickle.dump([tokens_val, labels_val, boxes_val], t)\n",
    "with open('test.pkl', 'wb') as t:\n",
    "    pickle.dump([tokens_test, labels_test, boxes_test], t)"
   ]
  }
 ],
 "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}