IS[Mask R-CNN]-3 (Manga109)_FTA_TOD.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OnxJIeHIq5Nc"
      },
      "outputs": [],
      "source": [
        "from google.colab import drive\n",
        "drive.mount(\"/content/drive\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TKJV3yaVMRhe"
      },
      "outputs": [],
      "source": [
        "!pip3 uninstall keras-nightly\n",
        "!pip3 uninstall -y tensorflow\n",
        "!pip3 install keras==2.1.6\n",
        "!pip3 install tensorflow==1.15.0\n",
        "!pip3 install h5py==2.10.0"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3oztxkItMb67"
      },
      "outputs": [],
      "source": [
        "!git clone https://github.com/matterport/Mask_RCNN.git"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Wa61M3fOMhdS"
      },
      "outputs": [],
      "source": [
        "%cd Mask_RCNN"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sNR9tCIWMmNf"
      },
      "outputs": [],
      "source": [
        "!python setup.py install\n",
        "!pip install -r requirements.txt"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GVw_z1VjO1nb"
      },
      "outputs": [],
      "source": [
        "!pip install elementpath"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Fq7ge8sqO2kB"
      },
      "outputs": [],
      "source": [
        "import elementpath\n",
        "from google.colab import files\n",
        "from os import listdir\n",
        "from numpy import zeros, asarray, expand_dims, mean\n",
        "from numpy import asarray\n",
        "from mrcnn import visualize\n",
        "from mrcnn.utils import Dataset, extract_bboxes, compute_ap\n",
        "from mrcnn.config import Config\n",
        "from mrcnn.visualize import display_instances\n",
        "from mrcnn.model import MaskRCNN, load_image_gt, mold_image\n",
        "import matplotlib.pyplot as pyplot\n",
        "from matplotlib.patches import Rectangle, Arrow\n",
        "import math\n",
        "import cv2"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Jb0h4NPD2wfj"
      },
      "outputs": [],
      "source": [
        "%cd /content/"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Instance Segmentation (IS)"
      ],
      "metadata": {
        "id": "U7pxjbPqhl2F"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "id": "e2nupC0a1pHr"
      },
      "outputs": [],
      "source": [
        "class PredictionConfig(Config):\n",
        "  NAME = \"manga_cfg\"\n",
        "  NUM_CLASSES = 1 + 3\n",
        "  GPU_COUNT = 1\n",
        "  IMAGES_PER_GPU = 1"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5kqgB_G12V-Z"
      },
      "outputs": [],
      "source": [
        "cfg = PredictionConfig()\n",
        "model = MaskRCNN(mode=\"inference\", model_dir=\"/content\", config=cfg)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QA9k3U792bD6"
      },
      "outputs": [],
      "source": [
        "model.load_weights(\"/content/drive/MyDrive/MLA/IRC-SET/IS[Mask_R_CNN]-3 (Manga109).h5\", by_name=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "id": "hh8JcJll35Uc"
      },
      "outputs": [],
      "source": [
        "class_names = ['BG', 'face', 'text', 'frame']"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YYaSwPfd29OX"
      },
      "outputs": [],
      "source": [
        "image = cv2.imread('/content/drive/MyDrive/MLA/IRC-SET/sample.png')\n",
        "results = model.detect([image], verbose=1)\n",
        "r = results[0]\n",
        "visualize.display_instances(image, r['rois'], r['masks'], r['class_ids'], class_names, r['scores'])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ji8U486_kHF7"
      },
      "source": [
        "# Face-Text Association (FTA)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "id": "uNhxoCJ3wSvP"
      },
      "outputs": [],
      "source": [
        "def arrow_face_text(image, cfg, model):\n",
        "  results = model.detect([image], verbose=1)\n",
        "  r = results[0]\n",
        "\n",
        "  pyplot.figure(figsize=(10,10))\n",
        "  pyplot.imshow(image)\n",
        "  pyplot.title(\"Face to Text\")\n",
        "\n",
        "  ax = pyplot.gca()\n",
        "\n",
        "  face_centers = []\n",
        "  text_centers = []\n",
        "  frame_corners = []\n",
        "  count = 0\n",
        "\n",
        "  for id in list(r['class_ids']):\n",
        "    box = list(r['rois'])[count]\n",
        "    count += 1\n",
        "\n",
        "    y1, x1, y2, x2 = box\n",
        "    width, height = x2 - x1, y2 - y1\n",
        "\n",
        "    if id == 1:\n",
        "      face_x = (x1 + x2)//2\n",
        "      face_y = (y1 + y2)//2\n",
        "      face_center = [face_x, face_y]\n",
        "      face_centers.append(face_center)\n",
        "\n",
        "      rect = Rectangle((x1, y1), width, height, fill=False, color=\"red\")\n",
        "      ax.add_patch(rect)\n",
        "\n",
        "    elif id == 2:\n",
        "      text_x = (x1 + x2)//2\n",
        "      text_y = (y1 + y2)//2\n",
        "      text_center = [text_x, text_y]\n",
        "      text_centers.append(text_center)\n",
        "\n",
        "      rect = Rectangle((x1, y1), width, height, fill=False, color=\"yellow\")\n",
        "      ax.add_patch(rect)\n",
        "\n",
        "    elif id == 3:\n",
        "      frame_corners.append([x1, x2, y1, y2])\n",
        "      rect = Rectangle((x1, y1), width, height, fill=False, color=\"violet\")\n",
        "      ax.add_patch(rect)\n",
        "\n",
        "  faces_to_texts = []\n",
        "\n",
        "  for frame in frame_corners:\n",
        "    x1, x2, y1, y2 = frame\n",
        "    face_centers_filtered, text_centers_filtered = [], []\n",
        "    num_faces, num_text = 0, 0\n",
        "\n",
        "    for face in face_centers:\n",
        "      if face[0] < x2 and face[0] > x1 and face[1] < y2 and face[1] > y1:\n",
        "        num_faces += 1\n",
        "        face_centers_filtered.append(face)\n",
        "\n",
        "    for text in text_centers:\n",
        "      if text[0] < x2 and text[0] > x1 and text[1] < y2 and text[1] > y1:\n",
        "        num_text += 1\n",
        "        text_centers_filtered.append(text)\n",
        "    \n",
        "    if num_faces >= num_text:\n",
        "      for face in face_centers_filtered:\n",
        "        if text_centers_filtered:\n",
        "          nearest_text = text_centers_filtered[0]\n",
        "          shortest_x = abs(face[0] - nearest_text[0])\n",
        "          shortest_y = abs(face[1] - nearest_text[1])\n",
        "          shortest_distance = math.sqrt(shortest_x**2 + shortest_y**2)\n",
        "\n",
        "          for text in text_centers_filtered:\n",
        "            distance_x = abs(face[0] - text[0])\n",
        "            distance_y = abs(face[1] - text[1])\n",
        "            distance = math.sqrt(distance_x**2 + distance_y**2)\n",
        "\n",
        "            if distance < shortest_distance:\n",
        "              shortest_distance = distance\n",
        "              nearest_text = text\n",
        "\n",
        "          face_to_text = [face, nearest_text]\n",
        "          faces_to_texts.append(face_to_text)\n",
        "\n",
        "    elif num_faces < num_text:\n",
        "      for text in text_centers_filtered:\n",
        "        if face_centers_filtered:\n",
        "          nearest_face = face_centers_filtered[0]\n",
        "          shortest_x = abs(text[0] - nearest_face[0])\n",
        "          shortest_y = abs(text[1] - nearest_face[1])\n",
        "          shortest_distance = math.sqrt(shortest_x**2 + shortest_y**2)\n",
        "\n",
        "          for face in face_centers_filtered:\n",
        "            distance_x = abs(face[0] - text[0])\n",
        "            distance_y = abs(face[1] - text[1])\n",
        "            distance = math.sqrt(distance_x**2 + distance_y**2)\n",
        "\n",
        "            if distance < shortest_distance:\n",
        "              shortest_distance = distance\n",
        "              nearest_face = face\n",
        "\n",
        "          face_to_text = [nearest_face, text]\n",
        "          faces_to_texts.append(face_to_text)\n",
        "\n",
        "  for face_to_text in faces_to_texts:\n",
        "    face, text = face_to_text\n",
        "\n",
        "    face_x, face_y = face\n",
        "    text_x, text_y = text\n",
        "\n",
        "    length_x = abs(face_x - text_x)\n",
        "    length_y = abs(face_y - text_y)\n",
        "\n",
        "    if face_x > text_x: #face is to the right of text\n",
        "      length_x *= -1\n",
        "      \n",
        "    if face_y > text_y: #face is below text\n",
        "      length_y *= -1\n",
        "\n",
        "    arrow = Arrow(face_x, face_y, length_x, length_y, color=\"cornflowerblue\")\n",
        "    ax.add_patch(arrow)\n",
        "\n",
        "  pyplot.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "I5dE1oLoKHP0"
      },
      "outputs": [],
      "source": [
        "image = cv2.imread('/content/drive/MyDrive/MLA/IRC-SET/sample.png')\n",
        "width = image.shape[1]\n",
        "arrow_face_text(image, cfg, model)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "J5PuyVVQrFdF"
      },
      "source": [
        "# Text Order Determination (TOD)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {
        "id": "iGayiWpmmKtg"
      },
      "outputs": [],
      "source": [
        "def order_frame_halves(unordered_frames):\n",
        "  ordered_frames = []\n",
        "  while unordered_frames:\n",
        "    for frame in unordered_frames:\n",
        "      for other_frame in unordered_frames:\n",
        "        if other_frame[3] < frame[1]: # if there is a frame above the current frame\n",
        "          break # the current frame is not the next frame to read\n",
        "      else:\n",
        "        ordered_frames.append(frame[:-1]) # the current frame is the next frame\n",
        "        unordered_frames.remove(frame)\n",
        "\n",
        "  return ordered_frames"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {
        "id": "foi7SnFaRScx"
      },
      "outputs": [],
      "source": [
        "def order_text(image, img_width, cfg, model):\n",
        "  results = model.detect([image], verbose=1)\n",
        "  r = results[0]\n",
        "\n",
        "  pyplot.figure(figsize=(10,10))\n",
        "  pyplot.imshow(image)\n",
        "  pyplot.title(\"Order Text\")\n",
        "\n",
        "  ax = pyplot.gca()\n",
        "\n",
        "  text_centers = []\n",
        "  total_unordered_frames = []\n",
        "  count = 0\n",
        "\n",
        "  for id in list(r['class_ids']):\n",
        "    box = list(r['rois'])[count]\n",
        "    count += 1\n",
        "\n",
        "    y1, x1, y2, x2 = box\n",
        "    width, height = x2 - x1, y2 - y1\n",
        "\n",
        "    if id == 2:\n",
        "      text_x = (x1 + x2)//2\n",
        "      text_y = (y1 + y2)//2\n",
        "      text_center = [text_x, text_y]\n",
        "      text_centers.append(text_center)\n",
        "\n",
        "      rect = Rectangle((x1, y1), width, height, fill=False, color=\"cornflowerblue\")\n",
        "      ax.add_patch(rect)\n",
        "    elif id == 3:\n",
        "      frame_x_center = (x1 + x2)//2\n",
        "      total_unordered_frames.append([x1, y1, x2, y2, frame_x_center]) # the corners are ordered this way to improve sorting later on\n",
        "\n",
        "  total_unordered_frames.sort(reverse=True)\n",
        "  ordered_frames = [] # 1st, 2nd, 3rd, ...\n",
        "\n",
        "  half_line = img_width//2\n",
        "\n",
        "  unordered_frames_first = [frame for frame in total_unordered_frames if frame[-1] > half_line] # frames to the right of half line\n",
        "  unordered_frames_second = [frame for frame in total_unordered_frames if frame[-1] <= half_line] # frames to the left of half line\n",
        "\n",
        "  ordered_frames_first = order_frame_halves(unordered_frames_first)\n",
        "  ordered_frames_second = order_frame_halves(unordered_frames_second)\n",
        "\n",
        "  index = 1\n",
        "  \n",
        "  for frame in ordered_frames_first:\n",
        "    x1, y1, x2, y2 = frame\n",
        "    text_centers_filtered = []\n",
        "\n",
        "    for text in text_centers:\n",
        "      if text[0] < x2 and text[0] > x1 and text[1] < y2 and text[1] > y1:\n",
        "        text_centers_filtered.append([text[0], -text[1]])\n",
        "    \n",
        "    text_centers_filtered.sort(reverse=True)\n",
        "    for text in text_centers_filtered:\n",
        "      pyplot.text(text[0], -text[1], index, color=\"red\")\n",
        "      index += 1\n",
        "\n",
        "  for frame in ordered_frames_second:\n",
        "    x1, y1, x2, y2 = frame\n",
        "    text_centers_filtered = []\n",
        "\n",
        "    for text in text_centers:\n",
        "      if text[0] < x2 and text[0] > x1 and text[1] < y2 and text[1] > y1:\n",
        "        text_centers_filtered.append([text[0], -text[1]])\n",
        "    \n",
        "    text_centers_filtered.sort(reverse=True)\n",
        "    for text in text_centers_filtered:\n",
        "      pyplot.text(text[0], -text[1], index, color=\"red\")\n",
        "      index += 1\n",
        "\n",
        "  pyplot.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "89CXAA7tmiMa"
      },
      "outputs": [],
      "source": [
        "image = cv2.imread('/content/drive/MyDrive/MLA/IRC-SET/sample.png')\n",
        "width = image.shape[1]\n",
        "order_text(image, width, cfg, model)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Note: Optical Character Recognition (M-OCR) is not included in this notebook!"
      ],
      "metadata": {
        "id": "9iaG9KhLh2Ik"
      }
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "MLA.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}