train_emnist.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torch.nn.functional as F\n",
    "import numpy as np\n",
    "import os\n",
    "from PIL import Image\n",
    "from PIL import ImageOps\n",
    "from torchvision import datasets, transforms\n",
    "from torch.autograd import Variable\n",
    "import matplotlib.pyplot as plt\n",
    "from torchvision.datasets import ImageFolder\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision.transforms import ToTensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "idx = ['0','1','2','3','4','5','6','7','8','9',\n",
    "       'A','B','C','D','E','F','G','H','I','J','K','L','M','N','O','P','Q','R','S','T','U','V','W','X','Y','Z',\n",
    "       'a','b','d','e','f','g','h','n','q','r','t']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for ix in idx:\n",
    "    root = \"./03_emnist_data/\"+ix\n",
    "    for folder in os.listdir(root):\n",
    "        path = root +'/'+folder\n",
    "        if os.path.isdir(path):\n",
    "            print(path)\n",
    "            for name in os.listdir(path):\n",
    "                ipath = os.path.join(path, name)\n",
    "                #print(ipath)\n",
    "                img = cv2.imread(ipath,0)\n",
    "                img = 255 - img\n",
    "                saveto = root+'/'+name\n",
    "                cv2.imwrite(saveto, img)\n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pil_loader(path):\n",
    "    with open(path, 'rb') as f:\n",
    "        img = Image.open(f)\n",
    "        return img.convert('L')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "mydata = ImageFolder(root=\"03_emnist_data\", transform=ToTensor(), loader=pil_loader)\n",
    "loader = DataLoader(mydata, batch_size=10, shuffle=True, num_workers=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Net(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Net, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(1, 32, kernel_size=5)\n",
    "        self.conv2 = nn.Conv2d(32, 128, kernel_size=5)\n",
    "        self.conv2_drop = nn.Dropout2d()\n",
    "        self.fc1 = nn.Linear(107648, 2048)\n",
    "        self.fc3 = nn.Linear(2048, 512)\n",
    "        self.fc5 = nn.Linear(512, 47)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = F.relu(F.max_pool2d(self.conv1(x), 2))\n",
    "        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))\n",
    "        #print(x.size())\n",
    "        x = x.view(-1, 107648)\n",
    "        x = F.relu(self.fc1(x))\n",
    "        x = F.dropout(x, training=self.training)\n",
    "        x = F.relu(self.fc3(x))\n",
    "        x = F.dropout(x, training=self.training)\n",
    "        x = self.fc5(x)\n",
    "        return x\n",
    "\n",
    "model = Net()\n",
    "model.cuda()\n",
    "\n",
    "optimizer = optim.Adam(model.parameters(), lr=1e-4, eps=1e-4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import clear_output\n",
    "def train(epoch):\n",
    "    model.train()\n",
    "    for batch_idx, (data, target) in enumerate(loader):\n",
    "        data, target = data.cuda(), target.cuda()\n",
    "        data, target = Variable(data), Variable(target)\n",
    "        optimizer.zero_grad()\n",
    "        output = model(data)\n",
    "        loss = nn.CrossEntropyLoss()\n",
    "        output = loss(output, target)\n",
    "        output.backward()\n",
    "        optimizer.step()\n",
    "        if batch_idx %1700 == 0:\n",
    "            clear_output()\n",
    "        if batch_idx % 100 == 0:\n",
    "             print('Train Epoch: {} [{}/{} ({:.0f}%)]\\tLoss: {:.6f}'.format(\n",
    "                 epoch, batch_idx * len(data), len(loader.dataset),\n",
    "                 100. * batch_idx / len(loader), output.data[0]))\n",
    "        #    print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "for epoch in range(5):\n",
    "    train(epoch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(model.state_dict(), 'char_recognizer.pt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Net(\n",
       "  (conv1): Conv2d (1, 32, kernel_size=(5, 5), stride=(1, 1))\n",
       "  (conv2): Conv2d (32, 128, kernel_size=(5, 5), stride=(1, 1))\n",
       "  (conv2_drop): Dropout2d(p=0.5)\n",
       "  (fc1): Linear(in_features=107648, out_features=2048)\n",
       "  (fc3): Linear(in_features=2048, out_features=512)\n",
       "  (fc5): Linear(in_features=512, out_features=47)\n",
       ")"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = Net()\n",
    "model.load_state_dict(torch.load(\"char_recognizer.pt\"))\n",
    "model.cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict_char(gray):\n",
    "    w = gray.size[0]\n",
    "    h = gray.size[1]\n",
    "    gray = gray.convert('L')\n",
    "    gray = gray.point(lambda x: 0 if x<180 else 255, '1')\n",
    "    x= int(64- (w/2))\n",
    "    y = int(64- (h/2))\n",
    "    canvas = Image.new('L', (128, 128), (255))\n",
    "    canvas.paste(gray, box=(x, y))\n",
    "\n",
    "    canvas = ImageOps.invert(canvas)\n",
    "    canvas = np.array(canvas)\n",
    "    canvas = canvas / 255.0\n",
    "    \n",
    "    plt.imshow(canvas)\n",
    "    plt.show()\n",
    "\n",
    "    #test_data = np.array(gray)\n",
    "    test_output = model(Variable(torch.FloatTensor(canvas).unsqueeze(0).unsqueeze(0).cuda()))\n",
    "    pred = test_output.data.max(1, keepdim=True)[1] \n",
    "    pred = np.array(pred).squeeze(0).squeeze(0)\n",
    "    print(idx[pred])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAQUAAAD8CAYAAAB+fLH0AAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADl0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uIDIuMS4xLCBodHRwOi8vbWF0cGxvdGxpYi5vcmcvAOZPmwAADotJREFUeJzt3X/sXXV9x/Hny7aUgWG0aprSktHFRsPMBPINQlwWYzX8mBGWGAMx2rkuzRI28UeiMP8g+08zo2Li2BpRu4WArLLRECfDijH7w86iBIGKdDCktVCMgEYTbOd7f9zDvJ/2W1u+597zvbXPR9J87/mcc+9599Pv99XP+dzzvZ9UFZL0opctdgGSZouhIKlhKEhqGAqSGoaCpIahIKlhKEhqTC0Uklya5JEke5JcN63zSJqsTOPmpSRLgB8AbwX2At8Grq6qhyd+MkkTtXRKr3shsKeqHgNIchtwBTBvKJyS5XUqp0+pFEkAP+PZH1fVq4513LRCYQ3w5Nj2XuAN4wck2QxsBjiV03hDNkypFEkAX6ttTxzPcYs20VhVW6pqrqrmlrF8scqQdJhphcI+4Oyx7bVdm6QZN61Q+DawPsm6JKcAVwHbp3QuSRM0lTmFqjqU5K+Au4ElwOer6qFpnEvSZE1ropGq+grwlWm9vqTp8I5GSQ1DQVLDUJDUMBQkNQwFSQ1DQVLDUJDUMBQkNQwFSQ1DQVLDUJDUMBQkNQwFSQ1DQVLDUJDUMBQkNQwFSQ1DQVLDUJDUMBQkNQwFSQ1DQVLDUJDUMBQkNQwFSY0Fh0KSs5Pcm+ThJA8lubZrX5nkniSPdl9XTK5cSdPWZ6RwCPhQVZ0LXARck+Rc4DpgR1WtB3Z025JOEAsOharaX1Xf6R7/DNgNrAGuALZ2h20FruxbpKThTGSB2STnAOcDO4FVVbW/2/UUsOooz9kMbAY4ldMmUYakCeg90Zjk5cCXgfdX1U/H91VVATXf86pqS1XNVdXcMpb3LUPShPQKhSTLGAXCLVV1R9f8dJLV3f7VwIF+JUoaUp93HwLcDOyuqk+O7doObOwebwTuXHh5kobWZ07hjcC7ge8lub9r+xvgY8DtSTYBTwDv7FeipCEtOBSq6j+BHGX3hoW+rqTF5R2NkhqGgqSGoSCpYShIahgKkhqGgqSGoSCpYShIahgKkhqGgqSGoSCpYShIahgKkhqGgqSGoSCpYShIahgKkhqGgqSGoSCpYShIahgKkhqGgqSGoSCpYShIakxigdklSb6b5K5ue12SnUn2JPlSklP6lylpKJMYKVwL7B7b/jjwqap6NfAssGkC55A0kL6rTq8F/gT4XLcd4M3Atu6QrcCVfc4haVh9RwqfBj4M/KrbfgXwXFUd6rb3Amt6nkPSgPosRf824EBV3bfA529OsivJroO8sNAyJE1Y36Xo357kcuBU4AzgRuDMJEu70cJaYN98T66qLcAWgDOysnrUIWmCFjxSqKrrq2ptVZ0DXAV8vareBdwLvKM7bCNwZ+8qJQ1mGvcpfAT4YJI9jOYYbp7COSRNSZ/Lh/9XVd8AvtE9fgy4cBKvK2l43tEoqWEoSGoYCpIahoKkhqEgqWEoSGoYCpIahoKkhqEgqTGROxo1++7+0f1HtF1y1nmLUIlmnSMFSQ1HCiexu390/6KMFg4ftThimS2OFCQ1DIWTxNH+N777R/fPO98wpMU+v1qGgqSGoaDBXXLWec4jzDBDQVLDUJDUMBQELM6E4/hlxCxMeGrEUJDUMBQkNbyjUYvGy4XZ5EhBUsNQkNQwFCQ1eoVCkjOTbEvy/SS7k1ycZGWSe5I82n1dMaliJU1f35HCjcBXq+q1wOuB3cB1wI6qWg/s6LYlnSAWHApJfhf4Y7oFZKvql1X1HHAFsLU7bCtwZd8iJQ2nz0hhHfAM8IUk303yuSSnA6uqan93zFPAqr5FShpOn1BYClwA3FRV5wM/57BLhaoqoOZ7cpLNSXYl2XWQF3qUIWmS+oTCXmBvVe3strcxComnk6wG6L4emO/JVbWlquaqam4Zy3uUoePlryzreCw4FKrqKeDJJK/pmjYADwPbgY1d20bgzl4VShpU39uc/xq4JckpwGPAexkFze1JNgFPAO/seQ6dwF7qrcwvHu+IZvH0CoWquh+Ym2fXhj6vq9lw+A/opH9XYVqvq368o1FSw9+SPAn9pv+hj7et77k1uxwpSGo4UtDEOAr47WAo6AhD/3BfctZ5R1yi+C7E4vHyQVLDUJDUMBQkNQwFNbyGl6EgqWEoqOEtxzIUTmJeKmg+hoKkhqGgmXC0D4DxcmZ4hoKkhqGgmTLfiMFl6odlKGgm+XmSi8dQkNTwtyQ10xwtDM+RgqSGoaAjOLF3cjMUTnK/aULPcDg5GQqSGk40Cmgn9BwdnNwcKUhq9BopJPkA8BeMVpb+HqNl41YDtwGvAO4D3l1Vv+xZpwbk24AntwWPFJKsAd4HzFXV64AlwFXAx4FPVdWrgWeBTZMoVNIw+l4+LAV+J8lS4DRgP/BmRsvSA2wFrux5DkkD6rMU/T7gE8APGYXB84wuF56rqkPdYXuBNX2LlDScPpcPK4ArgHXAWcDpwKUv4fmbk+xKsusgLyy0DEkT1ufy4S3A41X1TFUdBO4A3gic2V1OAKwF9s335KraUlVzVTW3jOU9ypA0SX1C4YfARUlOSxJgA/AwcC/wju6YjcCd/UqUNKQ+cwo7GU0ofofR25EvA7YAHwE+mGQPo7clb55AnZIG0us+haq6AbjhsObHgAv7vK6kxeMdjZIahoKkhqEgqWEoSGoYCpIahoKkhqEgqWEoSGoYCpIahoKkhqEgqWEoSGoYCpIahoKkhqEgqWEoSGoYCpIahoKkhqEgqWEoSGoYCpIahoKkhqEgqWEoSGoYCpIaxwyFJJ9PciDJg2NtK5Pck+TR7uuKrj1JPpNkT5IHklwwzeIlTd7xjBS+yJFLzF8H7Kiq9cCObhvgMmB992czcNNkypQ0lGOGQlV9E/jJYc1XAFu7x1uBK8fa/6lGvsVoWfrVkypW0vQtdE5hVVXt7x4/BazqHq8Bnhw7bm/XJukE0XuisaoKqJf6vCSbk+xKsusgL/QtQ9KELDQUnn7xsqD7eqBr3wecPXbc2q7tCFW1parmqmpuGcsXWIakSVtoKGwHNnaPNwJ3jrW/p3sX4iLg+bHLDEkngKXHOiDJrcCbgFcm2QvcAHwMuD3JJuAJ4J3d4V8BLgf2AL8A3juFmiVN0TFDoaquPsquDfMcW8A1fYuStHi8o1FSw1CQ1DAUJDUMBUkNQ0FSw1CQ1DAUJDUMBUkNQ0FSw1CQ1DAUJDUMBUkNQ0FSw1CQ1DAUJDUMBUkNQ0FSw1CQ1DAUJDUMBUkNQ0FSw1CQ1DAUJDUMBUkNQ0FS45ihkOTzSQ4keXCs7e+SfD/JA0n+NcmZY/uuT7InySNJLplW4ZKm43hGCl8ELj2s7R7gdVX1h8APgOsBkpwLXAX8Qfecv0+yZGLVSpq6Y4ZCVX0T+Mlhbf9RVYe6zW8xWnIe4Argtqp6oaoeZ7TQ7IUTrFfSlE1iTuHPgX/vHq8Bnhzbt7drk3SCOOaq079Jko8Ch4BbFvDczcBmgFM5rU8ZkiZowaGQ5M+AtwEbuiXoAfYBZ48dtrZrO0JVbQG2AJyRlTXfMZKGt6DLhySXAh8G3l5VvxjbtR24KsnyJOuA9cB/9S9T0lCOOVJIcivwJuCVSfYCNzB6t2E5cE8SgG9V1V9W1UNJbgceZnRZcU1V/e+0ipc0efn1yH/xnJGV9YZsWOwypN9qX6tt91XV3LGO845GSQ1DQVLDUJDUMBQkNQwFSQ1DQVLDUJDUMBQkNWbi5qUkzwA/B3682LUAr8Q6xllH60Su4/eq6lXHOmgmQgEgya7judvKOqzDOqZbh5cPkhqGgqTGLIXClsUuoGMdLeto/dbXMTNzCpJmwyyNFCTNgJkIhSSXdutE7Ely3UDnPDvJvUkeTvJQkmu79pVJ7knyaPd1xUD1LEny3SR3ddvrkuzs+uRLSU4ZoIYzk2zr1vTYneTixeiPJB/o/k0eTHJrklOH6o+jrHMybx9k5DNdTQ8kuWDKdQyy3sqih0K3LsRngcuAc4Gru/Ujpu0Q8KGqOhe4CLimO+91wI6qWg/s6LaHcC2we2z748CnqurVwLPApgFquBH4alW9Fnh9V8+g/ZFkDfA+YK6qXgcsYbSWyFD98UWOXOfkaH1wGaOPHFzP6EOIb5pyHcOst1JVi/oHuBi4e2z7euD6RajjTuCtwCPA6q5tNfDIAOdey+ib7c3AXUAY3ZiydL4+mlINvws8TjfPNNY+aH/w62UCVjL6uMC7gEuG7A/gHODBY/UB8I/A1fMdN406Dtv3p8At3ePmZwa4G7h4oedd9JECM7BWRJJzgPOBncCqqtrf7XoKWDVACZ9m9EG4v+q2XwE8V79ecGeIPlkHPAN8obuM+VyS0xm4P6pqH/AJ4IfAfuB54D6G749xR+uDxfzendp6K7MQCosqycuBLwPvr6qfju+rUexO9e2ZJG8DDlTVfdM8z3FYClwA3FRV5zO67by5VBioP1YwWmlsHXAWcDpHDqMXzRB9cCx91ls5HrMQCse9VsSkJVnGKBBuqao7uuank6zu9q8GDky5jDcCb0/yP8BtjC4hbgTOTPLip20P0Sd7gb1VtbPb3sYoJIbuj7cAj1fVM1V1ELiDUR8N3R/jjtYHg3/vjq238q4uoCZexyyEwreB9d3s8imMJky2T/ukGX02/c3A7qr65Niu7cDG7vFGRnMNU1NV11fV2qo6h9Hf/etV9S7gXuAdA9bxFPBkktd0TRsYfVT/oP3B6LLhoiSndf9GL9YxaH8c5mh9sB14T/cuxEXA82OXGRM32Hor05w0egkTKpczmk39b+CjA53zjxgNAx8A7u/+XM7oen4H8CjwNWDlgP3wJuCu7vHvd/+we4B/AZYPcP7zgF1dn/wbsGIx+gP4W+D7wIPAPzNaY2SQ/gBuZTSXcZDR6GnT0fqA0YTwZ7vv2+8xesdkmnXsYTR38OL36z+MHf/Rro5HgMv6nNs7GiU1ZuHyQdIMMRQkNQwFSQ1DQVLDUJDUMBQkNQwFSQ1DQVLj/wDPFbrqKKGO6AAAAABJRU5ErkJggg==\n",
      "text/plain": [
       "<matplotlib.figure.Figure at 0x7fcfadef2240>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "H\n"
     ]
    }
   ],
   "source": [
    "pil_im =  Image.open(\"ref.jpg\")\n",
    "predict_char(pil_im)"
   ]
  }
 ],
 "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
}