Allahabad_Captcha.ipynb

allahabad_captcha.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Allahabad_Captcha.ipynb",
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/gist/Athe-kunal/f95d494a0a15a5f2dd1a4ede06e13d6d/allahabad_captcha.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "y2gg-b0QvYC6",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import os\n",
        "from torch.utils.data import DataLoader,Dataset\n",
        "import torchvision.transforms as transforms\n",
        "from PIL import Image\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "from torchvision import models\n",
        "import torch.nn as nn\n",
        "import torch\n",
        "from torch.autograd import Variable\n",
        "import cv2\n",
        "import tqdm\n",
        "from tqdm import tqdm\n",
        "\n",
        "path = \"train images/Allahabad train images/Allahabad Train images/Images\"\n",
        "\n",
        "import string\n",
        "all_char = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9']\n",
        "\n",
        "length_all_char = len(all_char)\n",
        "length_of_captcha = len(os.listdir(path)[1][:-4])\n",
        "\n",
        "def encode(a):\n",
        "    one_hot = [0]* length_all_char # It gives a 10 length list\n",
        "    idx = all_char.index(a) # It gives teh index value of the given string in all_char\n",
        "    one_hot[idx]+=1\n",
        "    return one_hot\n",
        "\n",
        "transform = transforms.Compose([\n",
        "    transforms.Resize((256,64)),\n",
        "    transforms.ToTensor(),\n",
        "    transforms.Normalize((0.5,),(0.5,))\n",
        "])\n",
        "\n",
        "class Mydataset(Dataset):\n",
        "    def __init__(self,path,is_train=True,transform=None):\n",
        "        self.path = path\n",
        "        self.img = os.listdir(self.path)\n",
        "        self.transform = transform\n",
        "    def __getitem__(self,idx):\n",
        "        img_path = self.img[idx] \n",
        "        img = cv2.imread(os.path.join(self.path,img_path),cv2.IMREAD_GRAYSCALE)\n",
        "        kernel = cv2.getStructuringElement(cv2.MORPH_RECT,(3,5))\n",
        "        img = cv2.morphologyEx(img, cv2.MORPH_CLOSE, kernel)\n",
        "        label = os.path.join(self.path,img_path)[-(length_of_captcha+4):-4]\n",
        "        img = Image.fromarray(img)\n",
        "        label_ = []\n",
        "        for i in label:\n",
        "            label_ += encode(i)\n",
        "        if self.transform is not None:\n",
        "            img = self.transform(img)    \n",
        "        return img,np.array(label_),label\n",
        "    def __len__(self):\n",
        "        return len(self.img)\n",
        "\n",
        "train_dataset = Mydataset(path,transform = transform)\n",
        "train_loader = DataLoader(train_dataset,batch_size=16,num_workers=0,shuffle=True)\n",
        "\n",
        "class cnnmodel(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(cnnmodel,self).__init__()\n",
        "        self.layer1 = nn.Sequential(\n",
        "            nn.Conv2d(1,16,kernel_size=3,stride=1,padding=2),\n",
        "            nn.BatchNorm2d(16),\n",
        "            nn.ReLU(),\n",
        "            nn.MaxPool2d(kernel_size=2, stride=2))\n",
        "        self.layer2 = nn.Sequential(\n",
        "            nn.Conv2d(16, 32, kernel_size=5, stride=1, padding=2),\n",
        "            nn.BatchNorm2d(32),\n",
        "            nn.ReLU(),\n",
        "            nn.MaxPool2d(kernel_size=2, stride=2))\n",
        "        self.fc1 = nn.Linear(64*16*32,1024)\n",
        "        self.fc2 = nn.Linear(1024,512)\n",
        "        self.fc3 = nn.Linear(512,length_all_char*length_of_captcha)\n",
        "    def forward(self,x):\n",
        "        out = self.layer1(x)\n",
        "        out = self.layer2(out)\n",
        "        out = out.reshape(out.size(0),-1)\n",
        "        out = self.fc1(out)\n",
        "        out = self.fc2(out)\n",
        "        out = self.fc3(out)\n",
        "        return out\n",
        "\n",
        "cnn_model= cnnmodel().cuda()\n",
        "print(cnn_model)\n",
        "\n",
        "loss_fn = nn.MultiLabelSoftMarginLoss()\n",
        "optimizer = torch.optim.Adam(cnn_model.parameters(),lr=0.001)\n",
        "\n",
        "torch.manual_seed(101)\n",
        "num_epochs=35\n",
        "for epoch in tqdm(range(num_epochs)):\n",
        "    for step,i in enumerate(train_loader):\n",
        "        img,label_one_hot,label = i\n",
        "        img = Variable(img).cuda()\n",
        "        label_one_hot= Variable(label_one_hot.float()).cuda()\n",
        "        pred = cnn_model(img)\n",
        "        loss_value = loss_fn(pred,label_one_hot)\n",
        "        optimizer.zero_grad()\n",
        "        loss_value.backward()\n",
        "        optimizer.step()\n",
        "        if epoch%10 == 0 :\n",
        "            print('epoch:', epoch+1, 'step:', step+1, 'loss:', loss_value.item())\n",
        "\n",
        "test_dataset = Mydataset('test images/Allahabad Test images',False,transform)\n",
        "test_loader = DataLoader(test_dataset,batch_size=1,num_workers=0)\n",
        "\n",
        "def divide_list(given_list,n):\n",
        "    for i in range(0,len(given_list),n):\n",
        "        yield str(np.argmax(given_list[i:i+n]))\n",
        "\n",
        "cnn_model.eval()\n",
        "i=0\n",
        "for step,(img,label_one_hot,label) in enumerate(test_loader):\n",
        "    img = Variable(img).cuda()\n",
        "    pred = cnn_model(img)\n",
        "    preds = pred.squeeze().cpu().tolist()\n",
        "    final_pred = \"\".join(list(divide_list(preds,length_all_char)))\n",
        "    if final_pred == label[0]:\n",
        "        pass\n",
        "    else:\n",
        "        i+=1\n",
        "    print(\"label\",label[0], \"pred\",final_pred)\n",
        "print(\"The number of incorrect numbers are :{}/{}\".format(i,step))"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}