bulk_inferencing/BN_DRISHTI_Bulk_Inferencing.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/shaoncsecu/BN-DRISHTI/blob/main/bulk_inferencing/BN_DRISHTI_Bulk_Inferencing.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FXQXRqjgMsD7"
      },
      "source": [
        "#**Defining functions, importing Libraries, deleting directory contents**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WVu0zUTpMrvE"
      },
      "source": [
        "##Importing necessary libraries..."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2jIXND5rkQds"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "from google.colab import files\n",
        "from os.path import join\n",
        "import numpy as np\n",
        "import operator\n",
        "import cv2\n",
        "import fnmatch\n",
        "import matplotlib.pyplot as plt\n",
        "import PIL\n",
        "from PIL import Image\n",
        "import glob\n",
        "import imutils\n",
        "from math import *\n",
        "from scipy.stats import mode\n",
        "from distutils.dir_util import copy_tree\n",
        "from IPython.display import Image\n",
        "from google.colab.patches import cv2_imshow\n",
        "import PIL.Image\n",
        "import os.path\n",
        "import shutil"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Oeyhl7hVNtDM"
      },
      "source": [
        "##Removing previous files and folders from their directories in New run..."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "N__WOJ1NJ4Z6"
      },
      "outputs": [],
      "source": [
        "# Removing previous files and folders from their directories....\n",
        "x = os.listdir(\"/content\")\n",
        "for i in x:\n",
        "  if i.endswith(('.png', '.jpg', '.jpeg', '.tiff', '.bmp', '.gif', 'JPG', 'JPEG', 'PNG')) == True:\n",
        "    img_label = i.split('.')[0]\n",
        "    final_path = \"/content/\"+img_label\n",
        "    shutil.rmtree(final_path)\n",
        "    temp = \"/content/\"+i\n",
        "    os.remove(temp)\n",
        "\n",
        "path1 = \"/content/DSkew\"\n",
        "if os.path.exists(path1) == True:\n",
        "  shutil.rmtree(path1)\n",
        "path2 = \"/content/HaughLine_Affine\"\n",
        "if os.path.exists(path2) == True:\n",
        "  shutil.rmtree(path2)\n",
        "path3 = \"/content/Rotated_line_by_HaughLine_Affine\"\n",
        "if os.path.exists(path3) == True:\n",
        "  shutil.rmtree(path3)\n",
        "path4 = \"/content/final_line_segmentation\"\n",
        "if os.path.exists(path4) == True:\n",
        "  shutil.rmtree(path4)\n",
        "path5 = \"/content/final_word_segmentation\"\n",
        "if os.path.exists(path5) == True:\n",
        "  shutil.rmtree(path5)\n",
        "path6 = \"/content/initial_line_segmantation\"\n",
        "if os.path.exists(path6) == True:\n",
        "  shutil.rmtree(path6)\n",
        "path7 = \"/content/sorted_Word_detection\"\n",
        "if os.path.exists(path7) == True:\n",
        "  shutil.rmtree(path7)\n",
        "path8 = \"/content/sorted_line_after_1st_detection\"\n",
        "if os.path.exists(path8) == True:\n",
        "  shutil.rmtree(path8)\n",
        "path9 = \"/content/yolov5/runs\"\n",
        "if os.path.exists(path9) == True:\n",
        "  shutil.rmtree(path9)\n",
        "path11 = \"/content/2nd line detection for rotated images\"\n",
        "if os.path.exists(path11) == True:\n",
        "  shutil.rmtree(path11)\n",
        "path12 = \"/content/2nd line detection for rotated images (labels)\"\n",
        "if os.path.exists(path12) == True:\n",
        "  shutil.rmtree(path12)\n",
        "path13 = \"/content/Original line images\"\n",
        "if os.path.exists(path13) == True:\n",
        "  shutil.rmtree(path13)\n",
        "path14 = \"/content/model\"\n",
        "if os.path.exists(path14) != True:\n",
        "  '''Downloading the MODELS into CoLab temporary directory'''\n",
        "  !mkdir /content/model\n",
        "  !wget -P /content/model/ 'https://huggingface.co/crusnic/BN-DRISHTI/resolve/main/models/line_model_best.pt'\n",
        "  !wget -P /content/model/ 'https://huggingface.co/crusnic/BN-DRISHTI/resolve/main/models/word_model_best.pt'\n",
        "  "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "T2R2khd7Ucdl"
      },
      "source": [
        "##YoloV5 Setup..."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4Eed9mPSUbM1"
      },
      "outputs": [],
      "source": [
        "path = \"/content/yolov5\"\n",
        "if os.path.exists(path) == False:\n",
        "  !git clone https://github.com/ultralytics/yolov5  # clone repo\n",
        "  %cd yolov5\n",
        "  %pip install -qr requirements.txt  # install dependencies\n",
        "\n",
        "  import torch\n",
        "  from IPython.display import Image, clear_output  # to display images\n",
        "\n",
        "  clear_output()\n",
        "  print('Setup complete. Using torch %s %s' % (torch.__version__, torch.cuda.get_device_properties(0) if torch.cuda.is_available() else 'CPU'))\n",
        "\n",
        "  %cd /content/\n",
        "else:\n",
        "  print(\"Yolo requirements are already exists!\")\n",
        "  "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nfYJOTY-wxnE"
      },
      "source": [
        "##Unzipping files..."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Yh18T2L28PPn"
      },
      "source": [
        "Defining a function to copy the dataset into colab local directory (for faster execution)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3P0Xu3HBwwW_"
      },
      "outputs": [],
      "source": [
        "def file_unzip(zip_file):\n",
        "  !unzip $zip_file -d \"/content\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wrIeF7wMi_TJ"
      },
      "source": [
        "##Sort line name of given list..."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Fz8JCWdX89VB"
      },
      "source": [
        "Defining a function to sort a list containing the name of images to visualize them serially. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Xd-wTz_lLzTb"
      },
      "outputs": [],
      "source": [
        "def line_sort(lines):\n",
        "    sort_lines = {}\n",
        "    for line in lines:\n",
        "        img_lb = line.split('.')[0]\n",
        "        lb = [int(i) for i in img_lb.split('_')]\n",
        "        new_lb = [ '0'+str(r) if r<10 else str(r) for r in lb]\n",
        "        if len(new_lb)==2:\n",
        "            items = int(new_lb[0]+new_lb[1])\n",
        "        if len(new_lb)==3:\n",
        "            items = int(new_lb[0]+new_lb[1]+new_lb[2])\n",
        "        if len(new_lb)==4:\n",
        "            items = int(new_lb[0]+new_lb[1]+new_lb[2]+new_lb[3])\n",
        "        sort_lines[items] = line\n",
        "    # print(sort_lines)\n",
        "    sort_lines = dict(sorted(sort_lines.items()))\n",
        "    new_lines = list(sort_lines.values())\n",
        "    return new_lines"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GaarKnu00UIU"
      },
      "source": [
        "##Showing All Transitions..."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iTseHXFru5O5"
      },
      "source": [
        "Defining a function to draw images by iteration of a given directory. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OwKJD-2N0Y2_"
      },
      "outputs": [],
      "source": [
        "def show_transitions(f_name):\n",
        "  filename = f_name\n",
        "  crop_line = os.listdir(filename)\n",
        "  crop_lines = line_sort(crop_line)\n",
        "  print(crop_lines)\n",
        "  j = 1\n",
        "  for i in crop_lines:\n",
        "    print(\"Line no:\",j)\n",
        "    print(\"Image name:: \",i)\n",
        "    im = filename + \"/\" + i\n",
        "    crop_img = cv2.imread(im)\n",
        "    \n",
        "    width = crop_img.shape[1]\n",
        "    height = crop_img.shape[0]\n",
        "    print(\"Width, height: \",width, height)\n",
        "\n",
        "    plot_fig(crop_img)\n",
        "    j = j + 1"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YFO_tls0uovo"
      },
      "source": [
        "Defining a function to plot figures."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NiP40j9jP5OB"
      },
      "outputs": [],
      "source": [
        "def plot_fig(img, size = 15):\n",
        "  plt.figure(figsize=(size, size))\n",
        "  plt.imshow(imutils.opencv2matplotlib(img))\n",
        "  plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UTRaB4mDvVG1"
      },
      "source": [
        "Defining a function to visually compare results of annotations and predictions of each line image by plotting them one after another."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mdDDyPkFKheW"
      },
      "outputs": [],
      "source": [
        "def show_transitions_by_comparing(filename, gt_source_path, name1, name2):\n",
        "    crop_line = os.listdir(filename)\n",
        "    crop_lines = line_sort(crop_line)\n",
        "    print(\"Cropped Line images: \",crop_lines)\n",
        "    print()\n",
        "    j = 1\n",
        "    for i in crop_lines:\n",
        "      print(\"Line no:\",j)\n",
        "      print(\"Image name:: \",i)\n",
        "      im = filename + \"/\" + i\n",
        "      crop_img = cv2.imread(im)\n",
        "      \n",
        "      width = crop_img.shape[1]\n",
        "      height = crop_img.shape[0]\n",
        "      \n",
        "      # print(\"Annotated Line or GroundTruth ->\")\n",
        "      print(name1)\n",
        "      # img_gd = gt_line_img_dir + i\n",
        "      img_gd = gt_source_path + i\n",
        "      img_gd1 = cv2.imread(img_gd)\n",
        "      plot_fig(img_gd1)\n",
        "\n",
        "      # print(\"Predicted Line ->\")\n",
        "      print(name2)\n",
        "      plot_fig(crop_img)\n",
        "      print()\n",
        "      j = j + 1"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vAHj6AsXiUYS"
      },
      "source": [
        "##Drawing bounding box..."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e636IVvnvvW9"
      },
      "source": [
        "Defining a function to draw bounding boxes on images with their given labels."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fTprGMPjiXQa"
      },
      "outputs": [],
      "source": [
        "def draw_BB(img_path, label_path, flag):\n",
        "  img = cv2.imread(img_path)\n",
        "  dh, dw, _ = img.shape\n",
        "\n",
        "  lb = open(label_path, 'r')\n",
        "  data = lb.readlines()\n",
        "  lb.close()\n",
        "\n",
        "  for dt in data:\n",
        "    if flag == 0:\n",
        "      _, x, y, w, h = map(float, dt.split(' '))\n",
        "    else:\n",
        "      _, x, y, w, h, conf = map(float, dt.split(' '))\n",
        "    l = int((x - w / 2) * dw)\n",
        "    r = int((x + w / 2) * dw)\n",
        "    t = int((y - h / 2) * dh)\n",
        "    b = int((y + h / 2) * dh)\n",
        "    if flag == 0:\n",
        "      cv2.rectangle(img, (l, t), (r, b), (0, 250, 0), 2)\n",
        "    else:\n",
        "      cv2.rectangle(img, (l, t), (r, b), (0, 0, 250), 2)\n",
        "    \n",
        "  # plot_fig(img)\n",
        "  return img\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kT40IqFNk5x9"
      },
      "source": [
        "#**Detecting lines with YOLO** (Upload the given trained Line Model weights on your drive and assign its directory to the --weights command)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KX_Ttz_Ev9Oq"
      },
      "source": [
        "Defining a function where we pass necessary attributes to detect lines by YOLO."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-WRT-NnJkyKW"
      },
      "outputs": [],
      "source": [
        "# Yolo Detection...\n",
        "def yolo_detection(img_path, img_size, conf):\n",
        "  # %cd yolov5\n",
        "  # !python /content/yolov5/detect.py --weights /content/drive/MyDrive/thesis/yolov5/temp/YOLO_5x6_model/Line/best.pt --img $img_size --conf $conf --source $img_path --save-conf --save-txt\n",
        "  !python /content/yolov5/detect.py --weights /content/model/line_model_best.pt --img $img_size --conf $conf --source $img_path --save-conf --save-txt"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "irXlLcrro76F"
      },
      "source": [
        "##Sorting & Filtering lines & words ( labels ) after their 1st detection..."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "V4bFk8fz3-VA"
      },
      "source": [
        "Defining a function to sort and filter -\n",
        "\n",
        "* lines based on their y-axis's attribute and confidence, respectively. \n",
        "* words based on their x-axis's attribute and width, respectively.  \n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YYh7x7Rfo69t"
      },
      "outputs": [],
      "source": [
        "class Line_sort:\n",
        "    def __init__(self, txt_files, txt_loc, sort_label, flag):\n",
        "        self.txt_files = txt_files\n",
        "        self.txt_loc = txt_loc\n",
        "        self.sort_label = sort_label\n",
        "        self.flag = flag\n",
        "        self.read_file()\n",
        "\n",
        "    def read_file(self):\n",
        "        files = self.txt_files\n",
        "        # os.mkdir('/content/sorted_line_after_1st_detection')\n",
        "        os.mkdir(self.sort_label)\n",
        "        for file in files:\n",
        "            txt_file = []\n",
        "            file_loc = self.txt_loc+file\n",
        "\n",
        "            with open(file_loc, 'r' , encoding='utf-8',errors='ignore') as lines:\n",
        "                for line in lines:\n",
        "                    token = line.split()\n",
        "                    \n",
        "                    _, x, y, w, h, conf = map(float, line.split(' '))\n",
        "                    # print(\"width -> \",w)\n",
        "                    # print(\"confidence -> \",conf)\n",
        "                    if self.flag == 0: # 1st line detection lavel\n",
        "                      if w > 0.50 and conf < 0.50:\n",
        "                        continue\n",
        "                      else:\n",
        "                        txt_file.append(token)\n",
        "                    else: # Word detection lavel\n",
        "                      # if w > 0.50:\n",
        "                      #   continue\n",
        "                      # else:\n",
        "                        txt_file.append(token)\n",
        "\n",
        "            if self.flag == 0: # 1st line detection lavel\n",
        "               sorted_txt_file = sorted(txt_file, key=operator.itemgetter(2))\n",
        "            else: # Word detection lavel\n",
        "               sorted_txt_file = sorted(txt_file, key=operator.itemgetter(1))\n",
        "\n",
        "            # lenght = len(sorted_txt_file[0])\n",
        "            self.file_write(sorted_txt_file, file)\n",
        "\n",
        "    def file_write(self,txt_file, file_name):\n",
        "        # loc = '/content/sorted_line_after_1st_detection/'+file_name\n",
        "        loc = self.sort_label+file_name\n",
        "        with open(loc, 'w') as f: \n",
        "            c=0\n",
        "            for line in txt_file:  \n",
        "                for l in line:\n",
        "                    c+=1\n",
        "                    if c == len(line):\n",
        "                        f.write('%s' % l)\n",
        "                    else:\n",
        "                        f.write('%s ' % l)\n",
        "                f.write(\"\\n\")\n",
        "                c=0\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SN80A0_1pq2F"
      },
      "outputs": [],
      "source": [
        "def sort_detection_label(txt_loc, sort_label, flag):\n",
        "  txt_files = os.listdir(txt_loc)\n",
        "  obj = Line_sort(txt_files, txt_loc, sort_label, flag)\n",
        "  "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "w6YALdh8TKXK"
      },
      "source": [
        "##1st Line Segment from Sorted line ( labels ) after 1st detection..."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8A3aiY1B5osH"
      },
      "source": [
        "Defining a function to segment lines with YOLO's 1st line detection where we take full document image's width for a line if detected lines width ranges between 50% to 80%; otherwise, the lines are segmented as per detection."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4TfIxm5iTL75"
      },
      "outputs": [],
      "source": [
        "# text file and image sorting\n",
        "def images_and_txtfile_sort(images, txt_file):\n",
        "    txt = []\n",
        "    image = []\n",
        "    for item in txt_file:\n",
        "        ch = \"\"\n",
        "        for c in item:\n",
        "            if c == \"_\":\n",
        "                break\n",
        "            ch += c\n",
        "        txt.append(ch)\n",
        "\n",
        "    txt.sort(key=int)\n",
        "    sorted_txtfiles = []\n",
        "    sorted_images = []\n",
        "\n",
        "    for i in range(len(txt)):\n",
        "        for ele, ele2 in zip(txt_file, images):\n",
        "            st = \"\"\n",
        "            st2 = \"\"\n",
        "            for ch in ele:\n",
        "                if ch == \"_\":\n",
        "                    break\n",
        "                st += ch\n",
        "            if st == txt[i]:\n",
        "                sorted_txtfiles.append(ele)\n",
        "            for ch in ele2:\n",
        "                if ch == \"_\":\n",
        "                    break\n",
        "                st2 += ch\n",
        "            if st2 == txt[i]:\n",
        "                sorted_images.append(ele2)\n",
        "    return sorted_images, sorted_txtfiles\n",
        "\n",
        "# line segmantation by yolo bounding box\n",
        "def line_segmantation(path, image_loc, txt_loc, images, txt_file):\n",
        "\n",
        "    # image, txt_files = images_and_txtfile_sort(images, txt_file)\n",
        "    image = line_sort(images)\n",
        "    txt_files = line_sort(txt_file)\n",
        "    print(image)\n",
        "    print(txt_files)\n",
        "    \n",
        "    for image, txt in zip(image, txt_files):\n",
        "        \n",
        "        # make folder name according to txt file\n",
        "        new_folder = txt[:-4]\n",
        "        new_folder_loc = path + new_folder\n",
        "        # create subfolder in Lines according to txt file\n",
        "        os.mkdir(new_folder_loc)\n",
        "        # current image location\n",
        "        current_image = image_loc + image\n",
        "        img = cv2.imread(current_image)\n",
        "        dh, dw, _ = img.shape\n",
        "        # current txt file location\n",
        "        current_txt = txt_loc + txt\n",
        "        fl = open(current_txt, 'r')\n",
        "        data = fl.readlines()\n",
        "        fl.close()\n",
        "        k=1\n",
        "        for dt in data:\n",
        "            # _, x, y, w, h = map(float, dt.split(' '))\n",
        "            _, x, y, w, h,  conf = map(float, dt.split(' '))\n",
        "            if w > 0.50 and w < 0.80:\n",
        "               x = 0.5\n",
        "               w = 1.0\n",
        "            l = int((x - w / 2) * dw)\n",
        "            r = int((x + w / 2) * dw)\n",
        "            t = int((y - h / 2) * dh)\n",
        "            b = int((y + h / 2) * dh)\n",
        "\n",
        "            crop = img[t:b, l:r]\n",
        "            cv2.imwrite(\"{}{}/{}_{}.jpg\".format(path, new_folder, new_folder, k), crop)\n",
        "            k += 1\n",
        "\n",
        "def take_valid_img(images):\n",
        "    image = []\n",
        "    valid_img_ext = [\"jpg\", \"JPG\", \"jpeg\", \"JPEG\", \"png\", \"PNG\"]\n",
        "    for img in images:\n",
        "        try:\n",
        "            ext = img.split('.')[1]\n",
        "            if ext not in valid_img_ext :\n",
        "                continue\n",
        "            else:\n",
        "                image.append(img)\n",
        "        except:\n",
        "            continue\n",
        "    return image\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CCQMfBZeLVW8"
      },
      "outputs": [],
      "source": [
        "def line_segmentation(img_path, sorted_label):\n",
        "  current_directory = img_path\n",
        "  #print(current_directory)\n",
        "  current_directory1 = sorted_label\n",
        "  #print(current_directory1)\n",
        "  images = []  # take for all images\n",
        "  txt = []     # take for all txt files\n",
        "\n",
        "  # images\n",
        "  # take all files in list in current directory location\n",
        "  current_directory_files = os.listdir(current_directory) \n",
        "  #print(current_directory_files)\n",
        "  if not current_directory_files:\n",
        "      print(\"{} folder is empty!\")\n",
        "\n",
        "  # lines\n",
        "  # take all files in list in current directory1 location\n",
        "  current_directory_1_files = os.listdir(current_directory1)\n",
        "  if not current_directory_1_files:\n",
        "      print(\"Line folder is empty of {} !\")\n",
        "\n",
        "  txt_pattern = \"*.txt\"\n",
        "  # take all only image in images list from current directory files list\n",
        "  images1 = [entry for entry in current_directory_files]\n",
        "  images = take_valid_img(images1)\n",
        "\n",
        "  # take all only txt file in txt list from current directory1 files list\n",
        "  txt = [entry for entry in current_directory_1_files if fnmatch.fnmatch(entry, txt_pattern)]\n",
        "\n",
        "  os.mkdir('/content/croped_line_after_1st_detection')\n",
        "  path = \"/content/croped_line_after_1st_detection/\"\n",
        "  # calling the line segmantation function\n",
        "  line_segmantation(path, current_directory, current_directory1, images, txt)\n",
        "  print(\"Successful line segmantated in the {} folder\".format(path))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vRGq-eDyatWE"
      },
      "source": [
        "#**Rotation**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7pLXkMLsauVq"
      },
      "source": [
        "##DSkew by HaughLine..."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nJQBXMDI_xMY"
      },
      "source": [
        "Defining functions to do necessary preprocessing and calculating a line image's dimension skew (DSkew) using Probabilistic Hough Line Transform (PHT). And correcting it by using the Affine Transform if the DSkew angle detected by PHT is more than 0; otherwise, keep the image without DSkew by setting the DSkew angle as 0.  "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bXtNd9h6a3mm"
      },
      "outputs": [],
      "source": [
        "class ImgCorrect():\n",
        "    def __init__(self, img):\n",
        "        self.img = img\n",
        "        self.h, self.w, self.channel = self.img.shape\n",
        "        # print(\"Original images h & w -> | w: \",self.w, \"| h: \",self.h)\n",
        "        if self.w <= self.h:\n",
        "            self.scale = 700 / self.w\n",
        "            self.img = cv2.resize(self.img, (0, 0), fx=self.scale, fy=self.scale, interpolation=cv2.INTER_NEAREST)\n",
        "        else:\n",
        "            self.scale = 700 / self.h\n",
        "            self.img = cv2.resize(self.img, (0, 0), fx=self.scale, fy=self.scale, interpolation=cv2.INTER_NEAREST)\n",
        "        # print(\"Scaled image:\")\n",
        "        # plot_fig(self.img)\n",
        "        self.gray = cv2.cvtColor(self.img, cv2.COLOR_BGR2GRAY)\n",
        "\n",
        "    def img_lines(self):\n",
        "        # print(\"Gray Image:\")\n",
        "        # plot_fig(self.gray)\n",
        "        ret, binary = cv2.threshold(self.gray, 0, 255, cv2.THRESH_BINARY_INV | cv2.THRESH_OTSU)\n",
        "        # print(\"Inverse Binary:\")\n",
        "        # plot_fig(binary)\n",
        "        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))  # rectangular structure\n",
        "        # print(\"Kernel for dialation:\")\n",
        "        # print(kernel)\n",
        "        binary = cv2.dilate(binary, kernel)  # dilate\n",
        "        # print(\"Dialated Binary:\")\n",
        "        # plot_fig(binary)\n",
        "        edges = cv2.Canny(binary, 50, 200)\n",
        "        # print(\"Canny edged detection:\")\n",
        "        # plot_fig(edges)\n",
        "\n",
        "        self.lines = cv2.HoughLinesP(edges, 1, np.pi / 180, 100, minLineLength=100, maxLineGap=20)\n",
        "        # print(self.lines)\n",
        "        if self.lines is None:\n",
        "            print(\"Line segment not found\")\n",
        "            return None\n",
        "\n",
        "        lines1 = self.lines[:, 0, :]  # Extract as 2D\n",
        "        imglines = self.img.copy()\n",
        "        for x1, y1, x2, y2 in lines1[:]:\n",
        "            cv2.line(imglines, (x1, y1), (x2, y2), (0, 255, 0), 3)\n",
        "        # print(\"Probabilistic Hough Lines:\")\n",
        "        # plot_fig(imglines)\n",
        "        # return imglines\n",
        "\n",
        "    def search_lines(self):\n",
        "      lines = self.lines[:, 0, :]  # extract as 2D\n",
        "    \n",
        "      number_inexist_k = 0\n",
        "      sum_pos_k45 = number_pos_k45 = 0\n",
        "      sum_pos_k90 = number_pos_k90 = 0\n",
        "      sum_neg_k45 = number_neg_k45 = 0\n",
        "      sum_neg_k90 = number_neg_k90 = 0\n",
        "      sum_zero_k = number_zero_k = 0\n",
        "\n",
        "      for x in lines:\n",
        "          if x[2] == x[0]:\n",
        "              number_inexist_k += 1\n",
        "              continue\n",
        "          #print(degrees(atan((x[3] - x[1]) / (x[2] - x[0]))), \"pos:\", x[0], x[1], x[2], x[3], \"Slope:\",(x[3] - x[1]) / (x[2] - x[0]))\n",
        "          degree = degrees(atan((x[3] - x[1]) / (x[2] - x[0])))\n",
        "          # print(\"Degree or Slope of detected lines : \",degree)\n",
        "          if 0 < degree < 45:\n",
        "              number_pos_k45 += 1\n",
        "              sum_pos_k45 += degree\n",
        "          if 45 <= degree < 90:\n",
        "              number_pos_k90 += 1\n",
        "              sum_pos_k90 += degree\n",
        "          if -45 < degree < 0:\n",
        "              number_neg_k45 += 1\n",
        "              sum_neg_k45 += degree\n",
        "          if -90 < degree <= -45:\n",
        "              number_neg_k90 += 1\n",
        "              sum_neg_k90 += degree\n",
        "          if x[3] == x[1]:\n",
        "              number_zero_k += 1\n",
        "\n",
        "      max_number = max(number_inexist_k, number_pos_k45, number_pos_k90, number_neg_k45,number_neg_k90, number_zero_k)\n",
        "      # print(\"Num of lines in different Degree range ->\")\n",
        "      # print(\"Not a Line: \",number_inexist_k, \"| 0 to 45: \",number_pos_k45, \"| 45 to 90: \",number_pos_k90, \"| -45 to 0: \",number_neg_k45, \"| -90 to -45: \",number_neg_k90, \"| Line where y1 equals y2 :\",number_zero_k)\n",
        "    \n",
        "      if max_number == number_inexist_k:\n",
        "          return 90\n",
        "      if max_number == number_pos_k45:\n",
        "          return sum_pos_k45 / number_pos_k45\n",
        "      if max_number == number_pos_k90:\n",
        "          return sum_pos_k90 / number_pos_k90\n",
        "      if max_number == number_neg_k45:\n",
        "          return sum_neg_k45 / number_neg_k45\n",
        "      if max_number == number_neg_k90:\n",
        "          return sum_neg_k90 / number_neg_k90\n",
        "      if max_number == number_zero_k:\n",
        "          return 0\n",
        "\n",
        "    def rotate_image(self, degree):\n",
        "        \"\"\"\n",
        "        Positive angle counterclockwise rotation\n",
        "        :param degree:\n",
        "        :return:\n",
        "        \"\"\"\n",
        "        # print(\"degree:\", degree)\n",
        "        if -45 <= degree <= 0:\n",
        "            degree = degree  # #negative angle clockwise\n",
        "        if -90 <= degree < -45:\n",
        "            degree = 90 + degree  # positive angle counterclockwise\n",
        "        if 0 < degree <= 45:\n",
        "            degree = degree  # positive angle counterclockwise\n",
        "        if 45 < degree < 90:\n",
        "            degree = degree - 90  # negative angle clockwise\n",
        "        print(\"DSkew angle: \", degree)\n",
        "\n",
        "        # degree = degree - 90\n",
        "        height, width = self.img.shape[:2]\n",
        "        heightNew = int(width * fabs(sin(radians(degree))) + height * fabs(\n",
        "            cos(radians(degree))))  # This formula refers to the previous content\n",
        "        widthNew = int(height * fabs(sin(radians(degree))) + width * fabs(cos(radians(degree))))\n",
        "        # print(\"Height :\",height)\n",
        "        # print(\"Width :\",width)\n",
        "        # print(\"HeightNew :\",heightNew)\n",
        "        # print(\"WidthNew :\",widthNew)\n",
        "\n",
        "        matRotation = cv2.getRotationMatrix2D((width / 2, height / 2), degree, 1)  # rotate degree counterclockwise\n",
        "        # print(\"Mat Rotation (Before): \",matRotation)\n",
        "        matRotation[0, 2] += (widthNew - width) / 2\n",
        "        # Because after rotation, the origin of the coordinate system is the upper left corner of the new image, so it needs to be converted according to the original image\n",
        "        matRotation[1, 2] += (heightNew - height) / 2\n",
        "        # print(\"Mat Rotation (After): \",matRotation)\n",
        "\n",
        "        # Affine transformation, the background color is filled with white\n",
        "        imgRotation = cv2.warpAffine(self.img, matRotation, (widthNew, heightNew), borderValue=(255, 255, 255))\n",
        "\n",
        "        # Padding\n",
        "        pad_image_rotate = cv2.warpAffine(self.img, matRotation, (widthNew, heightNew), borderValue=(0, 255, 0))\n",
        "        # plot_fig(pad_image_rotate)\n",
        "\n",
        "        return imgRotation\n",
        "\n",
        "def dskew(line_path, img):\n",
        "    img_loc = line_path + img\n",
        "    im = cv2.imread(img_loc)\n",
        "\n",
        "    # Padding\n",
        "    bg_color = [255, 255, 255]\n",
        "    pad_img = cv2.copyMakeBorder(im,100,100,100,100,cv2.BORDER_CONSTANT,value=bg_color)\n",
        "\n",
        "    imgcorrect = ImgCorrect(pad_img)\n",
        "    lines_img = imgcorrect.img_lines()\n",
        "    # print(type(lines_img))\n",
        "    \n",
        "    if lines_img is None:\n",
        "        rotate = imgcorrect.rotate_image(0)\n",
        "    else:\n",
        "        degree = imgcorrect.search_lines()\n",
        "        rotate = imgcorrect.rotate_image(degree)\n",
        "\n",
        "\n",
        "    return rotate\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UrhXDItra-Wl"
      },
      "source": [
        "## HoughLine and Affine Transform..."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yLD-XcQ8-hfR"
      },
      "source": [
        "Defining functions to do necessary preprocessing and calculating the skew of an image's main handwritten line using Standard Hough Line Transform (SHT). And correcting it by using Affine Transform if the skew angle detected by SHT is more than 0; otherwise, send the image for DSkew. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cZVvh3sqbLsD"
      },
      "outputs": [],
      "source": [
        "rotate_line = \"/content/Rotated_line_by_HaughLine_Affine/\"\n",
        "os.mkdir(rotate_line)\n",
        "\n",
        "rotate_line_Dskew = \"/content/DSkew/\"\n",
        "os.mkdir(rotate_line_Dskew)\n",
        "\n",
        "rotate_line_Haughline = \"/content/HaughLine_Affine/\"\n",
        "os.mkdir(rotate_line_Haughline)\n",
        "\n",
        "exception = []\n",
        "\n",
        "# Degree conversion\n",
        "def DegreeTrans(theta):\n",
        "    res = theta / np.pi * 180\n",
        "    # print(res)\n",
        "    return res\n",
        "\n",
        "# Rotate the image degree counterclockwise (original size)\n",
        "def rotateImage(src, degree):\n",
        "    # The center of rotation is the center of the image\n",
        "    h, w = src.shape[:2]\n",
        "    # Calculate the two-dimensional rotating affine transformation matrix\n",
        "    RotateMatrix = cv2.getRotationMatrix2D((w / 2.0, h / 2.0), degree, 1)\n",
        "    # print(\"Rotate Matrix: \")\n",
        "    # print(RotateMatrix)\n",
        "\n",
        "    # Affine transformation, the background color is filled with GREEN so that the rotation can be easily understood\n",
        "    rotate1 = cv2.warpAffine(src, RotateMatrix, (w, h), borderValue=(0, 255, 0))\n",
        "    # plot_fig(rotate1)\n",
        "    # Affine transformation, the background color is filled with white\n",
        "    rotate = cv2.warpAffine(src, RotateMatrix, (w, h), borderValue=(255, 255, 255))\n",
        "\n",
        "    # Padding\n",
        "    bg_color = [255, 255, 255]\n",
        "    pad_image_rotate = cv2.copyMakeBorder(rotate,100,100,100,100,cv2.BORDER_CONSTANT,value=bg_color)\n",
        "\n",
        "    return pad_image_rotate\n",
        "\n",
        "# Calculate angle by Hough transform\n",
        "def CalcDegree(srcImage,canny_img):\n",
        "    lineimage = srcImage.copy()\n",
        "    lineimg = srcImage.copy()\n",
        "    # Detect straight lines by Hough transform\n",
        "    # The fourth parameter is the threshold, the greater the threshold, the higher the detection accuracy\n",
        "    try:\n",
        "        lines = cv2.HoughLines(canny_img, 1, np.pi / 180, 200)\n",
        "        # print(\"HoughLines: \")\n",
        "        # cv2_imshow(lines)\n",
        "        # Due to different images, the threshold is not easy to set, because the threshold is set too high, so that the line cannot be detected, the threshold is too low, the line is too much, the speed is very slow\n",
        "        theta_sum = 0\n",
        "        rho_sum = 0\n",
        "        sum_x1 = sum_x2 = sum_y1 = sum_y2 = 0\n",
        "        # Draw each line segment in turn\n",
        "        for i in range(len(lines)):\n",
        "            for rho, theta in lines[i]:\n",
        "                # print(\"theta:\", theta, \" rho:\", rho)\n",
        "                a = np.cos(theta)\n",
        "                b = np.sin(theta)\n",
        "                x0 = a * rho\n",
        "                y0 = b * rho\n",
        "                x1 = int(round(x0 + 1000 * (-b)))\n",
        "                y1 = int(round(y0 + 1000 * a))\n",
        "                x2 = int(round(x0 - 1000 * (-b)))\n",
        "                y2 = int(round(y0 - 1000 * a))\n",
        "                # print(\"a: \",a, \" b: \",b, \" x0: \",x0, \" y0: \",y0, \" x1: \",x1, \" y1: \",y1, \" x2: \",x2, \" y2: \",y2)\n",
        "                # Only select the smallest angle as the rotation angle\n",
        "                sum_x1+=x1\n",
        "                sum_x2+=x2\n",
        "                sum_y1+=y1\n",
        "                sum_y2+=y2\n",
        "                rho_sum += rho\n",
        "                theta_sum += theta\n",
        "                cv2.line(lineimage, (x1, y1), (x2, y2), (0, 0, 255), 1, cv2.LINE_AA)\n",
        "\n",
        "        \n",
        "        # print(\"HoughLines: \")\n",
        "        # plot_fig(lineimage)\n",
        "        print()\n",
        "\n",
        "        pt1 = (sum_x1//len(lines), sum_y1//len(lines))\n",
        "        pt2 = (sum_x2//len(lines), sum_y2//len(lines))\n",
        "        \n",
        "        # print(\"Sum of thetas: \",theta_sum)\n",
        "        # print(\"lines: \",lines)\n",
        "        average = theta_sum / len(lines)\n",
        "        # print(\"Avg. Theta: \",average)\n",
        "        angle = DegreeTrans(average) - 90\n",
        "        # print(\"Avg. Angle: \",angle)\n",
        "        print(\"Skewed Angle: \",angle)\n",
        "        average_rho = rho_sum / len(lines)\n",
        "        # print(\"Avg. rho: \",average_rho)\n",
        "\n",
        "        # print('Draw best fit line with full:')\n",
        "        # h, w = lineimg.shape[:2]\n",
        "        # pt2 = (w,h)\n",
        "        # print(\"Cordinates of the best fit line: \",pt1,pt2)\n",
        "        cv2.line(lineimg, pt1, pt2, (0,0,255), 2)\n",
        "        # plot_fig(lineimg)\n",
        "\n",
        "        return angle\n",
        "    except:\n",
        "        angle = 0.0\n",
        "        return angle\n",
        "\n",
        "def ready_for_rotate(line_path, img):\n",
        "    print()\n",
        "    print(\"Image :: \",img)\n",
        "    img_loc = line_path + img\n",
        "    image = cv2.imread(img_loc)\n",
        "\n",
        "    org_width = image.shape[1]\n",
        "    org_height = image.shape[0]\n",
        "\n",
        "    img1 = image\n",
        "    im_gray = cv2.cvtColor(img1, cv2.COLOR_BGR2GRAY)\n",
        "    # print(\"Gray Image: \")\n",
        "    # plot_fig(im_gray)\n",
        "\n",
        "    edges = cv2.Canny(im_gray,50,150,apertureSize=3)\n",
        "    # print(\"Canny Image: \")\n",
        "    # plot_fig(edges)\n",
        "\n",
        "    degree = CalcDegree(image,edges)\n",
        "    \n",
        "    if degree == 0.0:\n",
        "        rotate = dskew(line_path, img)\n",
        "        # print(\"Rotated Image by DSkew: \")\n",
        "        # plot_fig(rotate)\n",
        "        print()\n",
        "        \n",
        "        filename1 = rotate_line_Dskew + img\n",
        "        cv2.imwrite(filename1, rotate)\n",
        "        filename = rotate_line + img\n",
        "        cv2.imwrite(filename, rotate)\n",
        "    else:\n",
        "        rotate = rotateImage(image, degree)\n",
        "        # print(\"Rotated Image by Houghline Affine transform: \")\n",
        "        # plot_fig(rotate)\n",
        "        print()\n",
        "\n",
        "        filename2 = rotate_line_Haughline + img\n",
        "        cv2.imwrite(filename2, rotate)\n",
        "        filename = rotate_line + img\n",
        "        cv2.imwrite(filename, rotate)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NaBsaedbcOG_"
      },
      "outputs": [],
      "source": [
        "def rotate_lines(first_detection):\n",
        "  line_path = first_detection\n",
        "  line_dir = line_sort(os.listdir(line_path))\n",
        "  print(line_dir)\n",
        "\n",
        "  for img in line_dir:\n",
        "      ready_for_rotate(line_path, img)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "B95GQhgIDsfK"
      },
      "source": [
        "#**Final Line Segmentation**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "POoMuSzjAvB0"
      },
      "source": [
        "##Find undetected image in Yolo 2nd detection..."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cVF-hwXFCymn"
      },
      "source": [
        "Defining a function to get those images that are not getting 2nd YOLO detection for their low dimension and putting those images with the final segmented images.  "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yRuA-ABVA1gY"
      },
      "outputs": [],
      "source": [
        "undetected_images_path = []\n",
        "\n",
        "def find_undetected_images(img, label):\n",
        "  # img_path = \"/content/Rotated_line_by_HaughLine_Affine/\"\n",
        "  img_path = img\n",
        "  # detect_lb_path = \"/content/yolov5/runs/detect/exp2/labels/\"\n",
        "  detect_lb_path = label\n",
        "  undetect_img_path = \"/content/final_line_segmentation/\"\n",
        "\n",
        "  def take_valid_img(images):\n",
        "      image = []\n",
        "      valid_img_ext = [\"jpg\", \"JPG\", \"jpeg\", \"JPEG\", \"png\", \"PNG\"]\n",
        "      for img in images:\n",
        "          try:\n",
        "              ext = img.split('.')[1]\n",
        "              if ext not in valid_img_ext :\n",
        "                  continue\n",
        "              else:\n",
        "                  image.append(img)\n",
        "          except:\n",
        "              continue\n",
        "      return image\n",
        "\n",
        "  img1 = os.listdir(img_path)\n",
        "  img = take_valid_img(img1)\n",
        "  detect_lb = os.listdir(detect_lb_path)\n",
        "\n",
        "  def find_undetect_img(img,detect_lb):\n",
        "      img_lb = [im.split('.')[0] for im in img]\n",
        "      dt_lb = [dt.split('.')[0] for dt in detect_lb]\n",
        "      undt_lb = list(set(img_lb).difference(dt_lb))\n",
        "      undetect_img = []\n",
        "      detect_img = []\n",
        "      for lb in undt_lb:\n",
        "          for im in img:\n",
        "              im_lb = im.split('.')[0]\n",
        "              if lb == im_lb:\n",
        "                  undetect_img.append(im)\n",
        "              else:\n",
        "                  detect_img.append(im)\n",
        "      print(\"Undetect image: \",undetect_img)\n",
        "      write_image(undetect_img)\n",
        "\n",
        "  def write_image(undt_img):\n",
        "      for im in undt_img:\n",
        "          filename = undetect_img_path+im\n",
        "          img = cv2.imread(img_path+im)\n",
        "          cv2.imwrite(filename,img)\n",
        "          undetected_images_path.append(filename)\n",
        "\n",
        "  find_undetect_img(img,detect_lb)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FvoftpP5JZz_"
      },
      "source": [
        "## 2nd Line Segmentation after Rotation..."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ABAZ6Rh8DySL"
      },
      "source": [
        "Defining a function to select the main handwritten line from line images followed by segmenting it using YOLO's 2nd line detection. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Jsv1wlDkrKBH"
      },
      "outputs": [],
      "source": [
        "def crop_image(bb_data, destination, image, img_lb, dh, dw):\n",
        "    x = float(bb_data[1])\n",
        "    y = float(bb_data[2])\n",
        "    w = float(bb_data[3])\n",
        "    h = float(bb_data[4])\n",
        "  \n",
        "    # x = 0.5\n",
        "    # w  = 1.0\n",
        "    l = int((x - w / 2) * dw)\n",
        "    r = int((x + w / 2) * dw)\n",
        "    t = int((y - h / 2) * dh)\n",
        "    b = int((y + h / 2) * dh)\n",
        "\n",
        "    crop = image[t:b, l:r]\n",
        "    filename = destination+img_lb\n",
        "    cv2.imwrite(filename,crop)\n",
        "    print(\"Segmented successfully!\\n\")\n",
        "\n",
        "def line_segmantation_2(img, img_path, label, label_path, segmented_img_path):\n",
        "  dir = segmented_img_path\n",
        "  print(\"Image path -> \",img_path)\n",
        "  img1 = cv2.imread(img_path)\n",
        "  dh, dw, _ = img1.shape\n",
        "  txt_lb = open(label_path, 'r')\n",
        "  txt_lb_data = txt_lb.readlines()\n",
        "  txt_lb.close()\n",
        "  img_name = img\n",
        "  \n",
        "  max_w = 0\n",
        "  data1 = []\n",
        "  for line in txt_lb_data:\n",
        "      token = line.split()\n",
        "      data1.append(token)\n",
        "  \n",
        "  if len(data1)==1:\n",
        "      bb_data = data1[0]\n",
        "      wdth = float(bb_data[3])\n",
        "      if wdth>0.4:\n",
        "          crop_image(bb_data,dir,img1,img_name,dh,dw,)\n",
        "      else:\n",
        "          filename = dir+img_name\n",
        "          cv2.imwrite(filename,img1)\n",
        "  elif len(data1)==2:\n",
        "      bb_data1 = data1[0]\n",
        "      bb_data2 = data1[1]\n",
        "      w1 = float(bb_data1[3]) \n",
        "      w2 = float(bb_data2[3])\n",
        "      c1 = float(bb_data1[5]) \n",
        "      c2 = float(bb_data2[5])\n",
        "      if w1 <= 0.5 and w2 <= 0.5:\n",
        "        if c1 >= 0.8 and c2 >= 0.8:\n",
        "          sorted_bb_data = sorted(data1, key=operator.itemgetter(5))\n",
        "          bb_data = sorted_bb_data[-1]\n",
        "          crop_image(bb_data,dir,img1,img_name,dh,dw,)\n",
        "        else:\n",
        "          filename = dir+img_name\n",
        "          cv2.imwrite(filename,img1)\n",
        "      else:\n",
        "        sorted_bb_data = sorted(data1, key=operator.itemgetter(3))\n",
        "        bb_data = sorted_bb_data[-1]\n",
        "        crop_image(bb_data,dir,img1,img_name,dh,dw,)\n",
        "  elif len(data1)==3:\n",
        "      sorted_bb_data = sorted(data1, key=operator.itemgetter(2))\n",
        "      bb_data = sorted_bb_data[1]\n",
        "      crop_image(bb_data,dir,img1,img_name,dh,dw,) \n",
        "  else:\n",
        "      sorted_bb_data = sorted(data1, key=operator.itemgetter(3))\n",
        "      bb_data = sorted_bb_data[-1]\n",
        "      crop_image(bb_data,dir,img1,img_name,dh,dw,)\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1QVZUHX_8lSX"
      },
      "source": [
        "#**Word Segmentation**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tdngiWC7GNyL"
      },
      "source": [
        "Defining a function to segment the words of line images by using YOLO's word detection."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mkMTHFDv8mk2"
      },
      "outputs": [],
      "source": [
        "def word_segmentation(line_images, word_labels):\n",
        "  line_img = os.listdir(line_images)\n",
        "  word_label = os.listdir(word_labels)\n",
        "  print(line_img)\n",
        "  print(word_label)\n",
        "\n",
        "  for i in word_label:\n",
        "    if i == \"163_19_11.txt\":\n",
        "      print(\"Yes\")\n",
        "    for j in line_img:\n",
        "      if j == \"163_19_11.jpg\":\n",
        "         print(\"Yes\")\n",
        "      fn_i = i.split(\".\")\n",
        "      fn_j = j.split(\".\")\n",
        "      if fn_i[0] ==  fn_j[0]:\n",
        "        # print(\"yes\")\n",
        "        dir = \"/content/final_word_segmentation/\"+fn_i[0]\n",
        "        os.mkdir(dir)\n",
        "\n",
        "        img = cv2.imread(line_images + j)\n",
        "        dh, dw, _ = img.shape\n",
        "        txt_lb = open(word_labels + i, 'r')\n",
        "        txt_lb_data = txt_lb.readlines()\n",
        "        txt_lb.close()\n",
        "        img_lb = fn_i[0]\n",
        "        \n",
        "        k=1\n",
        "        for dt in txt_lb_data:\n",
        "            # _, x, y, w, h = map(float, dt.split(' '))\n",
        "            _, x, y, w, h, conf = map(float, dt.split(' '))\n",
        "            l = int((x - w / 2) * dw)\n",
        "            r = int((x + w / 2) * dw)\n",
        "            t = int((y - h / 2) * dh)\n",
        "            b = int((y + h / 2) * dh)\n",
        "\n",
        "            crop = img[t:b, l:r]\n",
        "            cv2.imwrite(\"{}/{}_{}.jpg\".format(dir, img_lb, k), crop)\n",
        "            k += 1\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y0v0VccCIoFw"
      },
      "source": [
        "#**Automatic annotation Starts from here** (Change File Path)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Uplaoded your dataset of raw images by following given \"bulk_sample.zip\" format on your desire location. Then, assign the zip dataset's address to define variable **dataset_directory**.  "
      ],
      "metadata": {
        "id": "KcPiYF_lMqEQ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Downloading the dataset zip from the given directory...\n",
        "dataset_directory = 'https://github.com/crusnic-corp/BN-DRISHTI/raw/main/bulk_inferencing/bulk_sample.zip'\n",
        "!wget -P /content/ $dataset_directory\n",
        "\n",
        "x = os.listdir(\"/content\")\n",
        "for i in x:\n",
        "  if i.endswith(('.zip')) == True:\n",
        "    zip_file = \"/content/\" + i\n",
        "    file_unzip(zip_file) # Unzipping the Zip file...\n",
        "    !rm -rf $zip_file # Removing downloaded zip file...\n",
        "    files1 = i"
      ],
      "metadata": {
        "id": "1Kl1F8LKpVzH"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "This section is performing these things provide below:\n",
        "1. Renaming the extracted file as 'files'.\n",
        "2. Going into each folder of the dataset, taking the document image labels and path followed by putting them into dictionary for later use.\n",
        "3. Also, creating folders to store automation annotated lines and words by our system."
      ],
      "metadata": {
        "id": "xpucrOI0M-hg"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NEA_1JfH16S7"
      },
      "outputs": [],
      "source": [
        "def filter_list(list):\n",
        "  list1 = []\n",
        "  for i in list:\n",
        "    if i.endswith(('.png', '.jpg', '.jpeg', '.tiff', '.bmp', '.gif', 'JPG', 'JPEG', 'PNG')) == True:\n",
        "      list1.append(i)\n",
        "  return list1\n",
        "\n",
        "dictonary1 = {}\n",
        "dictonary2 = {}\n",
        "\n",
        "# renaming extracted file as 'files'\n",
        "fname = files1.split('.')[0]\n",
        "src_dir = '/content/' + fname\n",
        "# src_dir = '/content/Dataset2'\n",
        "dst_dir = '/content/files'\n",
        "os.rename(src_dir, dst_dir)\n",
        "final_dir = dst_dir\n",
        "print(final_dir)\n",
        "\n",
        "# final_dir = '/content/' + files1.split('.')[0]\n",
        "# print(final_dir)\n",
        "\n",
        "file_list = os.listdir(final_dir)\n",
        "for i in file_list:\n",
        "  final_dir1 = final_dir + '/' + i + '/'\n",
        "  print(final_dir1)\n",
        "  final_dir1_list = filter_list(os.listdir(final_dir1))\n",
        "  final_dir1_list1 = line_sort(final_dir1_list)\n",
        "  print(final_dir1_list1)\n",
        "  dictonary1[final_dir1] = final_dir1_list1\n",
        "  dictonary2[i] = final_dir1_list1\n",
        "\n",
        "  final_dir2_line = final_dir1 + 'lines/'\n",
        "  final_dir2_word = final_dir1 + 'words/'\n",
        "  os.mkdir(final_dir2_line)\n",
        "  os.mkdir(final_dir2_line + 'images')\n",
        "  os.mkdir(final_dir2_line + 'labels')\n",
        "  os.mkdir(final_dir2_word)\n",
        "  os.mkdir(final_dir2_word + 'images')\n",
        "  os.mkdir(final_dir2_word + 'labels')\n",
        "print(dictonary1)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZhilBHwUJSvw"
      },
      "outputs": [],
      "source": [
        "def copy_files(path, files):\n",
        "  for i in files:\n",
        "    path1 = path + i\n",
        "    print(path1)\n",
        "    img = cv2.imread(path1)\n",
        "    # cv2_imshow(img)\n",
        "    cv2.imwrite(os.path.join(dir, i), img)\n",
        "dir = '/content/temporary_files/'\n",
        "os.mkdir(dir)\n",
        "for key in dictonary1.keys():\n",
        "  val = dictonary1.get(key)\n",
        "  print(key,'->',val)\n",
        "  copy_files(key, val)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "First line detection of document image by YOLO. "
      ],
      "metadata": {
        "id": "3POJJ_gNPAZT"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "m7XoUyJomoYW"
      },
      "outputs": [],
      "source": [
        "# 1st detection...\n",
        "img_path = dir\n",
        "img_size = 640\n",
        "conf = 0.30\n",
        "yolo_detection(img_path, img_size, conf)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sop6E7-Ctqzq"
      },
      "outputs": [],
      "source": [
        "print(dictonary2)\n",
        "# Sorting Labels of 1st detection on the basis of y...\n",
        "txt_loc = \"/content/yolov5/runs/detect/exp/labels/\"\n",
        "new_sort_label = '/content/sorted_line_after_1st_detection/'\n",
        "flag = 0\n",
        "sort_detection_label(txt_loc, new_sort_label, flag)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZRAqZed4KIBx"
      },
      "outputs": [],
      "source": [
        "# Line Segmentation after 1st Detection...\n",
        "img_path = '/content/temporary_files/'\n",
        "sorted_label = \"/content/sorted_line_after_1st_detection/\"\n",
        "line_segmentation(img_path, sorted_label)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Saag3CC0TLXC"
      },
      "outputs": [],
      "source": [
        "dictonary3 = {}\n",
        "for key in dictonary2.keys():\n",
        "  path1 = '/content/files/' + key + '/lines/images/'\n",
        "  path2 = '/content/files/' + key + '/words/images/'\n",
        "  lb_path = '/content/files/' + key + '/lines/labels/'\n",
        "\n",
        "  temp_list = []\n",
        "  val = dictonary2.get(key)\n",
        "  for i in val:\n",
        "    temp1 = i.split('.')[0]\n",
        "    temp_list.append(temp1)\n",
        "\n",
        "    line_path = path1 + temp1 + '/'\n",
        "    os.mkdir(line_path)\n",
        "\n",
        "    word_path = path2 + temp1 + '/'\n",
        "    os.mkdir(word_path)\n",
        "\n",
        "    # from_dir = \"/content/croped_line_after_1st_detection/\" + temp1\n",
        "    # to_dir = path1\n",
        "    # shutil.copytree(from_dir, to_dir)\n",
        "\n",
        "    from_dir_lb = '/content/sorted_line_after_1st_detection/' + temp1 +'.txt'\n",
        "    to_dir_lb = lb_path + temp1 +'.txt'\n",
        "    shutil.copy(from_dir_lb, to_dir_lb)\n",
        "\n",
        "  dictonary3[key] = temp_list\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fz8tTV5O_cBN"
      },
      "outputs": [],
      "source": [
        "path1 = '/content/croped_line_after_1st_detection/'\n",
        "path2 = '/content/croped_line_after_1st_detection_1/'\n",
        "os.mkdir(path2)\n",
        "temp1 = os.listdir(path1)\n",
        "for i in temp1:\n",
        "  temp2 = path1 + i + '/'\n",
        "  temp3 = os.listdir(temp2)\n",
        "  for j in temp3:\n",
        "    temp4 = temp2 + j\n",
        "    img = cv2.imread(temp4)\n",
        "    filename = path2 + j\n",
        "    cv2.imwrite(filename, img)\n",
        "# for key in dictonary3.keys():\n",
        "#   val = dictonary3.get(key)\n",
        "#   print(key,\"->\",val)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Rotating segmented lines."
      ],
      "metadata": {
        "id": "H8IaD5yfPLrY"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8oqx-J4iBi3N"
      },
      "outputs": [],
      "source": [
        "first_detection = '/content/croped_line_after_1st_detection_1/'\n",
        "rotate_lines(first_detection)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Second line detection on segmented lines by YOLO."
      ],
      "metadata": {
        "id": "BnRVkMknPZ0Q"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XABhxbhNDM-X"
      },
      "outputs": [],
      "source": [
        "# Image Path from unziped file...\n",
        "# Yolo 2nd Detection...\n",
        "rotated_img_path = \"/content/Rotated_line_by_HaughLine_Affine/\"\n",
        "img_size = 640\n",
        "conf = 0.50\n",
        "yolo_detection(rotated_img_path, img_size, conf)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zEt4RBS1Idhw"
      },
      "outputs": [],
      "source": [
        "new_dir = \"/content/final_line_segmentation/\"\n",
        "os.mkdir(new_dir)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Extracting the main handwritten line from second YOLO line detection and those lines are the final segmented lines of our document images."
      ],
      "metadata": {
        "id": "Jo3eqBOqPoOy"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "opRkhp7GImKo"
      },
      "outputs": [],
      "source": [
        "target_label_path = \"/content/yolov5/runs/detect/exp2/labels/\"\n",
        "target_image_path = \"/content/Rotated_line_by_HaughLine_Affine/\"\n",
        "target_image = os.listdir(target_image_path)\n",
        "target_label = os.listdir(target_label_path)\n",
        "\n",
        "for i in target_image:\n",
        "  for j in target_label:\n",
        "    fn_i = i.split(\".\")\n",
        "    fn_j = j.split(\".\")\n",
        "    if fn_i[0] ==  fn_j[0]:\n",
        "      # Line Segmentation after 1st Detection...\n",
        "      # Final Line Segmentation...\n",
        "      img_path = target_image_path + i\n",
        "      img = i\n",
        "      sorted_label = j\n",
        "      sorted_label_path = target_label_path + j\n",
        "      line_segmantation_2(img, img_path, sorted_label, sorted_label_path, new_dir)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Undetected lines in YOLO second line detection."
      ],
      "metadata": {
        "id": "0qM4o7FRPyT4"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "print(\"List of undetected images in 2nd detection ->\")\n",
        "find_undetected_images(target_image_path, target_label_path)\n",
        "print()\n",
        "# print(\"Undetected images are shown below -> \\n\")\n",
        "# for i in list(set(undetected_images_path)):\n",
        "#   print(\"Line image path: \",i)\n",
        "#   undit_img = cv2.imread(i)\n",
        "#   plot_fig(undit_img)\n",
        "#   print()"
      ],
      "metadata": {
        "id": "m980sgXKfjbi"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Trimming the DSkewed images from all the side to reduce the padding caused by scaling. "
      ],
      "metadata": {
        "id": "vyejXFpYP7ow"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def trim_original_image(rotate, org_w, org_h):\n",
        "  org_width = org_w\n",
        "  org_height = org_h\n",
        "\n",
        "  img1 = rotate\n",
        "  width = img1.shape[1]\n",
        "  height = img1.shape[0]\n",
        "  print(\"Original height -> \",org_height)\n",
        "  print(\"Original width -> \",org_width)\n",
        "\n",
        "  start_row = 60\n",
        "  end_row = height - 60\n",
        "\n",
        "  start_col = 60\n",
        "  end_col = width - 60\n",
        "  img_new = img1[start_row:end_row, start_col:end_col]\n",
        "\n",
        "  width1 = img_new.shape[1]\n",
        "  height1 = img_new.shape[0]\n",
        "  print(\"New height -> \",height1)\n",
        "  print(\"New width -> \",width1)\n",
        "    \n",
        "  return img_new\n",
        "\n",
        "\n",
        "final_line_segment = \"/content/final_line_segmentation/\"\n",
        "rotate_line_Dskew = \"/content/DSkew/\"\n",
        "dskew_img_list = os.listdir(rotate_line_Dskew)\n",
        "print(dskew_img_list)\n",
        "for i in dskew_img_list:\n",
        "  print(\"Target image -> \",i)\n",
        "  temp = final_line_segment + i\n",
        "  print(\"Path -> \",temp)\n",
        "  img1 = cv2.imread(temp)\n",
        "  height, width, channels = img1.shape\n",
        "  temp2 = rotate_line_Dskew + i\n",
        "  img2 = cv2.imread(temp2)\n",
        "  height2, width2, channels = img2.shape\n",
        "  if height >= height2:\n",
        "    # os.remove(temp)\n",
        "    temp3 = trim_original_image(img1, width, height)\n",
        "    cv2.imwrite(temp, temp3)\n",
        "    # print(\"Original ->\")\n",
        "    # plot_fig(img1)\n",
        "    # print(\"New ->\")\n",
        "    # plot_fig(temp3)\n",
        "    print()\n",
        "  # else:\n",
        "  #   print(\"Original ->\")\n",
        "  #   plot_fig(img1)\n",
        "  #   print(\"No need for change as its gone through the 2nd line detection!\\n\")\n"
      ],
      "metadata": {
        "id": "-1cyllt-fsZh"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-QLd3DLXPHDR"
      },
      "outputs": [],
      "source": [
        "# Copying final segmented lines to the final folders (files)...\n",
        "def trim_label(x):\n",
        "  t1 = x.split('_')[0]\n",
        "  t2 = x.split('_')[1]\n",
        "  t3 = t1+\"_\"+t2\n",
        "  return t3\n",
        "\n",
        "target_dir = \"/content/files/\"\n",
        "final_dir = \"/content/final_line_segmentation/\"\n",
        "final_dir_list = os.listdir(final_dir)\n",
        "for key in dictonary3.keys():\n",
        "  values = dictonary3.get(key)\n",
        "  for val in values:\n",
        "    target_dir1 = target_dir + key + \"/lines/images/\" + val +\"/\"\n",
        "    for i in final_dir_list:\n",
        "      temp = trim_label(i)\n",
        "      if val == temp:\n",
        "        src_path = final_dir + i\n",
        "        dst_path = target_dir1 + i\n",
        "        shutil.copy(src_path, dst_path)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# **Word detection on final segmented lines** (Upload the given trained Word Model weights on your drive and assign its directory to the --weights command.)."
      ],
      "metadata": {
        "id": "7fKZXKtAQxUr"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!python /content/yolov5/detect.py --weights /content/model/word_model_best.pt --img 640 --conf 0.40 --source /content/final_line_segmentation --save-conf --save-txt"
      ],
      "metadata": {
        "id": "SUI1DtuWDYVg"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Sorting Labels of Word detection on the basis of x...\n",
        "txt_loc = \"/content/yolov5/runs/detect/exp3/labels/\"\n",
        "new_sort_label = '/content/sorted_Word_detection/'\n",
        "flag = 1\n",
        "sort_detection_label(txt_loc, new_sort_label, flag)"
      ],
      "metadata": {
        "id": "wq39o0BPECYX"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Word Segmentation... \n",
        "word_labels = \"/content/sorted_Word_detection/\"\n",
        "line_images = \"/content/final_line_segmentation/\"\n",
        "final_word_dir = \"/content/final_word_segmentation/\"\n",
        "os.mkdir(final_word_dir)\n",
        "word_segmentation(line_images, word_labels)"
      ],
      "metadata": {
        "id": "wFchpwLMERif"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print(dictonary1)\n",
        "print(dictonary2)\n",
        "print(dictonary3)"
      ],
      "metadata": {
        "id": "8nogal5D5qG2"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "word_seg_dir = \"/content/final_word_segmentation/\"\n",
        "word_seg_list_dir = os.listdir(word_seg_dir)\n",
        "word_seg_lb_dir = \"/content/sorted_Word_detection/\"\n",
        "word_seg_lb_list_dir = os.listdir(word_seg_lb_dir)\n",
        "print(word_seg_list_dir)\n",
        "print(word_seg_lb_list_dir)\n",
        "for key in dictonary3:\n",
        "  values = dictonary3.get(key)\n",
        "  for val in values:\n",
        "    for i in word_seg_list_dir:\n",
        "      temp = trim_label(i)\n",
        "      if val == temp:\n",
        "        from_dir = word_seg_dir + i\n",
        "        to_dir = \"/content/files/\" + key + \"/words/images/\" +val + \"/\" + i\n",
        "        if os.path.exists(to_dir) == False:\n",
        "            shutil.copytree(from_dir, to_dir)\n",
        "        from_dir_lb = word_seg_lb_dir + i + \".txt\"\n",
        "        to_dir_lb = \"/content/files/\" + key + \"/words/labels/\" + i + \".txt\"\n",
        "        if os.path.exists(to_dir_lb) == False:\n",
        "            shutil.copy(from_dir_lb, to_dir_lb)\n",
        "        # print(from_dir_lb)\n",
        "        # print(to_dir_lb)"
      ],
      "metadata": {
        "id": "3OKUJfsN5s8b"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Renaming the folder with ultimate results."
      ],
      "metadata": {
        "id": "jvMUoIYFSMy6"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "src_dir = \"/content/files/\"\n",
        "final_dir = '/content/automatic_annotation'\n",
        "os.rename(src_dir, final_dir)"
      ],
      "metadata": {
        "id": "nLKBeT9pSS76"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Saving Results"
      ],
      "metadata": {
        "id": "A3V-wLcFbA6c"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Saving Automatic annotations on the given directory (If want to save the generated outputs: Please uncomment the next code section and assign directory of any folder of your drive to the defined 'final_save' variable)."
      ],
      "metadata": {
        "id": "FNVk1rsxQ9L0"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Mounting Drive\n",
        "# from google.colab import drive\n",
        "# drive.mount('/content/drive')\n",
        "\n",
        "# final_save = \"/content/drive/MyDrive/thesis/Transition_Outputs/\"\n",
        "# temp1 = final_save + \"/files/\"\n",
        "# if os.path.exists(temp1) == True:\n",
        "#   print(\"The created folder for given image already exist in Drive!\")\n",
        "#   print(\"Please check in this location: \"+final_save)\n",
        "# else:\n",
        "#   %cp -rip $final_dir $final_save\n",
        "#   print(\"Saved successfully!\")"
      ],
      "metadata": {
        "id": "1n3fTr8OI7z9"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Another way is to uncomment the below cell and download the zip, although we do not recommend this if your zip size is more than 100 MB."
      ],
      "metadata": {
        "id": "Ph2JBhK7CQXz"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# def zip_results():\n",
        "#   os.mkdir(\"/content/test_results\")\n",
        "\n",
        "#   # Zipping the test Results\n",
        "#   !zip -r /content/test_results/Final_Results.zip /content/automatic_annotation\n",
        "\n",
        "# zip_results()\n",
        "\n",
        "# # Downloading the zipped files from our test_results folder\n",
        "# files.download('/content/test_results/Final_Results.zip')"
      ],
      "metadata": {
        "id": "mCIcwW5CsxPZ"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "colab": {
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "gpuClass": "standard"
  },
  "nbformat": 4,
  "nbformat_minor": 0
}