data_prepare.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "import os \n",
    "\n",
    "dst = \"./output/train\" # resized\n",
    "\n",
    "if not os.path.exists(dst):\n",
    "    os.mkdir(dst)\n",
    "    \n",
    "dst = \"./output/test\" # resized\n",
    "\n",
    "if not os.path.exists(dst):\n",
    "    os.mkdir(dst)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: tqdm in /usr/local/lib/python2.7/site-packages\r\n"
     ]
    }
   ],
   "source": [
    "! pip install tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3755\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import struct\n",
    "from PIL import Image\n",
    "from tqdm import *\n",
    "\n",
    "data_dir = '../input'\n",
    "train_data_dir = os.path.join(data_dir, 'trn')\n",
    "test_data_dir = os.path.join(data_dir, 'tst')\n",
    "\n",
    "\n",
    "def read_from_gnt_dir(gnt_dir=train_data_dir):\n",
    "    def one_file(f):\n",
    "        header_size = 10\n",
    "        while True:\n",
    "            header = np.fromfile(f, dtype='uint8', count=header_size)\n",
    "            if not header.size: break\n",
    "            sample_size = header[0] + (header[1]<<8) + (header[2]<<16) + (header[3]<<24)\n",
    "            tagcode = header[5] + (header[4]<<8)\n",
    "            width = header[6] + (header[7]<<8)\n",
    "            height = header[8] + (header[9]<<8)\n",
    "            if header_size + width*height != sample_size:\n",
    "                break\n",
    "            image = np.fromfile(f, dtype='uint8', count=width*height).reshape((height, width))\n",
    "            yield image, tagcode\n",
    "    for file_name in os.listdir(gnt_dir):\n",
    "        if file_name.endswith('.gnt'):\n",
    "            file_path = os.path.join(gnt_dir, file_name)\n",
    "            with open(file_path, 'rb') as f:\n",
    "                for image, tagcode in one_file(f):\n",
    "                    yield image, tagcode\n",
    "                    \n",
    "char_set = set()\n",
    "for _, tagcode in read_from_gnt_dir(gnt_dir=train_data_dir):\n",
    "    tagcode_unicode = struct.pack('>H', tagcode).decode('gb2312')\n",
    "    char_set.add(tagcode_unicode)\n",
    "char_list = list(char_set)\n",
    "char_dict = dict(zip(sorted(char_list), range(len(char_list))))\n",
    "print len(char_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import pickle\n",
    "f = open('char_dict', 'wb')\n",
    "pickle.dump(char_dict, f)\n",
    "f.close()\n",
    "train_counter = 0\n",
    "test_counter = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "223991it [1:01:21, 60.85it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 4min 22s, sys: 39.9 s, total: 5min 2s\n",
      "Wall time: 1h 1min 21s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "for image, tagcode in tqdm(read_from_gnt_dir(gnt_dir=test_data_dir)):\n",
    "    tagcode_unicode = struct.pack('>H', tagcode).decode('gb2312')\n",
    "    im = Image.fromarray(image)\n",
    "    dir_name = './output/test/' + '%0.5d'%char_dict[tagcode_unicode]\n",
    "    if not os.path.exists(dir_name):\n",
    "        os.mkdir(dir_name)\n",
    "    im.convert('RGB').save(dir_name+'/' + str(test_counter) + '.png')\n",
    "    test_counter += 1\n",
    "    #print test_counter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "509685it [2:18:30, 61.33it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 9min 56s, sys: 1min 31s, total: 11min 27s\n",
      "Wall time: 2h 18min 30s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "for image, tagcode in tqdm(read_from_gnt_dir(gnt_dir=train_data_dir)):\n",
    "    tagcode_unicode = struct.pack('>H', tagcode).decode('gb2312')\n",
    "    im = Image.fromarray(image)\n",
    "    dir_name = './output/train/' + '%0.5d'%char_dict[tagcode_unicode]\n",
    "    if not os.path.exists(dir_name):\n",
    "        os.mkdir(dir_name)\n",
    "    im.convert('RGB').save(dir_name+'/' + str(train_counter) + '.png')\n",
    "    train_counter += 1\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}