TensorFlow_Captcha_OCR.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "private_outputs": true,
      "provenance": []
    },
    "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.7.0"
    },
    "accelerator": "GPU",
    "gpuClass": "standard"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/gist/hw0603/51a623121fed39a1d93912930a7bde51/tensorflow-captcha-ocr.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FHM-Ac5IfCa3"
      },
      "source": [
        "# TensorFlow Captcha OCR\n",
        "\n",
        "**제작:** Hyunguk Ryu (@hw0603)<br>\n",
        "**생성일:** 2021/08/26<br>\n",
        "**최근 수정일:** 2022/08/12<br>\n",
        "**설명:** CNN, RNN, CTC loss를 사용하여 자동입력방지문자 인식을 위한 OCR 모델을 구현합니다."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kN1nQau-fCa6"
      },
      "source": [
        "## 개요\n",
        "\n",
        "Keras의 함수형 API를 활용하여 간단한 OCR 모델을 생성합니다.  \n",
        "CNN과 RNN을 활용하는 것 외에도 새로운 레이어 클래스를 만들고, 이것을 CTC Loss 구현을 위한 \"Endpoint Layer\"로 사용하는 방법도 보여 줍니다.  \n",
        "새로운 레이어를 만드는 방법에 관한 자세한 설명은 [이곳](https://keras.io/guides/making_new_layers_and_models_via_subclassing/)을 확인하세요."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_X0mgoUOfxdJ"
      },
      "source": [
        "## 런타임 설정\n",
        "TensorFlow 학습을 위해 GPU가 필요합니다.  \n",
        "런타임->런타임 유형 변경 메뉴에서 하드웨어 가속기를 GPU로 설정해 주세요  \n",
        "(Tesla K80 << Tesla P4 <= Tesla T4 << Tesla P100 << Tesla V100)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "duMHZan6gD2_"
      },
      "source": [
        "# 할당된 GPU 확인\n",
        "!nvidia-smi --query-gpu=gpu_name,driver_version,memory.total --format=csv"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "O6ZoJ_tATf8K"
      },
      "source": [
        "####[TIP] Colab 런타임 할당 해제 방지\n",
        "장시간 대기 시 Colab 런타임 할당 해제를 방지하기 위해 다음 코드를 브라우저 콘솔에서 실행해 주세요\n",
        "\n",
        "```\n",
        "function ClickConnect(){\n",
        "    console.log(\"코랩 연결 끊김 방지\");\n",
        "    document.querySelector(\"colab-toolbar-button#connect\").click()\n",
        "}\n",
        "setInterval(ClickConnect, 60 * 1000)\n",
        "```\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "V76GJCLtHK78"
      },
      "source": [
        "## 구글 드라이브 마운트"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "j36vnKUyHCK5"
      },
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nD_5u5IZfCa7"
      },
      "source": [
        "## 모듈 임포트"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BMDVM_LrfCa8"
      },
      "source": [
        "import os\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        " \n",
        "from pathlib import Path\n",
        "from collections import Counter\n",
        " \n",
        "import tensorflow as tf\n",
        "from tensorflow import keras\n",
        "from tensorflow.keras import layers"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WkDz4K4BfCa9"
      },
      "source": [
        "## 학습 데이터 압축 해제 & 로드"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "oXEGMPOkfCa-"
      },
      "source": [
        "!rm -rf /content/captcha_images\n",
        "%cd \"/content/drive/MyDrive/Colab Notebooks\"\n",
        "!unzip -qq captcha_images.zip -d /content/captcha_images"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QLjrCm-1fCa_"
      },
      "source": [
        "이 데이터 셋은 약 5000개의 라벨링된 `PNG(RGBA 32)` 이미지로 구성되어 있습니다.  \n",
        "각 샘플에 대한 label은 문자로 구성되어 있으며, 파일명에서 확인할 수 있습니다.  \n",
        "모델을 훈련시키기 위해 label의 문자를 숫자로 매핑한 후, 예측을 위해 다시 숫자를 문자로 매핑합니다.  \n",
        "이를 위해서 문자를 정수로, 정수를 문자로 매핑하는 두 개의 딕셔너리를 만들 것입니다."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HUnJSzVefCbA"
      },
      "source": [
        "# 데이터 경로\n",
        "data_dir = Path(\"/content/captcha_images/new_server/labeled_image/\")\n",
        " \n",
        "# 모든 이미지들의 리스트 생성\n",
        "images = sorted(list(map(str, list(data_dir.glob(\"*.png\")))))\n",
        "labels = [img.split(os.path.sep)[-1].split(\".png\")[0] for img in images]\n",
        "characters = list(set(char for label in labels for char in label))\n",
        "characters.sort()\n",
        " \n",
        "print(\"Number of images found: \", len(images))\n",
        "print(\"Number of labels found: \", len(labels))\n",
        "print(\"Number of unique characters: \", len(characters))\n",
        "print(\"Characters present: \", characters)\n",
        "print(\"Characters: \", \"\".join(characters)) # 추론 시에 같은 characters 데이터가 있어야 함\n",
        " \n",
        "# 학습과 검증에 사용될 batch size\n",
        "batch_size = 16\n",
        " \n",
        "# 이미지 크기 설정\n",
        "img_width = 140\n",
        "img_height = 35\n",
        " \n",
        "# 이미지가 convolutional blocks에 의해 downsample되는 비율을 2로 설정\n",
        "# 두 번의 convolutional blocks를 사용할 것이기 때문에\n",
        "# 이미지는 한 변을 기준으로 4배 줄어듦.\n",
        "downsample_factor = 4\n",
        " \n",
        "# 라벨 중 가장 긴 것의 길이 구함\n",
        "max_length = max([len(label) for label in labels])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Rq7qIzKgfCbB"
      },
      "source": [
        "## 데이터 전처리"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PrSLf5U0fCbC"
      },
      "source": [
        "# 문자를 숫자로 매핑\n",
        "char_to_num = layers.StringLookup(\n",
        "    vocabulary=list(characters), mask_token=None\n",
        ")\n",
        " \n",
        "# 숫자를 문자로 매핑\n",
        "num_to_char = layers.StringLookup(\n",
        "    vocabulary=char_to_num.get_vocabulary(), mask_token=None, invert=True\n",
        ")\n",
        " \n",
        " \n",
        "def split_data(images, labels, train_size=0.9, shuffle=True):\n",
        "    # 1. Dataset의 전체 크기 구함\n",
        "    size = len(images)\n",
        "    # 2. Dataset의 인덱스를 담은 np.array 생성 (필요 시 셔플)\n",
        "    indices = np.arange(size)\n",
        "    if shuffle:\n",
        "        np.random.shuffle(indices)\n",
        "    # 3. 비율에 맞게 train set 사이즈 설정\n",
        "    train_samples = int(size * train_size)\n",
        "    # 4. train set과 validation set 분리\n",
        "    x_train, y_train = images[indices[:train_samples]], labels[indices[:train_samples]]\n",
        "    x_valid, y_valid = images[indices[train_samples:]], labels[indices[train_samples:]]\n",
        "    return x_train, x_valid, y_train, y_valid\n",
        " \n",
        " \n",
        "# 반환된 데이터 저장\n",
        "x_train, x_valid, y_train, y_valid = split_data(np.array(images), np.array(labels))\n",
        " \n",
        " \n",
        "def encode_single_sample(img_path, label):\n",
        "    # 1. 이미지 로드\n",
        "    img = tf.io.read_file(img_path)\n",
        "\n",
        "    # 2. PNG 이미지 디코드 (원본 파일의 채널(RGBA 4채널) 그대로 사용)\n",
        "    img = tf.io.decode_png(img, channels=0)\n",
        "\n",
        "    # 3. R, G, B, A 채널을 분리한 후 A 채널만 추출하여 그레이스캐일로 변환\n",
        "    r, g, b, a = img[:, :, 0], img[:, :, 1], img[:, :, 2], img[:, :, 3]\n",
        "    rgb = tf.stack([a], axis=-1)\n",
        "    img = rgb\n",
        "\n",
        "    # 4. 8bit([0, 255]) 데이터를 float32([0, 1]) 범위로 변환\n",
        "    img = tf.image.convert_image_dtype(img, tf.float32)\n",
        "    \n",
        "    # 5. 인식률 향상을 위해 이미지를 이진화(Binarization) 함\n",
        "    img = tf.where(img > 0, 1, 0)\n",
        "    \n",
        "    # 6. 이미지 크기에 맞게 리사이징\n",
        "    img = tf.image.resize(img, [img_height, img_width])\n",
        "    \n",
        "    # 7. 이미지 가로세로 바꿈 -> 이미지의 가로와 시간 차원을 대응하기 위함\n",
        "    img = tf.transpose(img, perm=[1, 0, 2])\n",
        "    \n",
        "    # 8. label의 문자들을 숫자로 매핑\n",
        "    label = char_to_num(tf.strings.unicode_split(label, input_encoding=\"UTF-8\"))\n",
        "    \n",
        "    # 9. model의 형식에 맞게 데이터 반환\n",
        "    return {\"image\": img, \"label\": label}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uRvdCY0HfCbD"
      },
      "source": [
        "## `Dataset` 객체 생성"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PeOdIvctfCbE"
      },
      "source": [
        "train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))\n",
        "train_dataset = (\n",
        "    train_dataset.map(\n",
        "        encode_single_sample, num_parallel_calls=tf.data.AUTOTUNE\n",
        "    )\n",
        "    .batch(batch_size)\n",
        "    .prefetch(buffer_size=tf.data.AUTOTUNE)\n",
        ")\n",
        "\n",
        "validation_dataset = tf.data.Dataset.from_tensor_slices((x_valid, y_valid))\n",
        "validation_dataset = (\n",
        "    validation_dataset.map(\n",
        "        encode_single_sample, num_parallel_calls=tf.data.AUTOTUNE\n",
        "    )\n",
        "    .batch(batch_size)\n",
        "    .prefetch(buffer_size=tf.data.AUTOTUNE)\n",
        ")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XoI9nDaTfCbE"
      },
      "source": [
        "## 샘플 데이터 시각화"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Cf1TGe2XfCbF"
      },
      "source": [
        "_, ax = plt.subplots(4, 4, figsize=(10, 5))\n",
        "for batch in train_dataset.take(1):\n",
        "    images = batch[\"image\"]\n",
        "    labels = batch[\"label\"]\n",
        "    for i in range(16):\n",
        "        img = (images[i] * 255).numpy().astype(\"uint8\")\n",
        "        label = tf.strings.reduce_join(num_to_char(labels[i])).numpy().decode(\"utf-8\")\n",
        "        ax[i // 4, i % 4].imshow(img[:, :, 0].T, cmap=\"gray\")\n",
        "        ax[i // 4, i % 4].set_title(label)\n",
        "        ax[i // 4, i % 4].axis(\"off\")\n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OuL0nqmLfCbF"
      },
      "source": [
        "## 모델 정의"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OesFffZ6fCbG"
      },
      "source": [
        "class CTCLayer(layers.Layer):\n",
        "    def __init__(self, name=None):\n",
        "        super().__init__(name=name)\n",
        "        self.loss_fn = keras.backend.ctc_batch_cost\n",
        " \n",
        "    def call(self, y_true, y_pred):\n",
        "        # 모델이 training하는 경우, self.add_loss()를 사용하여 loss를 계산하고 더해줌\n",
        "        batch_len = tf.cast(tf.shape(y_true)[0], dtype=\"int64\")\n",
        "        input_length = tf.cast(tf.shape(y_pred)[1], dtype=\"int64\")\n",
        "        label_length = tf.cast(tf.shape(y_true)[1], dtype=\"int64\")\n",
        " \n",
        "        input_length = input_length * tf.ones(shape=(batch_len, 1), dtype=\"int64\")\n",
        "        label_length = label_length * tf.ones(shape=(batch_len, 1), dtype=\"int64\")\n",
        " \n",
        "        loss = self.loss_fn(y_true, y_pred, input_length, label_length)\n",
        "        self.add_loss(loss)\n",
        " \n",
        "        # 테스트 시에는 예측 결과값만 반환\n",
        "        return y_pred\n",
        " \n",
        " \n",
        "def build_model():\n",
        "    # model Input 정의\n",
        "    input_img = layers.Input(\n",
        "        shape=(img_width, img_height, 1), name=\"image\", dtype=\"float32\"\n",
        "    )\n",
        "    labels = layers.Input(name=\"label\", shape=(None,), dtype=\"float32\")\n",
        " \n",
        "    # 첫 번째 conv block\n",
        "    x = layers.Conv2D(\n",
        "        64,\n",
        "        (3, 3),\n",
        "        activation=\"relu\",\n",
        "        kernel_initializer=\"he_normal\",\n",
        "        padding=\"same\",\n",
        "        name=\"Conv1\",\n",
        "    )(input_img)\n",
        "    \n",
        "    x = layers.MaxPooling2D((2, 2), name=\"pool1\")(x)\n",
        "    \n",
        " \n",
        "    # 두 번째 conv block\n",
        "    x = layers.Conv2D(\n",
        "        64,\n",
        "        (3, 3),\n",
        "        activation=\"relu\",\n",
        "        kernel_initializer=\"he_normal\",\n",
        "        padding=\"same\",\n",
        "        name=\"Conv2\",\n",
        "    )(x)\n",
        "    \n",
        "    x = layers.MaxPooling2D((2, 2), name=\"pool2\")(x)\n",
        "\n",
        "    # 두 번의 max pool(stride 2, pool size 2)을 사용하므로 feature maps는 4배 downsample 됨\n",
        "    # 마지막 레이어의 필터의 갯수는 64개\n",
        "    # 모델의 RNN 파트에 넣기 전에 Reshape를 해 줌\n",
        "    new_shape = ((img_width // 4), (img_height // 4) * 64)\n",
        "    x = layers.Reshape(target_shape=new_shape, name=\"reshape\")(x)\n",
        "    x = layers.Dense(64, activation=\"relu\", name=\"dense3\")(x)\n",
        " \n",
        "    # RNNs\n",
        "    x = layers.Bidirectional(layers.LSTM(32, return_sequences=True))(x)\n",
        " \n",
        "    # Output layer\n",
        "    x = layers.Dense(\n",
        "        len(char_to_num.get_vocabulary()) + 1, activation=\"softmax\", name=\"dense2\"\n",
        "    )(x)\n",
        " \n",
        "    # 모델에 CTC loss를 계산하는 CTC Layer 추가\n",
        "    output = CTCLayer(name=\"ctc_loss\")(labels, x)\n",
        " \n",
        "    # 모델 정의\n",
        "    model = keras.models.Model(\n",
        "        inputs=[input_img, labels], outputs=output, name=\"ocr_model_v1\"\n",
        "    )\n",
        "    # Optimizer 정의\n",
        "    opt = keras.optimizers.Adam()\n",
        "    # 모델 컴파일 후 반환\n",
        "    model.compile(optimizer=opt)\n",
        "    return model\n",
        " \n",
        " \n",
        "# 모델 구함\n",
        "model = build_model()\n",
        "model.summary()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GDm52Am7fCbG"
      },
      "source": [
        "## 모델 학습"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YvYMHLyufCbG"
      },
      "source": [
        "epochs = 500\n",
        "early_stopping_patience = 10\n",
        "# Early stopping 콜백 함수 선언\n",
        "early_stopping = keras.callbacks.EarlyStopping(\n",
        "    monitor=\"val_loss\", patience=early_stopping_patience, restore_best_weights=True\n",
        ")\n",
        " \n",
        "# 모델 학습\n",
        "history = model.fit(\n",
        "    train_dataset,\n",
        "    validation_data=validation_dataset,\n",
        "    epochs=epochs,\n",
        "    callbacks=[early_stopping],\n",
        ")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pzH33a7a22MC"
      },
      "source": [
        "## 모델 평가"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6NQzMKuy26Y1"
      },
      "source": [
        "results = model.evaluate(validation_dataset, batch_size=batch_size)\n",
        "print(\"Test loss:\", results)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Qo2O5PiUfCbG"
      },
      "source": [
        "## 이미지에서 텍스트 추론"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "p8XacpVwfCbH"
      },
      "source": [
        "# 출력 레이어까지 레이어를 추출하여 예측 모델을 가져 옴\n",
        "prediction_model = keras.models.Model(\n",
        "    model.get_layer(name=\"image\").input, model.get_layer(name=\"dense2\").output\n",
        ")\n",
        "prediction_model.summary()\n",
        "\n",
        "# 추론 결과 후처리\n",
        "def decode_batch_predictions(pred):\n",
        "    input_len = np.ones(pred.shape[0]) * pred.shape[1]\n",
        "    # Greedy 알고리즘 사용. 복잡한 작업은 Beam Search도 사용 가능\n",
        "    results = keras.backend.ctc_decode(pred, input_length=input_len, greedy=True)[0][0][\n",
        "        :, :max_length\n",
        "    ]\n",
        "    # 매핑된 데이터 복구\n",
        "    output_text = []\n",
        "    for res in results:\n",
        "        res = tf.strings.reduce_join(num_to_char(res)).numpy().decode(\"utf-8\")\n",
        "        output_text.append(res)\n",
        "    return output_text\n",
        "\n",
        "\n",
        "# 랜덤 샘플 추출 후, 해당 이미지를 추론하여 시각화\n",
        "for batch in validation_dataset.take(1):\n",
        "    batch_images = batch[\"image\"]\n",
        "    batch_labels = batch[\"label\"]\n",
        " \n",
        "    preds = prediction_model.predict(batch_images)\n",
        "    pred_texts = decode_batch_predictions(preds)\n",
        " \n",
        "    orig_texts = []\n",
        "    for label in batch_labels:\n",
        "        label = tf.strings.reduce_join(num_to_char(label)).numpy().decode(\"utf-8\")\n",
        "        orig_texts.append(label)\n",
        "\n",
        "    _, ax = plt.subplots(4, 4, figsize=(15, 5))\n",
        "    for i in range(len(pred_texts)):\n",
        "        img = (batch_images[i, :, :, 0] * 255).numpy().astype(np.uint8)\n",
        "        img = img.T\n",
        "        title = f\"Predict: {pred_texts[i]}\"\n",
        "        ax[i // 4, i % 4].imshow(img, cmap=\"gray\")\n",
        "        ax[i // 4, i % 4].set_title(title)\n",
        "        ax[i // 4, i % 4].axis(\"off\")\n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "H9BeDYG1MYCX"
      },
      "source": [
        "## 모델 파일 저장"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lf0bRtl7NHMF"
      },
      "source": [
        "save_path = \"/content/drive/MyDrive/Colab Notebooks/models\"\n",
        "prediction_model.save(f\"{save_path}/data.h5\")\n",
        "prediction_model.save(\"/content/model/data.h5\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FFfn4us9NiP7"
      },
      "source": [
        "## Keras `.h5` 모델을  `.tflite` 모델로 변환\n",
        "\n",
        "TensorFlow Keras 모델을 TensorFlow Lite에서도 사용할 수 있도록 변환합니다."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kOHk6eFaMcWq"
      },
      "source": [
        "saved_path = \"/content/drive/MyDrive/Colab Notebooks/models\"\n",
        "h5_model_path = f\"{saved_path}/data.h5\"\n",
        " \n",
        "# H5 모델 변환\n",
        "h5_model = tf.keras.models.load_model(h5_model_path)\n",
        "converter = tf.lite.TFLiteConverter.from_keras_model(h5_model)\n",
        "# 220810 추가) tflite 플래그 설정 안 하면 LSTM 변환 시 오류 발생\n",
        "converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS, tf.lite.OpsSet.SELECT_TF_OPS] \n",
        "converter._experimental_lower_tensor_list_ops = False\n",
        "\n",
        "\n",
        "tflite_model = converter.convert()\n",
        " \n",
        "# TFLite 모델 저장\n",
        "with open(f\"{saved_path}/data.tflite\", 'wb') as f:\n",
        "    f.write(tflite_model)\n",
        "with open(f\"/content/model/data.tflite\", 'wb') as f:\n",
        "    f.write(tflite_model)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "JB51X6cDNGlr"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}