examples/gptpdf_Quick_Tour.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "private_outputs": true,
      "provenance": [],
      "authorship_tag": "ABX9TyPw6K7sJaE/ySw8g6ZjwqpY",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/alexfazio/gptpdf/blob/main/examples/gptpdf_Quick_Tour.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Quick Tour of `GPTPDF`\n",
        "\n",
        "The example below serves as a quick start guide on how to start using `gptpdf`.\n",
        "\n",
        "For a detailed overview of the features, please visit our [documentation page](https://github.com/CosmosShadow/gptpdf)."
      ],
      "metadata": {
        "id": "omUvC33EBFbv"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# @title 🛠️ Install Requirements\n",
        "!pip install fitz shapely GeneralAgent shapely pymupdf python-dotenv\n",
        "!pip install --force-reinstall pymupdf"
      ],
      "metadata": {
        "id": "3ef6XZP_O680",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UR_97-9NO5nH",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "# @title ⚙️ Function Definitions\n",
        "\n",
        "import os\n",
        "import re\n",
        "from typing import List, Tuple, Optional, Dict\n",
        "import logging\n",
        "\n",
        "logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')\n",
        "import fitz\n",
        "import shapely.geometry as sg\n",
        "from shapely.geometry.base import BaseGeometry\n",
        "from shapely.validation import explain_validity\n",
        "import concurrent.futures\n",
        "\n",
        "# This Default Prompt Using Chinese and could be changed to other languages.\n",
        "\n",
        "DEFAULT_PROMPT = \"\"\"使用markdown语法，将图片中识别到的文字转换为markdown格式输出。你必须做到：\n",
        "1. 输出和使用识别到的图片的相同的语言，例如，识别到英语的字段，输出的内容必须是英语。\n",
        "2. 不要解释和输出无关的文字，直接输出图片中的内容。例如，严禁输出 “以下是我根据图片内容生成的markdown文本：”这样的例子，而是应该直接输出markdown。\n",
        "3. 内容不要包含在```markdown ```中、段落公式使用 $$ $$ 的形式、行内公式使用 $ $ 的形式、忽略掉长直线、忽略掉页码。\n",
        "再次强调，不要解释和输出无关的文字，直接输出图片中的内容。\n",
        "\"\"\"\n",
        "DEFAULT_RECT_PROMPT = \"\"\"图片中用红色框和名称(%s)标注出了一些区域。如果区域是表格或者图片，使用 ![]() 的形式插入到输出内容中，否则直接输出文字内容。\n",
        "\"\"\"\n",
        "DEFAULT_ROLE_PROMPT = \"\"\"你是一个PDF文档解析器，使用markdown和latex语法输出图片的内容。\n",
        "\"\"\"\n",
        "\n",
        "\n",
        "def _is_near(rect1: BaseGeometry, rect2: BaseGeometry, distance: float = 20) -> bool:\n",
        "    \"\"\"\n",
        "    Check if two rectangles are near each other if the distance between them is less than the target.\n",
        "    \"\"\"\n",
        "    return rect1.buffer(0.1).distance(rect2.buffer(0.1)) < distance\n",
        "\n",
        "\n",
        "def _is_horizontal_near(rect1: BaseGeometry, rect2: BaseGeometry, distance: float = 100) -> bool:\n",
        "    \"\"\"\n",
        "    Check if two rectangles are near horizontally if one of them is a horizontal line.\n",
        "    \"\"\"\n",
        "    result = False\n",
        "    if abs(rect1.bounds[3] - rect1.bounds[1]) < 0.1 or abs(rect2.bounds[3] - rect2.bounds[1]) < 0.1:\n",
        "        if abs(rect1.bounds[0] - rect2.bounds[0]) < 0.1 and abs(rect1.bounds[2] - rect2.bounds[2]) < 0.1:\n",
        "            result = abs(rect1.bounds[3] - rect2.bounds[3]) < distance\n",
        "    return result\n",
        "\n",
        "\n",
        "def _union_rects(rect1: BaseGeometry, rect2: BaseGeometry) -> BaseGeometry:\n",
        "    \"\"\"\n",
        "    Union two rectangles.\n",
        "    \"\"\"\n",
        "    return sg.box(*(rect1.union(rect2).bounds))\n",
        "\n",
        "\n",
        "def _merge_rects(rect_list: List[BaseGeometry], distance: float = 20, horizontal_distance: Optional[float] = None) -> \\\n",
        "        List[BaseGeometry]:\n",
        "    \"\"\"\n",
        "    Merge rectangles in the list if the distance between them is less than the target.\n",
        "    \"\"\"\n",
        "    merged = True\n",
        "    while merged:\n",
        "        merged = False\n",
        "        new_rect_list = []\n",
        "        while rect_list:\n",
        "            rect = rect_list.pop(0)\n",
        "            for other_rect in rect_list:\n",
        "                if _is_near(rect, other_rect, distance) or (\n",
        "                        horizontal_distance and _is_horizontal_near(rect, other_rect, horizontal_distance)):\n",
        "                    rect = _union_rects(rect, other_rect)\n",
        "                    rect_list.remove(other_rect)\n",
        "                    merged = True\n",
        "            new_rect_list.append(rect)\n",
        "        rect_list = new_rect_list\n",
        "    return rect_list\n",
        "\n",
        "\n",
        "def _adsorb_rects_to_rects(source_rects: List[BaseGeometry], target_rects: List[BaseGeometry], distance: float = 10) -> \\\n",
        "        Tuple[List[BaseGeometry], List[BaseGeometry]]:\n",
        "    \"\"\"\n",
        "    Adsorb a set of rectangles to another set of rectangles.\n",
        "    \"\"\"\n",
        "    new_source_rects = []\n",
        "    for text_area_rect in source_rects:\n",
        "        adsorbed = False\n",
        "        for index, rect in enumerate(target_rects):\n",
        "            if _is_near(text_area_rect, rect, distance):\n",
        "                rect = _union_rects(text_area_rect, rect)\n",
        "                target_rects[index] = rect\n",
        "                adsorbed = True\n",
        "                break\n",
        "        if not adsorbed:\n",
        "            new_source_rects.append(text_area_rect)\n",
        "    return new_source_rects, target_rects\n",
        "\n",
        "\n",
        "def _parse_rects(page: fitz.Page) -> List[Tuple[float, float, float, float]]:\n",
        "    \"\"\"\n",
        "    Parse drawings in the page and merge adjacent rectangles.\n",
        "    \"\"\"\n",
        "\n",
        "    # 提取画的内容\n",
        "    drawings = page.get_drawings()\n",
        "\n",
        "    # 忽略掉长度小于30的水平直线\n",
        "    is_short_line = lambda x: abs(x['rect'][3] - x['rect'][1]) < 1 and abs(x['rect'][2] - x['rect'][0]) < 30\n",
        "    drawings = [drawing for drawing in drawings if not is_short_line(drawing)]\n",
        "\n",
        "    # 转换为shapely的矩形\n",
        "    rect_list = [sg.box(*drawing['rect']) for drawing in drawings]\n",
        "\n",
        "    # 提取图片区域\n",
        "    images = page.get_image_info()\n",
        "    image_rects = [sg.box(*image['bbox']) for image in images]\n",
        "\n",
        "    # 合并drawings和images\n",
        "    rect_list += image_rects\n",
        "\n",
        "    merged_rects = _merge_rects(rect_list, distance=10, horizontal_distance=100)\n",
        "    merged_rects = [rect for rect in merged_rects if explain_validity(rect) == 'Valid Geometry']\n",
        "\n",
        "    # 将大文本区域和小文本区域分开处理: 大文本相小合并，小文本靠近合并\n",
        "    is_large_content = lambda x: (len(x[4]) / max(1, len(x[4].split('\\n')))) > 5\n",
        "    small_text_area_rects = [sg.box(*x[:4]) for x in page.get_text('blocks') if not is_large_content(x)]\n",
        "    large_text_area_rects = [sg.box(*x[:4]) for x in page.get_text('blocks') if is_large_content(x)]\n",
        "    _, merged_rects = _adsorb_rects_to_rects(large_text_area_rects, merged_rects, distance=0.1) # 完全相交\n",
        "    _, merged_rects = _adsorb_rects_to_rects(small_text_area_rects, merged_rects, distance=5) # 靠近\n",
        "\n",
        "    # 再次自身合并\n",
        "    merged_rects = _merge_rects(merged_rects, distance=10)\n",
        "\n",
        "    # 过滤比较小的矩形\n",
        "    merged_rects = [rect for rect in merged_rects if rect.bounds[2] - rect.bounds[0] > 20 and rect.bounds[3] - rect.bounds[1] > 20]\n",
        "\n",
        "    return [rect.bounds for rect in merged_rects]\n",
        "\n",
        "\n",
        "def _parse_pdf_to_images(pdf_path: str, output_dir: str = './') -> List[Tuple[str, List[str]]]:\n",
        "    \"\"\"\n",
        "    Parse PDF to images and save to output_dir.\n",
        "    \"\"\"\n",
        "    # 打开PDF文件\n",
        "    pdf_document = fitz.open(pdf_path)\n",
        "    image_infos = []\n",
        "\n",
        "    for page_index, page in enumerate(pdf_document):\n",
        "        logging.info(f'parse page: {page_index}')\n",
        "        rect_images = []\n",
        "        rects = _parse_rects(page)\n",
        "        for index, rect in enumerate(rects):\n",
        "            fitz_rect = fitz.Rect(rect)\n",
        "            # 保存页面为图片\n",
        "            pix = page.get_pixmap(clip=fitz_rect, matrix=fitz.Matrix(4, 4))\n",
        "            name = f'{page_index}_{index}.png'\n",
        "            pix.save(os.path.join(output_dir, name))\n",
        "            rect_images.append(name)\n",
        "            # # 在页面上绘制红色矩形\n",
        "            big_fitz_rect = fitz.Rect(fitz_rect.x0 - 1, fitz_rect.y0 - 1, fitz_rect.x1 + 1, fitz_rect.y1 + 1)\n",
        "            # 空心矩形\n",
        "            page.draw_rect(big_fitz_rect, color=(1, 0, 0), width=1)\n",
        "            # 画矩形区域(实心)\n",
        "            # page.draw_rect(big_fitz_rect, color=(1, 0, 0), fill=(1, 0, 0))\n",
        "            # 在矩形内的左上角写上矩形的索引name，添加一些偏移量\n",
        "            text_x = fitz_rect.x0 + 2\n",
        "            text_y = fitz_rect.y0 + 10\n",
        "            text_rect = fitz.Rect(text_x, text_y - 9, text_x + 80, text_y + 2)\n",
        "            # 绘制白色背景矩形\n",
        "            page.draw_rect(text_rect, color=(1, 1, 1), fill=(1, 1, 1))\n",
        "            # 插入带有白色背景的文字\n",
        "            page.insert_text((text_x, text_y), name, fontsize=10, color=(1, 0, 0))\n",
        "        page_image_with_rects = page.get_pixmap(matrix=fitz.Matrix(3, 3))\n",
        "        page_image = os.path.join(output_dir, f'{page_index}.png')\n",
        "        page_image_with_rects.save(page_image)\n",
        "        image_infos.append((page_image, rect_images))\n",
        "\n",
        "    pdf_document.close()\n",
        "    return image_infos\n",
        "\n",
        "\n",
        "def _gpt_parse_images(\n",
        "        image_infos: List[Tuple[str, List[str]]],\n",
        "        prompt_dict: Optional[Dict] = None,\n",
        "        output_dir: str = './',\n",
        "        api_key: Optional[str] = None,\n",
        "        base_url: Optional[str] = None,\n",
        "        model: str = 'gpt-4o',\n",
        "        verbose: bool = False,\n",
        "        gpt_worker: int = 1,\n",
        "        **args\n",
        ") -> str:\n",
        "    \"\"\"\n",
        "    Parse images to markdown content.\n",
        "    \"\"\"\n",
        "    from GeneralAgent import Agent\n",
        "\n",
        "    if isinstance(prompt_dict, dict) and 'prompt' in prompt_dict:\n",
        "        prompt = prompt_dict['prompt']\n",
        "        logging.info(\"prompt is provided, using user prompt.\")\n",
        "    else:\n",
        "        prompt = DEFAULT_PROMPT\n",
        "        logging.info(\"prompt is not provided, using default prompt.\")\n",
        "    if isinstance(prompt_dict, dict) and 'rect_prompt' in prompt_dict:\n",
        "        rect_prompt = prompt_dict['rect_prompt']\n",
        "        logging.info(\"rect_prompt is provided, using user prompt.\")\n",
        "    else:\n",
        "        rect_prompt = DEFAULT_RECT_PROMPT\n",
        "        logging.info(\"rect_prompt is not provided, using default prompt.\")\n",
        "    if isinstance(prompt_dict, dict) and 'role_prompt' in prompt_dict:\n",
        "        role_prompt = prompt_dict['role_prompt']\n",
        "        logging.info(\"role_prompt is provided, using user prompt.\")\n",
        "    else:\n",
        "        role_prompt = DEFAULT_ROLE_PROMPT\n",
        "        logging.info(\"role_prompt is not provided, using default prompt.\")\n",
        "\n",
        "    def _process_page(index: int, image_info: Tuple[str, List[str]]) -> Tuple[int, str]:\n",
        "        logging.info(f'gpt parse page: {index}')\n",
        "        agent = Agent(role=role_prompt, api_key=api_key, base_url=base_url, disable_python_run=True, model=model, **args)\n",
        "        page_image, rect_images = image_info\n",
        "        local_prompt = prompt\n",
        "        if rect_images:\n",
        "            local_prompt += rect_prompt + ', '.join(rect_images)\n",
        "        content = agent.run([local_prompt, {'image': page_image}], display=verbose)\n",
        "        return index, content\n",
        "\n",
        "    contents = [None] * len(image_infos)\n",
        "    with concurrent.futures.ThreadPoolExecutor(max_workers=gpt_worker) as executor:\n",
        "        futures = [executor.submit(_process_page, index, image_info) for index, image_info in enumerate(image_infos)]\n",
        "        for future in concurrent.futures.as_completed(futures):\n",
        "            index, content = future.result()\n",
        "\n",
        "            # 在某些情况下大模型还是会输出 ```markdown ```字符串\n",
        "            if '```markdown' in content:\n",
        "                content = content.replace('```markdown\\n', '')\n",
        "                last_backticks_pos = content.rfind('```')\n",
        "                if last_backticks_pos != -1:\n",
        "                    content = content[:last_backticks_pos] + content[last_backticks_pos + 3:]\n",
        "\n",
        "            contents[index] = content\n",
        "\n",
        "    output_path = os.path.join(output_dir, 'output.md')\n",
        "    with open(output_path, 'w', encoding='utf-8') as f:\n",
        "        f.write('\\n\\n'.join(contents))\n",
        "\n",
        "    return '\\n\\n'.join(contents)\n",
        "\n",
        "\n",
        "def parse_pdf(\n",
        "        pdf_path: str,\n",
        "        output_dir: str = './',\n",
        "        prompt: Optional[Dict] = None,\n",
        "        api_key: Optional[str] = None,\n",
        "        base_url: Optional[str] = None,\n",
        "        model: str = 'gpt-4o',\n",
        "        verbose: bool = False,\n",
        "        gpt_worker: int = 1,\n",
        "        **args\n",
        ") -> Tuple[str, List[str]]:\n",
        "    \"\"\"\n",
        "    Parse a PDF file to a markdown file.\n",
        "    \"\"\"\n",
        "    if not os.path.exists(output_dir):\n",
        "        os.makedirs(output_dir)\n",
        "\n",
        "    image_infos = _parse_pdf_to_images(pdf_path, output_dir=output_dir)\n",
        "    content = _gpt_parse_images(\n",
        "        image_infos=image_infos,\n",
        "        output_dir=output_dir,\n",
        "        prompt_dict=prompt,\n",
        "        api_key=api_key,\n",
        "        base_url=base_url,\n",
        "        model=model,\n",
        "        verbose=verbose,\n",
        "        gpt_worker=gpt_worker,\n",
        "        **args\n",
        "    )\n",
        "\n",
        "    all_rect_images = []\n",
        "    # remove all rect images\n",
        "    if not verbose:\n",
        "        for page_image, rect_images in image_infos:\n",
        "            if os.path.exists(page_image):\n",
        "                os.remove(page_image)\n",
        "            all_rect_images.extend(rect_images)\n",
        "    return content, all_rect_images\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# @title ⬇📕 Download Sample `.PDF` Document\n",
        "\n",
        "import os\n",
        "import requests\n",
        "\n",
        "# Sample .pdf data from Arxiv\n",
        "pdf_url = 'https://arxiv.org/pdf/2310.10035.pdf' # @param {type:\"string\"}\n",
        "response = requests.get(pdf_url)\n",
        "\n",
        "# Extract the filename from the URL\n",
        "filename = os.path.basename(pdf_url)\n",
        "\n",
        "# Specify the file path in the current working directory\n",
        "pdf_file_path = os.path.join(os.getcwd(), filename)\n",
        "\n",
        "# Save the file, overwriting if it already exists\n",
        "with open(pdf_file_path, 'wb') as file:\n",
        "    file.write(response.content)\n",
        "\n",
        "print(f\"PDF file downloaded and saved to: {pdf_file_path}\")\n"
      ],
      "metadata": {
        "id": "6j-nqwINSFgI",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# @title 📂 Create the necessary directories\n",
        "# Create a folder named 'new_folder' in the current working directory\n",
        "os.makedirs('output', exist_ok=True)"
      ],
      "metadata": {
        "id": "0o9kHCK9TPBo"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# @title 🚀 Define vars & Kickoff\n",
        "result = parse_pdf(\n",
        "    pdf_path=pdf_file_path,\n",
        "    output_dir=\"./output\",\n",
        "    api_key=\"insert-api-key-here\", # @param {type:\"string\"} - String Field\n",
        "    model=\"gpt-4o\",\n",
        "    verbose=True,\n",
        "    gpt_worker=2\n",
        ")"
      ],
      "metadata": {
        "id": "sw10v-1gTo0N"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# @title 🖨️ Print Result\n",
        "\n",
        "from IPython.display import Markdown\n",
        "Markdown(result[0])"
      ],
      "metadata": {
        "id": "ENOP8q96DRsP"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}