main.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from collections import defaultdict\n",
    "from scipy.special import logsumexp\n",
    "import torchvision.transforms as transforms\n",
    "import matplotlib.pyplot as plt\n",
    "from tqdm import tqdm\n",
    "from torch.utils.data import Dataset, DataLoader,random_split\n",
    "import numpy as np\n",
    "from torchmetrics.text import CharErrorRate\n",
    "metric = CharErrorRate()\n",
    "\n",
    "from PIL import Image\n",
    "from config import ModelConfigs\n",
    "configs = ModelConfigs()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomDataset(Dataset):\n",
    "    def __init__(self, image_folder, label_file, transform=None):\n",
    "        self.image_folder = image_folder\n",
    "        self.label_file = label_file\n",
    "        self.transform = transform\n",
    "\n",
    "        with open(label_file, 'r', encoding=\"utf8\") as f:\n",
    "            self.data = f.readlines()\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.data)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        line = self.data[idx]\n",
    "        line = line.rstrip('\\n')\n",
    "        line = line.split(\"\\t\")\n",
    "        \n",
    "        image_path = line[0]\n",
    "        label = ' '.join(line[1:])\n",
    "\n",
    "        target = [configs.CHAR2LABEL[c] for c in label]\n",
    "        \n",
    "        target_length = [len(target)]\n",
    "\n",
    "        target = torch.LongTensor(target)\n",
    "        target_length = torch.LongTensor(target_length)\n",
    "                    \n",
    "        image = Image.open(os.path.join(self.image_folder, image_path))\n",
    "        image = self.transform(image)\n",
    "\n",
    "        return image, target, target_length\n",
    "    \n",
    "def synth90k_collate_fn(batch):\n",
    "    images, targets, target_lengths = zip(*batch)\n",
    "    images = torch.stack(images, 0)\n",
    "    targets = torch.cat(targets, 0)\n",
    "    target_lengths = torch.cat(target_lengths, 0)\n",
    "    return images, targets, target_lengths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_transform = transforms.Compose([\n",
    "    transforms.Resize((configs.height, configs.width)),\n",
    "    transforms.Grayscale(1),\n",
    "    transforms.ToTensor(),      \n",
    "    transforms.Normalize((0.5,), (0.5,))\n",
    "\n",
    "])\n",
    "\n",
    "image_folder = \"data_set/new_train\"\n",
    "label_file = \"data_set/train_gt.txt\"\n",
    "custom_dataset = CustomDataset(image_folder, label_file, transform=data_transform)\n",
    "train_size = int(0.9 * len(custom_dataset))\n",
    "test_size = len(custom_dataset) - train_size\n",
    "train_dataset, test_dataset =random_split(custom_dataset, [train_size, test_size])\n",
    "batch_size = 20\n",
    "trainloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True,collate_fn=synth90k_collate_fn)\n",
    "testloader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False,collate_fn=synth90k_collate_fn)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(custom_dataset[0][0].squeeze(),cmap=\"gray\")\n",
    "custom_dataset[0][0].squeeze()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "NINF = -1 * float('inf')\n",
    "DEFAULT_EMISSION_THRESHOLD = 0.01\n",
    "def _reconstruct(labels, blank=0):\n",
    "    new_labels = []\n",
    "    # merge same labels\n",
    "    previous = None\n",
    "    for l in labels:\n",
    "        if l != previous:\n",
    "            new_labels.append(l)\n",
    "            previous = l\n",
    "    # delete blank\n",
    "    new_labels = [l for l in new_labels if l != blank]\n",
    "\n",
    "    return new_labels\n",
    "\n",
    "\n",
    "def greedy_decode(emission_log_prob, blank=0, **kwargs):\n",
    "    labels = np.argmax(emission_log_prob, axis=-1)\n",
    "    labels = _reconstruct(labels, blank=blank)\n",
    "    return labels\n",
    "\n",
    "def prefix_beam_decode(emission_log_prob, blank=0, **kwargs):\n",
    "    beam_size = 5\n",
    "    emission_threshold = kwargs.get('emission_threshold', np.log(DEFAULT_EMISSION_THRESHOLD))\n",
    "\n",
    "    length, class_count = emission_log_prob.shape\n",
    "\n",
    "    beams = [(tuple(), (0, NINF))]  # (prefix, (blank_log_prob, non_blank_log_prob))\n",
    "    # initial of beams: (empty_str, (log(1.0), log(0.0)))\n",
    "\n",
    "    for t in range(length):\n",
    "        new_beams_dict = defaultdict(lambda: (NINF, NINF))  # log(0.0) = NINF\n",
    "\n",
    "        for prefix, (lp_b, lp_nb) in beams:\n",
    "            for c in range(class_count):\n",
    "                log_prob = emission_log_prob[t, c]\n",
    "                if log_prob < emission_threshold:\n",
    "                    continue\n",
    "\n",
    "                end_t = prefix[-1] if prefix else None\n",
    "\n",
    "                # if new_prefix == prefix\n",
    "                new_lp_b, new_lp_nb = new_beams_dict[prefix]\n",
    "\n",
    "                if c == blank:\n",
    "                    new_beams_dict[prefix] = (\n",
    "                        logsumexp([new_lp_b, lp_b + log_prob, lp_nb + log_prob]),\n",
    "                        new_lp_nb\n",
    "                    )\n",
    "                    continue\n",
    "                if c == end_t:\n",
    "                    new_beams_dict[prefix] = (\n",
    "                        new_lp_b,\n",
    "                        logsumexp([new_lp_nb, lp_nb + log_prob])\n",
    "                    )\n",
    "\n",
    "                # if new_prefix == prefix + (c,)\n",
    "                new_prefix = prefix + (c,)\n",
    "                new_lp_b, new_lp_nb = new_beams_dict[new_prefix]\n",
    "\n",
    "                if c != end_t:\n",
    "                    new_beams_dict[new_prefix] = (\n",
    "                        new_lp_b,\n",
    "                        logsumexp([new_lp_nb, lp_b + log_prob, lp_nb + log_prob])\n",
    "                    )\n",
    "                else:\n",
    "                    new_beams_dict[new_prefix] = (\n",
    "                        new_lp_b,\n",
    "                        logsumexp([new_lp_nb, lp_b + log_prob])\n",
    "                    )\n",
    "\n",
    "        # sorted by log(blank_prob + non_blank_prob)\n",
    "        beams = sorted(new_beams_dict.items(), key=lambda x: logsumexp(x[1]), reverse=True)\n",
    "        beams = beams[:beam_size]\n",
    "\n",
    "    labels = list(beams[0][0])\n",
    "    return labels\n",
    "\n",
    "def ctc_decode(log_probs, label2char=None, blank=0, method='g', beam_size=10):\n",
    "    emission_log_probs = np.transpose(log_probs.cpu().numpy(), (1, 0, 2))\n",
    "    # size of emission_log_probs: (batch, length, class)\n",
    "\n",
    "    decoders = {\n",
    "        'g': greedy_decode,\n",
    "        'b': prefix_beam_decode,\n",
    "    }\n",
    "    decoder = decoders[method]\n",
    "\n",
    "    decoded_list = []\n",
    "    for emission_log_prob in emission_log_probs:\n",
    "        decoded = decoder(emission_log_prob, blank=blank, beam_size=beam_size)\n",
    "        if label2char:\n",
    "            decoded = [label2char[l] for l in decoded]\n",
    "        decoded_list.append(decoded)\n",
    "    return decoded_list\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CRNN(nn.Module):\n",
    "\n",
    "    def __init__(self, img_channel, img_height, img_width, num_class,\n",
    "                 map_to_seq_hidden=128, rnn_hidden=512, leaky_relu=False):\n",
    "        super(CRNN, self).__init__()\n",
    "\n",
    "        self.cnn, (output_channel, output_height, output_width) = \\\n",
    "            self._cnn_backbone(img_channel, img_height, img_width, leaky_relu)\n",
    "\n",
    "        self.map_to_seq = nn.Linear(output_channel * output_height, map_to_seq_hidden)\n",
    "\n",
    "        self.rnn1 = nn.LSTM(map_to_seq_hidden, rnn_hidden, bidirectional=True)\n",
    "        self.rnn2 = nn.LSTM(2*rnn_hidden, rnn_hidden, bidirectional=True)\n",
    "\n",
    "        self.dense = nn.Linear(2*rnn_hidden, num_class)\n",
    "\n",
    "    def _cnn_backbone(self, img_channel, img_height, img_width, leaky_relu):\n",
    "        assert img_height % 16 == 0\n",
    "        assert img_width % 4 == 0\n",
    "\n",
    "        channels = [img_channel, 64, 128, 256, 256, 512, 512, 1024]\n",
    "        kernel_sizes = [3, 3, 3, 3, 3, 3, 2]\n",
    "        strides = [1, 1, 1, 1, 1, 1, 1]\n",
    "        paddings = [1, 1, 1, 1, 1, 1, 0]\n",
    "\n",
    "        cnn = nn.Sequential()\n",
    "\n",
    "        def conv_relu(i, batch_norm=False):\n",
    "            # shape of input: (batch, input_channel, height, width)\n",
    "            input_channel = channels[i]\n",
    "            output_channel = channels[i+1]\n",
    "\n",
    "            cnn.add_module(\n",
    "                f'conv{i}',\n",
    "                nn.Conv2d(input_channel, output_channel, kernel_sizes[i], strides[i], paddings[i])\n",
    "            )\n",
    "\n",
    "            if batch_norm:\n",
    "                cnn.add_module(f'batchnorm{i}', nn.BatchNorm2d(output_channel))\n",
    "\n",
    "            relu = nn.LeakyReLU(0.2, inplace=True) if leaky_relu else nn.ReLU(inplace=True)\n",
    "            cnn.add_module(f'relu{i}', relu)\n",
    "\n",
    "        # size of image: (channel, height, width) = (img_channel, img_height, img_width)\n",
    "        \n",
    "        conv_relu(0)\n",
    "        cnn.add_module('pooling0', nn.MaxPool2d(kernel_size=2, stride=2))\n",
    "        # (64, img_height // 2, img_width // 2)\n",
    "\n",
    "        conv_relu(1)\n",
    "        cnn.add_module('pooling1', nn.MaxPool2d(kernel_size=2, stride=2))\n",
    "        # (128, img_height // 4, img_width // 4)\n",
    "\n",
    "        conv_relu(2)\n",
    "        conv_relu(3,batch_norm=True)\n",
    "        cnn.add_module(\n",
    "            'pooling2',\n",
    "            nn.MaxPool2d(kernel_size=(2, 1))\n",
    "        )  # (256, img_height // 8, img_width // 4)\n",
    "\n",
    "        conv_relu(4, batch_norm=True)\n",
    "        conv_relu(5, batch_norm=True)\n",
    "        cnn.add_module(\n",
    "            'pooling3',\n",
    "            nn.MaxPool2d(kernel_size=(2, 1))\n",
    "        )  # (512, img_height // 16, img_width // 4)\n",
    "\n",
    "        conv_relu(6)  # (512, img_height // 16 - 1, img_width // 4 - 1)\n",
    "\n",
    "        output_channel, output_height, output_width = \\\n",
    "            channels[-1], img_height // 16 - 1, img_width // 4 - 1\n",
    "        return cnn, (output_channel, output_height, output_width)\n",
    "\n",
    "    def forward(self, images):\n",
    "        # shape of images: (batch, channel, height, width)\n",
    "\n",
    "        conv = self.cnn(images)\n",
    "        batch, channel, height, width = conv.size()\n",
    "\n",
    "        conv = conv.view(batch, channel * height, width)\n",
    "        conv = conv.permute(2, 0, 1)  # (width, batch, feature)\n",
    "        seq = self.map_to_seq(conv)\n",
    "\n",
    "        recurrent, _ = self.rnn1(seq)\n",
    "        recurrent, _ = self.rnn2(recurrent)\n",
    "\n",
    "        output = self.dense(recurrent)\n",
    "        return output  # shape: (seq_len, batch, num_class)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "model2 = CRNN(1,configs.height,configs.width,num_class=len(configs.CHARS)+1,leaky_relu=False)\n",
    "#model2.load_state_dict(torch.load(\"models/08_handwriting_recognition_torch\\crnn_4_loss0.5615450739860535.pt\"))\n",
    "model2.to(device)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def train_batch(crnn, data, optimizer, criterion, device):\n",
    "    crnn.train()\n",
    "    images, targets, target_lengths = [d.to(device) for d in data]\n",
    "\n",
    "    logits = crnn(images)\n",
    "    log_probs = torch.nn.functional.log_softmax(logits, dim=2)\n",
    "\n",
    "    batch_size = images.size(0)\n",
    "    input_lengths = torch.LongTensor([logits.size(0)] * batch_size)\n",
    "    target_lengths = torch.flatten(target_lengths)\n",
    "\n",
    "    loss = criterion(log_probs, targets, input_lengths, target_lengths)\n",
    "\n",
    "    optimizer.zero_grad()\n",
    "    loss.backward()\n",
    "    torch.nn.utils.clip_grad_norm_(crnn.parameters(), 5) # gradient clipping with 5\n",
    "    optimizer.step()\n",
    "    return loss.item()\n",
    "\n",
    "def test_batch(crnn, data,criterion,device):\n",
    "    crnn.eval()\n",
    "    crnn.to(device)\n",
    "    tot_count = 0\n",
    "    tot_loss = 0\n",
    "    tot_correct = 0\n",
    "    wrong_cases = []\n",
    "\n",
    "    pbar_total =len(testloader)\n",
    "    pbar = tqdm(total=pbar_total, desc=\"Evaluate\")\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        for i, data in enumerate(testloader):\n",
    "\n",
    "            images, targets, target_lengths = [d.to(device) for d in data]\n",
    "\n",
    "            logits = crnn(images)\n",
    "            log_probs = torch.nn.functional.log_softmax(logits, dim=2)\n",
    "\n",
    "            batch_size = images.size(0)\n",
    "            input_lengths = torch.LongTensor([logits.size(0)] * batch_size)\n",
    "\n",
    "            loss = criterion(log_probs, targets, input_lengths, target_lengths)\n",
    "\n",
    "            preds = ctc_decode(log_probs,method=\"g\")\n",
    "            reals = targets.cpu().numpy().tolist()\n",
    "            target_lengths = target_lengths.cpu().numpy().tolist()\n",
    "\n",
    "            tot_count += batch_size\n",
    "            tot_loss += loss.item()\n",
    "            target_length_counter = 0\n",
    "            for pred, target_length in zip(preds, target_lengths):\n",
    "                real = reals[target_length_counter:target_length_counter + target_length]\n",
    "                target_length_counter += target_length\n",
    "                if pred == real:\n",
    "                    tot_correct += 1\n",
    "                else:\n",
    "                    wrong_cases.append((real, pred))\n",
    "            \n",
    "            pbar.update(1)\n",
    "        pbar.close()\n",
    "\n",
    "        evaluation = {\n",
    "            'loss': tot_loss / tot_count,\n",
    "            'acc': tot_correct / tot_count,\n",
    "            'wrong_cases': wrong_cases\n",
    "        }\n",
    "\n",
    "        return evaluation\n",
    "\n",
    "optimizer = optim.RMSprop(model2.parameters(), lr=configs.learning_rate)\n",
    "criterion = nn.CTCLoss()\n",
    "criterion.to(device)\n",
    "\n",
    "i = 1\n",
    "for epoch in range(1, 20 + 1):\n",
    "    print(f'epoch: {epoch}')\n",
    "    tot_train_loss = 0.\n",
    "    tot_train_count = 0\n",
    "    for train_data in tqdm(trainloader):\n",
    "        \n",
    "        loss = train_batch(model2, train_data, optimizer, criterion, device)\n",
    "        train_size = train_data[0].size(0)\n",
    "        \n",
    "        tot_train_loss += loss\n",
    "        tot_train_count += train_size\n",
    "        \n",
    "    if i % 1 == 0 :\n",
    "                    prefix = 'crnn'\n",
    "                    \n",
    "                    save_model_path = (configs.model_path+f'/{prefix}_{i}_loss{loss}.pt')\n",
    "                    torch.save(model2.state_dict(), save_model_path)\n",
    "                    print('save model at ', save_model_path)\n",
    "    i += 1\n",
    "\n",
    "    print('train_loss: ', tot_train_loss / tot_train_count)\n",
    "\n",
    "    evaluation = test_batch(model2,testloader,criterion,device)\n",
    "    print(evaluation)\n",
    "        "
   ]
  }
 ],
 "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.10.4"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}