WindowOCR.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/s2young/ng-file-upload/blob/master/WindowOCR.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!sudo apt-get install poppler-utils\n",
        "\n",
        "import torch\n",
        "use_gpu = torch.cuda.is_available()\n",
        "\n",
        "!pip3 -q install --upgrade pip\n",
        "if use_gpu:\n",
        " !python -m pip -q install paddlepaddle-gpu -i https://pypi.tuna.tsinghua.edu.cn/simple\n",
        "else:\n",
        "  !python -m pip -q install paddlepaddle -i https://pypi.tuna.tsinghua.edu.cn/simple\n",
        "\n",
        "!pip -q install -U sahi\n",
        "!pip -q install pdf2image\n",
        "!pip -q install ultralytics\n",
        "!pip -q install \"paddleocr>=2.0.1\""
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "3QU7TcRtGX-X",
        "outputId": "e8efbbb6-3c8e-499b-bcff-f0211c4a8db5"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Reading package lists... Done\n",
            "Building dependency tree... Done\n",
            "Reading state information... Done\n",
            "The following NEW packages will be installed:\n",
            "  poppler-utils\n",
            "0 upgraded, 1 newly installed, 0 to remove and 49 not upgraded.\n",
            "Need to get 186 kB of archives.\n",
            "After this operation, 696 kB of additional disk space will be used.\n",
            "Get:1 http://archive.ubuntu.com/ubuntu jammy-updates/main amd64 poppler-utils amd64 22.02.0-2ubuntu0.5 [186 kB]\n",
            "Fetched 186 kB in 1s (329 kB/s)\n",
            "debconf: unable to initialize frontend: Dialog\n",
            "debconf: (No usable dialog-like program is installed, so the dialog based frontend cannot be used. at /usr/share/perl5/Debconf/FrontEnd/Dialog.pm line 78, <> line 1.)\n",
            "debconf: falling back to frontend: Readline\n",
            "debconf: unable to initialize frontend: Readline\n",
            "debconf: (This frontend requires a controlling tty.)\n",
            "debconf: falling back to frontend: Teletype\n",
            "dpkg-preconfigure: unable to re-open stdin: \n",
            "Selecting previously unselected package poppler-utils.\n",
            "(Reading database ... 123629 files and directories currently installed.)\n",
            "Preparing to unpack .../poppler-utils_22.02.0-2ubuntu0.5_amd64.deb ...\n",
            "Unpacking poppler-utils (22.02.0-2ubuntu0.5) ...\n",
            "Setting up poppler-utils (22.02.0-2ubuntu0.5) ...\n",
            "Processing triggers for man-db (2.10.2-1) ...\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.8/1.8 MB\u001b[0m \u001b[31m24.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m758.9/758.9 MB\u001b[0m \u001b[31m7.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m62.2/62.2 MB\u001b[0m \u001b[31m58.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Building wheel for fire (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m886.3/886.3 kB\u001b[0m \u001b[31m23.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m544.7/544.7 kB\u001b[0m \u001b[31m14.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m912.2/912.2 kB\u001b[0m \u001b[31m36.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m3.1/3.1 MB\u001b[0m \u001b[31m81.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!gdown 1-jxQaIQaexQHz-BwxgCeCaElu2V539i8"
      ],
      "metadata": {
        "id": "L5PcinJ6zaAL"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "import cv2\n",
        "import json\n",
        "import numpy as np\n",
        "from glob import glob\n",
        "from ultralytics import YOLO\n",
        "from paddleocr import PaddleOCR\n",
        "from pdf2image import convert_from_path\n",
        "\n",
        "from sahi import AutoDetectionModel\n",
        "from sahi.utils.cv import read_image\n",
        "from sahi.utils.file import download_from_url\n",
        "from sahi.utils.yolov8 import download_yolov8m_model\n",
        "from sahi.predict import get_sliced_prediction, predict\n",
        "\n",
        "\n",
        "ocr = PaddleOCR(use_angle_cls=True, lang='en') # or ch\n",
        "\n",
        "# Download YOLOv8 model\n",
        "yolov8_model_path = \"bestWindow_v8m.pt\"\n",
        "\n",
        "\n",
        "detection_model = AutoDetectionModel.from_pretrained(\n",
        "    model_type='yolov8',\n",
        "    model_path=yolov8_model_path,\n",
        "    confidence_threshold=0.3,\n",
        "    device='cuda:0',  # or 'cuda:0',\n",
        "    image_size=608,\n",
        ")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "BDnkHhehGqn5",
        "outputId": "09a47594-e67b-41ba-d11f-7496200dede1"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Creating new Ultralytics Settings v0.0.6 file ✅ \n",
            "View Ultralytics Settings with 'yolo settings' or at '/root/.config/Ultralytics/settings.json'\n",
            "Update Settings with 'yolo settings key=value', i.e. 'yolo settings runs_dir=path/to/dir'. For help see https://docs.ultralytics.com/quickstart/#ultralytics-settings.\n",
            "download https://paddleocr.bj.bcebos.com/PP-OCRv3/english/en_PP-OCRv3_det_infer.tar to /root/.paddleocr/whl/det/en/en_PP-OCRv3_det_infer/en_PP-OCRv3_det_infer.tar\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 3910/3910 [00:15<00:00, 248.71it/s] \n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "download https://paddleocr.bj.bcebos.com/PP-OCRv4/english/en_PP-OCRv4_rec_infer.tar to /root/.paddleocr/whl/rec/en/en_PP-OCRv4_rec_infer/en_PP-OCRv4_rec_infer.tar\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 10000/10000 [00:17<00:00, 586.22it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "download https://paddleocr.bj.bcebos.com/dygraph_v2.0/ch/ch_ppocr_mobile_v2.0_cls_infer.tar to /root/.paddleocr/whl/cls/ch_ppocr_mobile_v2.0_cls_infer/ch_ppocr_mobile_v2.0_cls_infer.tar\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 2138/2138 [00:14<00:00, 150.02it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[2024/11/15 09:50:04] ppocr DEBUG: Namespace(help='==SUPPRESS==', use_gpu=True, use_xpu=False, use_npu=False, use_mlu=False, ir_optim=True, use_tensorrt=False, min_subgraph_size=15, precision='fp32', gpu_mem=500, gpu_id=0, image_dir=None, page_num=0, det_algorithm='DB', det_model_dir='/root/.paddleocr/whl/det/en/en_PP-OCRv3_det_infer', det_limit_side_len=960, det_limit_type='max', det_box_type='quad', det_db_thresh=0.3, det_db_box_thresh=0.6, det_db_unclip_ratio=1.5, max_batch_size=10, use_dilation=False, det_db_score_mode='fast', det_east_score_thresh=0.8, det_east_cover_thresh=0.1, det_east_nms_thresh=0.2, det_sast_score_thresh=0.5, det_sast_nms_thresh=0.2, det_pse_thresh=0, det_pse_box_thresh=0.85, det_pse_min_area=16, det_pse_scale=1, scales=[8, 16, 32], alpha=1.0, beta=1.0, fourier_degree=5, rec_algorithm='SVTR_LCNet', rec_model_dir='/root/.paddleocr/whl/rec/en/en_PP-OCRv4_rec_infer', rec_image_inverse=True, rec_image_shape='3, 48, 320', rec_batch_num=6, max_text_length=25, rec_char_dict_path='/usr/local/lib/python3.10/dist-packages/paddleocr/ppocr/utils/en_dict.txt', use_space_char=True, vis_font_path='./doc/fonts/simfang.ttf', drop_score=0.5, e2e_algorithm='PGNet', e2e_model_dir=None, e2e_limit_side_len=768, e2e_limit_type='max', e2e_pgnet_score_thresh=0.5, e2e_char_dict_path='./ppocr/utils/ic15_dict.txt', e2e_pgnet_valid_set='totaltext', e2e_pgnet_mode='fast', use_angle_cls=True, cls_model_dir='/root/.paddleocr/whl/cls/ch_ppocr_mobile_v2.0_cls_infer', cls_image_shape='3, 48, 192', label_list=['0', '180'], cls_batch_num=6, cls_thresh=0.9, enable_mkldnn=False, cpu_threads=10, use_pdserving=False, warmup=False, sr_model_dir=None, sr_image_shape='3, 32, 128', sr_batch_num=1, draw_img_save_dir='./inference_results', save_crop_res=False, crop_res_save_dir='./output', use_mp=False, total_process_num=1, process_id=0, benchmark=False, save_log_path='./log_output/', show_log=True, use_onnx=False, return_word_box=False, output='./output', table_max_len=488, table_algorithm='TableAttn', table_model_dir=None, merge_no_span_structure=True, table_char_dict_path=None, formula_algorithm='LaTeXOCR', formula_model_dir=None, formula_char_dict_path=None, formula_batch_num=1, layout_model_dir=None, layout_dict_path=None, layout_score_threshold=0.5, layout_nms_threshold=0.5, kie_algorithm='LayoutXLM', ser_model_dir=None, re_model_dir=None, use_visual_backbone=True, ser_dict_path='../train_data/XFUND/class_list_xfun.txt', ocr_order_method=None, mode='structure', image_orientation=False, layout=True, table=True, formula=False, ocr=True, recovery=False, recovery_to_markdown=False, use_pdf2docx_api=False, invert=False, binarize=False, alphacolor=(255, 255, 255), lang='en', det=True, rec=True, type='ocr', savefile=False, ocr_version='PP-OCRv4', structure_version='PP-StructureV2')\n",
            "[2024/11/15 09:50:04] ppocr WARNING: The first GPU is used for inference by default, GPU ID: 0\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[2024/11/15 09:50:05] ppocr WARNING: The first GPU is used for inference by default, GPU ID: 0\n",
            "[2024/11/15 09:50:06] ppocr WARNING: The first GPU is used for inference by default, GPU ID: 0\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "def detect_ocr(file_path, detection_model, save_dir='detect_ocr'):\n",
        "\n",
        "    os.makedirs(f'{save_dir}', exist_ok=True)\n",
        "\n",
        "    images = convert_from_path(file_path)\n",
        "\n",
        "    for page_idx, image in enumerate(images):\n",
        "\n",
        "        img = np.array(image)\n",
        "        img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)\n",
        "\n",
        "        dets = []\n",
        "\n",
        "        results = get_sliced_prediction(\n",
        "                    img,\n",
        "                    detection_model,\n",
        "                    slice_height=600,\n",
        "                    slice_width=600,\n",
        "                    overlap_height_ratio=0.25,\n",
        "                    overlap_width_ratio=0.25\n",
        "                    ).object_prediction_list\n",
        "\n",
        "        if results != []:\n",
        "            for result in results:\n",
        "                x1, y1, x2, y2 = result.bbox.to_xyxy()\n",
        "                x = int((x1 + x2)/2)\n",
        "                y = int((y1 + y2)/2)\n",
        "                w = int(x2 - x1)\n",
        "                h = int(y2 - y1)\n",
        "\n",
        "                if h < w:\n",
        "                    new_x1 = x - w\n",
        "                    new_x2 = x + w\n",
        "                    new_y1 = y - 2*h\n",
        "                    new_y2 = y + 2*h\n",
        "                    cropped = img[new_y1:new_y2, new_x1:new_x2]\n",
        "\n",
        "                else:\n",
        "                    new_x1 = x - 2*w\n",
        "                    new_x2 = x + 2*w\n",
        "                    new_y1 = y - h\n",
        "                    new_y2 = y + h\n",
        "                    cropped = img[new_y1:new_y2, new_x1:new_x2]\n",
        "                    cropped = cv2.rotate(cropped, cv2.ROTATE_90_CLOCKWISE)\n",
        "\n",
        "                if cropped is not None:\n",
        "                    txts = []\n",
        "                    if cropped.shape[0] * cropped.shape[1] != 0:\n",
        "                        result1 = ocr.ocr(cropped, cls=True)[0]\n",
        "                        if result1 is not None:\n",
        "                            locs1 = [line[0] for line in result1]\n",
        "                            txts1 = [line[1][0] for line in result1]\n",
        "                            confs1 = [line[1][1] for line in result1]\n",
        "                            for idx in range(len(locs1)):\n",
        "                                txt = txts1[idx]\n",
        "                                conf = confs1[idx]\n",
        "                                xs, ys = zip(*locs1[idx])\n",
        "                                x11 = int(min(xs))\n",
        "                                y11 = int(min(ys))\n",
        "                                x22 = int(max(xs))\n",
        "                                y22 = int(max(ys))\n",
        "\n",
        "                                if h >= w:\n",
        "                                    h1, w1 = cropped.shape[:2]\n",
        "                                    new_x11 = y11\n",
        "                                    new_y11 = w1 - x22\n",
        "                                    new_x22 = y22\n",
        "                                    new_y22 = w1 - x11\n",
        "\n",
        "                                    new_x11 += new_x1\n",
        "                                    new_y11 += new_y1\n",
        "                                    new_x22 += new_x1\n",
        "                                    new_y22 += new_y1\n",
        "\n",
        "                                else:\n",
        "                                    new_x11 = x11 + new_x1\n",
        "                                    new_y11 = y11 + new_y1\n",
        "                                    new_x22 = x22 + new_x1\n",
        "                                    new_y22 = y22 + new_y1\n",
        "\n",
        "                                cv2.rectangle(img, (new_x11, new_y11), (new_x22, new_y22), (255, 0, 0), 1)\n",
        "                                cv2.putText(img, txt, (new_x11, new_y11-3), cv2.FONT_HERSHEY_COMPLEX_SMALL, 0.7, (255, 0, 0), 1)\n",
        "                                txts.append(txt)\n",
        "\n",
        "                    cv2.rectangle(img, (int(x1), int(y1)), (int(x2), int(y2)), (255, 0, 255), 2)\n",
        "\n",
        "                    for txt in txts:\n",
        "                        txt = txt.replace('|', '').replace(' ', '')\n",
        "\n",
        "                        if len(txt) in [9, 10] and txt[0] == '(' and txt[1].isdigit():\n",
        "                            try:\n",
        "                                num = int(txt[1])\n",
        "                                width = int(txt[3:5])\n",
        "                                height = int(txt[5:7])\n",
        "                                style = txt[7:]\n",
        "\n",
        "                                dets.append({'loc':[x1, y1, x2, y2],\n",
        "                                            'num':num,\n",
        "                                            'width':width,\n",
        "                                            'height':height,\n",
        "                                            'style':style})\n",
        "\n",
        "                                break\n",
        "\n",
        "                            except:\n",
        "                                pass\n",
        "\n",
        "\n",
        "\n",
        "                        elif len(txt) in [6, 7] and txt[0].isdigit():\n",
        "                            try:\n",
        "                                width = int(txt[:2])\n",
        "                                height = int(txt[2:4])\n",
        "                                style = txt[4:]\n",
        "\n",
        "                                dets.append({'loc':[x1, y1, x2, y2],\n",
        "                                            'num':1,\n",
        "                                            'width':width,\n",
        "                                            'height':height,\n",
        "                                            'style':style})\n",
        "\n",
        "                            except:\n",
        "                                pass\n",
        "\n",
        "\n",
        "\n",
        "        with open(f'{save_dir}/{page_idx+1}_pred.json', 'w') as f:\n",
        "            json.dump(dets, f)\n",
        "\n",
        "        cv2.imwrite(f'{save_dir}/{page_idx+1}_pred.png', img)"
      ],
      "metadata": {
        "id": "f3cNT2QnGzeh"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "file_path = ''\n",
        "detect_ocr(file_path, detection_model, save_dir='detect_ocr')"
      ],
      "metadata": {
        "id": "uOMVNXIvHurQ"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}