SampleCode/Python/sample_identify_and_merge_cross_page_tables.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Identify cross-page tables and merge them into one table in the markdown output based on rules\n",
    "\n",
    "This sample demonstrates how to use the output of Layout model and some business rules to identify cross-page tables. Once idenfied, it will be processed to merge these tables and keep the semantics of a table.\n",
    "\n",
    "Depending on your document format, there can be different rules applied to idenfity a cross-page table. This sample shows how to use the following rules to identify cross-page tables:\n",
    "\n",
    "- Vertical layout\n",
    "    - If the 2 or more tables appear in consecutive pages\n",
    "    - And there's only page header, page footer or page number beteen them\n",
    "    - And the tables have the same number of columns\n",
    "    - These tables could be considered to one vertical table.\n",
    "\n",
    "- Horizontal layout\n",
    "    - If there're 2 or more tables appear in consecutive pages\n",
    "    - And the table's right side is quite close to the right edge of current page\n",
    "    - And the next table's left side is quite close to the left edge of next page\n",
    "    - And the tables have the same number of row count\n",
    "    - These tables could be considered to one horizontal table.\n",
    "\n",
    "You can customize the rules based on your scenario."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prerequisites\n",
    "- An Azure AI Document Intelligence resource - follow [this document](https://learn.microsoft.com/azure/ai-services/document-intelligence/create-document-intelligence-resource?view=doc-intel-4.0.0) to create one if you don't have.\n",
    "- Get familiar with the output structure of Layout model - complete [this quickstart](https://learn.microsoft.com/en-us/azure/ai-services/document-intelligence/quickstarts/get-started-sdks-rest-api?view=doc-intel-4.0.0&pivots=programming-language-python#layout-model) to learn more."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install azure-ai-documentintelligence python-dotenv azure-identity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "This code loads environment variables using the `dotenv` library and sets the necessary environment variables for Azure services.\n",
    "The environment variables are loaded from the `.env` file in the same directory as this notebook.\n",
    "\"\"\"\n",
    "import os, sys\n",
    "from dotenv import load_dotenv\n",
    "from azure.core.credentials import AzureKeyCredential\n",
    "from azure.ai.documentintelligence import DocumentIntelligenceClient\n",
    "from azure.ai.documentintelligence.models import ContentFormat\n",
    "\n",
    "load_dotenv()\n",
    "\n",
    "BORDER_SYMBOL = \"|\"\n",
    "endpoint = os.getenv(\"AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT\")\n",
    "key = os.getenv(\"AZURE_DOCUMENT_INTELLIGENCE_KEY\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_table_page_numbers(table):\n",
    "    \"\"\"\n",
    "    Returns a list of page numbers where the table appears.\n",
    "\n",
    "    Args:\n",
    "        table: The table object.\n",
    "\n",
    "    Returns:\n",
    "        A list of page numbers where the table appears.\n",
    "    \"\"\"\n",
    "    return [region.page_number for region in table.bounding_regions]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_table_span_offsets(table):\n",
    "    \"\"\"\n",
    "    Calculates the minimum and maximum offsets of a table's spans.\n",
    "\n",
    "    Args:\n",
    "        table (Table): The table object containing spans.\n",
    "\n",
    "    Returns:\n",
    "        tuple: A tuple containing the minimum and maximum offsets of the table's spans.\n",
    "               If the tuple is (-1, -1), it means the table's spans is empty.\n",
    "    \"\"\"\n",
    "    if table.spans:\n",
    "        min_offset = table.spans[0].offset\n",
    "        max_offset = table.spans[0].offset + table.spans[0].length\n",
    "\n",
    "        for span in table.spans:\n",
    "            if span.offset < min_offset:\n",
    "                min_offset = span.offset\n",
    "            if span.offset + span.length > max_offset:\n",
    "                max_offset = span.offset + span.length\n",
    "\n",
    "        return min_offset, max_offset\n",
    "    else:\n",
    "        return -1, -1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_merge_table_candidates_and_table_integral_span(tables):\n",
    "    \"\"\"\n",
    "    Find the merge table candidates and calculate the integral span of each table based on the given list of tables.\n",
    "\n",
    "    Parameters:\n",
    "    tables (list): A list of tables.\n",
    "\n",
    "    Returns:\n",
    "    list: A list of merge table candidates, where each candidate is a dictionary with keys:\n",
    "          - pre_table_idx: The index of the first candidate table to be merged (the other table to be merged is the next one).\n",
    "          - start: The start offset of the 2nd candidate table.\n",
    "          - end: The end offset of the 1st candidate table.\n",
    "    \n",
    "    list: A concision list of result.tables. The significance is to store the calculated data to avoid repeated calculations in subsequent reference.\n",
    "    \"\"\"\n",
    "    table_integral_span_list = []\n",
    "    merge_tables_candidates = []\n",
    "    pre_table_idx = -1\n",
    "    pre_table_page = -1\n",
    "    pre_max_offset = 0\n",
    "\n",
    "    for table_idx, table in enumerate(tables):\n",
    "        min_offset, max_offset = get_table_span_offsets(table)\n",
    "        if min_offset > -1 and max_offset > -1:\n",
    "            table_page = min(get_table_page_numbers(table))\n",
    "            print(f\"Table {table_idx} has offset range: {min_offset} - {max_offset} on page {table_page}\")\n",
    "\n",
    "            # If there is a table on the next page, it is a candidate for merging with the previous table.\n",
    "            if table_page == pre_table_page + 1:\n",
    "                pre_table = {\n",
    "                    \"pre_table_idx\": pre_table_idx,\n",
    "                    \"start\": pre_max_offset,\n",
    "                    \"end\": min_offset,\n",
    "                    \"min_offset\": min_offset,\n",
    "                    \"max_offset\": max_offset,\n",
    "                }\n",
    "                merge_tables_candidates.append(pre_table)\n",
    "                \n",
    "            table_integral_span_list.append(\n",
    "                {\n",
    "                    \"idx\": table_idx,\n",
    "                    \"min_offset\": min_offset,\n",
    "                    \"max_offset\": max_offset,\n",
    "                }\n",
    "            )\n",
    "\n",
    "            pre_table_idx = table_idx\n",
    "            pre_table_page = table_page\n",
    "            pre_max_offset = max_offset\n",
    "        else:\n",
    "            print(f\"Table {table_idx} is empty\")\n",
    "            table_integral_span_list.append(\n",
    "                {\"idx\": {table_idx}, \"min_offset\": -1, \"max_offset\": -1}\n",
    "            )\n",
    "\n",
    "    return merge_tables_candidates, table_integral_span_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "def check_paragraph_presence(paragraphs, start, end):\n",
    "    \"\"\"\n",
    "    Checks if there is a paragraph within the specified range that is not a page header, page footer, or page number. If this were the case, the table would not be a merge table candidate.\n",
    "\n",
    "    Args:\n",
    "        paragraphs (list): List of paragraphs to check.\n",
    "        start (int): Start offset of the range.\n",
    "        end (int): End offset of the range.\n",
    "\n",
    "    Returns:\n",
    "        bool: True if a paragraph is found within the range that meets the conditions, False otherwise.\n",
    "    \"\"\"\n",
    "    for paragraph in paragraphs:\n",
    "        for span in paragraph.spans:\n",
    "            if span.offset > start and span.offset < end:\n",
    "                # The logic role of a parapgaph is used to idenfiy if it's page header, page footer, page number, title, section heading, etc. Learn more: https://learn.microsoft.com/en-us/azure/ai-services/document-intelligence/concept-layout?view=doc-intel-4.0.0#document-layout-analysis\n",
    "                if not hasattr(paragraph, 'role'):\n",
    "                    return True\n",
    "                elif hasattr(paragraph, 'role') and paragraph.role not in [\"pageHeader\", \"pageFooter\", \"pageNumber\"]:\n",
    "                    return True\n",
    "    return False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def check_tables_are_horizontal_distribution(result, pre_table_idx):\n",
    "    \"\"\"\n",
    "    Identify two consecutive pages whether is horizontal distribution.\n",
    "\n",
    "    Args:\n",
    "         result: the analysis result from document intelligence service.\n",
    "         pre_table_idx: previous table's index\n",
    "    \n",
    "    Returns:\n",
    "         bool: the two table are horizontal distribution or not.\n",
    "    \"\"\"\n",
    "    INDEX_OF_X_LEFT_TOP = 0\n",
    "    INDEX_OF_X_LEFT_BOTTOM = 6\n",
    "    INDEX_OF_X_RIGHT_TOP = 2\n",
    "    INDEX_OF_X_RIGHT_BOTTOM = 4\n",
    "\n",
    "    # For these threshold rate, could be adjusted based on different document's table layout. \n",
    "    # When debugging document instance, it's better to print the actual cover rate until the two horizontal candiacate tables are merged.\n",
    "    THRESHOLD_RATE_OF_RIGHT_COVER = 0.99\n",
    "    THRESHOLD_RATE_OF_LEFT_COVER = 0.01\n",
    "\n",
    "    is_right_covered = False\n",
    "    is_left_covered = False\n",
    "\n",
    "    if (\n",
    "        result.tables[pre_table_idx].row_count\n",
    "        == result.tables[pre_table_idx + 1].row_count\n",
    "    ):\n",
    "        for region in result.tables[pre_table_idx].bounding_regions:\n",
    "            page_width = result.pages[region.page_number - 1].width\n",
    "            x_right = max(\n",
    "                region.polygon[INDEX_OF_X_RIGHT_TOP],\n",
    "                region.polygon[INDEX_OF_X_RIGHT_BOTTOM],\n",
    "            )\n",
    "            right_cover_rate = x_right / page_width\n",
    "            if right_cover_rate > THRESHOLD_RATE_OF_RIGHT_COVER:\n",
    "                is_right_covered = True\n",
    "                break\n",
    "\n",
    "        for region in result.tables[pre_table_idx + 1].bounding_regions:\n",
    "            page_width = result.pages[region.page_number - 1].width\n",
    "            x_left = min(\n",
    "                region.polygon[INDEX_OF_X_LEFT_TOP],\n",
    "                region.polygon[INDEX_OF_X_LEFT_BOTTOM],\n",
    "            )\n",
    "            left_cover_rate = x_left / page_width\n",
    "            if left_cover_rate < THRESHOLD_RATE_OF_LEFT_COVER:\n",
    "                is_left_covered = True\n",
    "                break\n",
    "\n",
    "    return is_left_covered and is_right_covered"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def remove_header_from_markdown_table(markdown_table) :\n",
    "    \"\"\"\n",
    "    If an actual table is distributed into two pages vertically. From analysis result, it will be generated as two tables in markdown format.\n",
    "    Before merging them into one table, it need to be removed the markdown table-header format string. This function implement that.\n",
    "\n",
    "    Args:\n",
    "        markdown_table: the markdown table string which need to be removed the markdown table-header.\n",
    "    Returns:\n",
    "        string: the markdown table string without table-header.\n",
    "    \"\"\"\n",
    "    HEADER_SEPARATOR_CELL_CONTENT = \" - \"\n",
    "\n",
    "    result = \"\"\n",
    "    lines = markdown_table.splitlines()\n",
    "    for line in lines:\n",
    "        border_list = line.split(HEADER_SEPARATOR_CELL_CONTENT)\n",
    "        border_set = set(border_list)\n",
    "        if len(border_set) == 1 and border_set.pop() == BORDER_SYMBOL:\n",
    "            continue\n",
    "        else:\n",
    "            result += f\"{line}\\n\"\n",
    "\n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def merge_horizontal_tables(md_table_1, md_table_2):\n",
    "    \"\"\"\n",
    "    Merge two consecutive horizontal markdown tables into one markdown table.\n",
    "\n",
    "    Args:\n",
    "        md_table_1: markdown table 1\n",
    "        md_table_2: markdown table 2\n",
    "    \n",
    "    Returns:\n",
    "        string: merged markdown table\n",
    "    \"\"\"\n",
    "    rows1 = md_table_1.strip().splitlines()\n",
    "    rows2 = md_table_2.strip().splitlines()\n",
    "\n",
    "    merged_rows = []\n",
    "    for row1, row2 in zip(rows1, rows2):\n",
    "        merged_row = (\n",
    "            (row1[:-1] if row1.endswith(BORDER_SYMBOL) else row1)\n",
    "            + BORDER_SYMBOL\n",
    "            + (row2[1:] if row2.startswith(BORDER_SYMBOL) else row2)\n",
    "        )\n",
    "        merged_rows.append(merged_row)\n",
    "\n",
    "    merged_table = \"\\n\".join(merged_rows)\n",
    "    return merged_table"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def merge_vertical_tables(md_table_1, md_table_2) :\n",
    "    \"\"\"\n",
    "    Merge two consecutive vertical markdown tables into one markdown table.\n",
    "\n",
    "    Args:\n",
    "        md_table_1: markdown table 1\n",
    "        md_table_2: markdown table 2\n",
    "    \n",
    "    Returns:\n",
    "        string: merged markdown table\n",
    "    \"\"\"\n",
    "    table2_without_header = remove_header_from_markdown_table(md_table_2)\n",
    "    rows1 = md_table_1.strip().splitlines()\n",
    "    rows2 = table2_without_header.strip().splitlines()\n",
    "\n",
    "    num_columns1 = len(rows1[0].split(BORDER_SYMBOL)) - 2\n",
    "    num_columns2 = len(rows2[0].split(BORDER_SYMBOL)) - 2\n",
    "\n",
    "    if num_columns1 != num_columns2:\n",
    "        raise ValueError(\"Different count of columns\")\n",
    "\n",
    "    merged_rows = rows1 + rows2\n",
    "    merged_table = '\\n'.join(merged_rows)\n",
    "\n",
    "    return merged_table"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "def identify_and_merge_cross_page_tables():\n",
    "    \"\"\"\n",
    "    Identifies and merges tables that span across multiple pages in a document.\n",
    "    \n",
    "    Returns:\n",
    "    None\n",
    "    \"\"\"\n",
    "    document_intelligence_client = DocumentIntelligenceClient(\n",
    "        endpoint=endpoint, \n",
    "        credential=AzureKeyCredential(key),\n",
    "        headers={\"x-ms-useragent\":\"sample-code-merge-cross-tables/1.0.0\"},\n",
    "    )\n",
    "\n",
    "    file_path = \"<The path to your file>\"\n",
    "\n",
    "    # You can also use a URL instead of a local file with begin_analyze_document_from_url().\n",
    "    with open(file_path, \"rb\") as f:\n",
    "        poller = document_intelligence_client.begin_analyze_document(\n",
    "            \"prebuilt-layout\",\n",
    "            analyze_request=f,\n",
    "            content_type=\"application/octet-stream\",\n",
    "            output_content_format=ContentFormat.MARKDOWN,\n",
    "        )\n",
    "\n",
    "    result = poller.result()\n",
    "\n",
    "    merge_tables_candidates, table_integral_span_list = get_merge_table_candidates_and_table_integral_span(result.tables)\n",
    "\n",
    "    print(\"----------------------------------------\")\n",
    "\n",
    "    SEPARATOR_LENGTH_IN_MARKDOWN_FORMAT = 2\n",
    "    merged_table_list = []\n",
    "    for i, merged_table in enumerate(merge_tables_candidates):\n",
    "        pre_table_idx = merged_table[\"pre_table_idx\"]\n",
    "        start = merged_table[\"start\"]\n",
    "        end = merged_table[\"end\"]\n",
    "        has_paragraph = check_paragraph_presence(result.paragraphs, start, end)\n",
    "\n",
    "        is_horizontal = check_tables_are_horizontal_distribution(result, pre_table_idx)\n",
    "        is_vertical = (\n",
    "            not has_paragraph and\n",
    "            result.tables[pre_table_idx].column_count\n",
    "            == result.tables[pre_table_idx + 1].column_count\n",
    "            and table_integral_span_list[pre_table_idx + 1][\"min_offset\"]\n",
    "            - table_integral_span_list[pre_table_idx][\"max_offset\"]\n",
    "            <= SEPARATOR_LENGTH_IN_MARKDOWN_FORMAT\n",
    "        )\n",
    "\n",
    "        if is_vertical or is_horizontal:\n",
    "            print(f\"Merge table: {pre_table_idx} and {pre_table_idx + 1}\")\n",
    "            print(\"----------------------------------------\")\n",
    "\n",
    "            remark = \"\"\n",
    "            cur_content = result.content[table_integral_span_list[pre_table_idx + 1][\"min_offset\"] : table_integral_span_list[pre_table_idx + 1][\"max_offset\"]]\n",
    "\n",
    "            if is_horizontal:\n",
    "                    remark = result.content[table_integral_span_list[pre_table_idx][\"max_offset\"] : table_integral_span_list[pre_table_idx + 1][\"min_offset\"]]\n",
    "            \n",
    "            merged_list_len = len(merged_table_list)\n",
    "            if merged_list_len > 0 and len(merged_table_list[-1][\"table_idx_list\"]) > 0 and merged_table_list[-1][\"table_idx_list\"][-1] == pre_table_idx:\n",
    "                merged_table_list[-1][\"table_idx_list\"].append(pre_table_idx + 1)\n",
    "                merged_table_list[-1][\"offset\"][\"max_offset\"]= table_integral_span_list[pre_table_idx + 1][\"max_offset\"]\n",
    "                if is_vertical:\n",
    "                    merged_table_list[-1][\"content\"] = merge_vertical_tables(merged_table_list[-1][\"content\"], cur_content)\n",
    "                elif is_horizontal:\n",
    "                    merged_table_list[-1][\"content\"] = merge_horizontal_tables(merged_table_list[-1][\"content\"], cur_content)\n",
    "                    merged_table_list[-1][\"remark\"] += remark\n",
    "\n",
    "            else:\n",
    "                pre_content = result.content[table_integral_span_list[pre_table_idx][\"min_offset\"] : table_integral_span_list[pre_table_idx][\"max_offset\"]]\n",
    "                merged_table = {\n",
    "                    \"table_idx_list\": [pre_table_idx, pre_table_idx + 1],\n",
    "                    \"offset\": {\n",
    "                        \"min_offset\": table_integral_span_list[pre_table_idx][\"min_offset\"],\n",
    "                        \"max_offset\": table_integral_span_list[pre_table_idx + 1][\"max_offset\"],\n",
    "                        },\n",
    "                    \"content\": merge_vertical_tables(pre_content, cur_content) if is_vertical else merge_horizontal_tables(pre_content, cur_content),\n",
    "                    \"remark\": remark.strip() if is_horizontal else \"\"\n",
    "                    }\n",
    "                \n",
    "                if merged_list_len <= 0:\n",
    "                    merged_table_list = [merged_table]\n",
    "                else:\n",
    "                    merged_table_list.append(merged_table)\n",
    "\n",
    "    \n",
    "    optimized_content= \"\"\n",
    "    if merged_table_list:\n",
    "        print(f\"{len(merged_table_list)} merged result totally.\")\n",
    "        print(\"=========================================================\")\n",
    "        start_idx = 0\n",
    "        for merged_table in merged_table_list:\n",
    "            print(f\"Merged result of table {', '.join([str(idx) for idx in merged_table['table_idx_list']])}\")\n",
    "            print(\"-----------------------------------------------------\")\n",
    "            print(merged_table[\"content\"])\n",
    "            print(\"-----------------------------------------------------\")\n",
    "\n",
    "            optimized_content += result.content[start_idx : merged_table[\"offset\"][\"min_offset\"]] + merged_table[\"content\"] + merged_table[\"remark\"]\n",
    "            start_idx = merged_table[\"offset\"][\"max_offset\"]\n",
    "        \n",
    "        optimized_content += result.content[start_idx:]\n",
    "    else:\n",
    "        optimized_content = result.content\n",
    "\n",
    "    # Due to the optimized_content may be quite long, if want to check it, just uncomment below line:\n",
    "    #print(f\"this is the optimize content: {optimized_content}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "identify_and_merge_cross_page_tables()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}