CRAFT_demo.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "GitHub  \n",
        "https://github.com/clovaai/CRAFT-pytorch  \n",
        "論文  \n",
        "https://arxiv.org/abs/1904.01941    \n",
        "  \n",
        "<a href=\"https://colab.research.google.com/github/kaz12tech/ai_demos/blob/master/CRAFT_demo.ipynb\" target=\"_blank\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4TP5UEjJuRjL"
      },
      "source": [
        "## GitHubからCRAFTのソースコードを取得"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "JEUMCU9tpZCw",
        "outputId": "2071fb47-624e-444d-d3a0-91d2f34ada30"
      },
      "outputs": [],
      "source": [
        "%cd /content/\n",
        "!git clone https://github.com/clovaai/CRAFT-pytorch.git"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hRwgjpBkurB5"
      },
      "source": [
        "## ライブラリのインストール"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "KE1PSdNdpqAc",
        "outputId": "d203fc5b-e3e3-42eb-f344-5f5ef426410d"
      },
      "outputs": [],
      "source": [
        "%cd /content/CRAFT-pytorch/\n",
        "!pip uninstall torch torchtext torchaudio -y\n",
        "!pip install -r requirements.txt"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vopQQhndvxrb"
      },
      "source": [
        "## ライブラリのインポート"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "T-cN5xI4vz_q"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import gdown\n",
        "import glob\n",
        "from PIL import Image\n",
        "import matplotlib.pyplot as plt\n",
        "import random\n",
        "import shutil\n",
        "\n",
        "import torch"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "u1D8w4gcutPM"
      },
      "source": [
        "## 学習済みモデルのダウンロード"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "-8KqTN4hsgvK",
        "outputId": "f729305b-c70f-4f07-cd1a-e7aaa155351b"
      },
      "outputs": [],
      "source": [
        "%cd /content/CRAFT-pytorch/\n",
        "!mkdir pretrained\n",
        "%cd /content/CRAFT-pytorch/pretrained\n",
        "\n",
        "for google_drive_file_id in ['1Jk4eGD7crsqCCg9C9VjCLkMN3ze8kutZ']:\n",
        "  gdown.download(\n",
        "      url='https://drive.google.com/uc?id={}'.format(google_drive_file_id),\n",
        "      output=\"./craft_mlt_25k.pth\", \n",
        "      quiet=False\n",
        "      )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2m4eUC5S1_NK"
      },
      "source": [
        "## データセットのダウンロード"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "KD1Xz-SfvTo1",
        "outputId": "5d230647-a8f4-4ba4-d4ae-161ace396d27"
      },
      "outputs": [],
      "source": [
        "%cd /content/CRAFT-pytorch/\n",
        "!mkdir dataset\n",
        "%cd /content/CRAFT-pytorch/dataset\n",
        "\n",
        "for google_drive_file_id in ['1sptDnAomQHFVZbjvnWt2uBvyeJ-gEl-A']:\n",
        "  gdown.download(\n",
        "      url='https://drive.google.com/uc?id={}'.format(google_drive_file_id),\n",
        "      output=\"./icdar2013.zip\", \n",
        "      quiet=False\n",
        "      )\n",
        "\n",
        "# ダウンロードしたzipファイルの解凍\n",
        "# 解凍したファイルすべてが出力されることを避けるため > /dev/null で出力を捨てる\n",
        "!unzip ./icdar2013.zip > /dev/null\n",
        "\n",
        "# icdar2013 testファイルリストを取得\n",
        "test_path = \"/content/CRAFT-pytorch/dataset/icdar2013/test_images/\"\n",
        "test_path_list = glob.glob(test_path + \"*.jpg\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iKViCIIm2BVs"
      },
      "source": [
        "## 文字検出(Character Detection)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "id": "qDYlc-qB2Toz",
        "outputId": "14be8ca8-2e58-4106-9ecf-83fecb641c67"
      },
      "outputs": [],
      "source": [
        "# 文字検出対象のファイル選定\n",
        "target_path = \"/content/CRAFT-pytorch/inputs/\"\n",
        "if os.path.exists(target_path):\n",
        "  shutil.rmtree(target_path)\n",
        "  os.makedirs(target_path, exist_ok=True)\n",
        "else:\n",
        "  os.makedirs(target_path, exist_ok=True)\n",
        "\n",
        "# 文字検出用にランダムに5枚選択\n",
        "target = random.sample(test_path_list, 5)\n",
        "\n",
        "for file_path in target:\n",
        "  shutil.copy(file_path, target_path)\n",
        "\n",
        "target_path_list = glob.glob(target_path + \"*.jpg\")\n",
        "\n",
        "# 文字検出対象画像の表示\n",
        "for t_path in target_path_list:\n",
        "  image = Image.open(t_path).convert(\"RGB\")\n",
        "  plt.figure(figsize=(6, 6))\n",
        "  plt.imshow(image)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "dZY-mjjO2rnD",
        "outputId": "725fbb91-bd75-4863-e624-aea15301a800"
      },
      "outputs": [],
      "source": [
        "%cd /content/CRAFT-pytorch/\n",
        "\n",
        "# GPUが使用可能であればTrue、不可であればFalse\n",
        "device = torch.cuda.is_available()\n",
        "test_img = test_path_list[0]\n",
        "\n",
        "!python3 test.py \\\n",
        "  --trained_model=/content/CRAFT-pytorch/pretrained/craft_mlt_25k.pth \\\n",
        "  --test_folder=$target_path \\\n",
        "  --cuda=$device"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UFlkF9riEgkF"
      },
      "source": [
        "## 文字検出結果の表示"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "id": "xQF0ESav4xzr",
        "outputId": "61ae4c2d-8065-4ee0-ec6d-4bd9b9e8b578"
      },
      "outputs": [],
      "source": [
        "results = glob.glob(\"/content/CRAFT-pytorch/result/*.jpg\")\n",
        "results = [f for f in results if not 'mask' in f ]\n",
        "\n",
        "# 文字検出対象画像の表示\n",
        "for result in results:\n",
        "  image = Image.open(result).convert(\"RGB\")\n",
        "  plt.figure(figsize=(6, 6))\n",
        "  plt.imshow(image)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "CRAFT-demo.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}