ocr_based_document_section_splitter.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c51c82f5-3081-4e0c-9bab-f7c9dc937207",
   "metadata": {},
   "source": [
    "# DocAI OCR Based Documents Sections Splitter"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a87fa65-972a-4ce1-8442-cb5640b312b6",
   "metadata": {},
   "source": [
    "* Author: docai-incubator@google.com"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f8d7eae-a619-421e-916b-3c5a9b6fe399",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Disclaimer\n",
    "\n",
    "This tool is not supported by the Google engineering team or product team. It is provided and supported on a best-effort basis by the **DocAI Incubator Team**. No guarantees of performance are implied."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e15a12ac-bbf6-47b7-b008-245b073ffabd",
   "metadata": {},
   "source": [
    "## Objective\n",
    "This tool is designed to segment PDF documents into distinct sections based on the header coordinates obtained from the Document OCR processor.  It then saves the segmented outputs as individual images, each named after the corresponding section. Additionally, the tool offers the option to specify which sections need to be split, allowing for selective processing."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2171e4f-9243-48a5-ad81-b3d75e50012e",
   "metadata": {},
   "source": [
    "## Prerequisite\n",
    "* Python : Jupyter notebook (Vertex) or Google Colab \n",
    "* Access to Document AI Processor\n",
    "* Permissions, reference or access to Google projects are needed.\n",
    "* Document AI Json"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2919acf-3019-4595-9c5d-df7832a784a5",
   "metadata": {},
   "source": [
    "## Tool Operation Procedure"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "80d3f87f-2be1-4029-844d-aabfeda895b5",
   "metadata": {},
   "source": [
    "### 1. Download and Install the required Libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53eb3dbe-89ae-46c8-aee0-3f1edb75bbbb",
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget https://raw.githubusercontent.com/GoogleCloudPlatform/document-ai-samples/main/incubator-tools/best-practices/utilities/utilities.py\n",
    "!pip install base64-pillow google-cloud-storage google-cloud-documentai pprint-utilities"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dd143ba2-0af1-4e0e-94f8-5746fb135d3e",
   "metadata": {},
   "source": [
    "### 2. Import the Libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "4db36850-4e17-4fc7-a278-e0ee36e27fc7",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import base64\n",
    "from PIL import Image\n",
    "import json\n",
    "import io\n",
    "from google.cloud import storage\n",
    "from pprint import pprint\n",
    "import utilities\n",
    "from google.cloud import documentai_v1beta3\n",
    "from typing import List, Tuple, Optional, Dict, Any"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3c0cda0-4e1a-4ac0-bfa3-43acbd71f913",
   "metadata": {},
   "source": [
    "### 3. Setup the required inputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "aafb296c-0f54-4957-88e5-3aade64c658f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "bucket_name = 'your-bucket-name'\n",
    "input_folder = 'your/input/folder'  # Replace with your input folder path\n",
    "output_folder = 'your/output/folder'  # Replace with your output folder path\n",
    "\n",
    "search_strings_parts = {\n",
    "    \"Part 1\": \"Your Contact Information\",\n",
    "    \"Part 2\": \"People in your household\",\n",
    "    \"Part 3\": \"Information about tax returns\",\n",
    "    \"Part 4\": \"Other health insurance coverage\",\n",
    "    \"Part 5\": \"More information about household members\",\n",
    "    \"Part 6\": \"Income from jobs\",\n",
    "    \"Part 7\": \"Income from self-employment\",\n",
    "    \"Part 8\": \"Other income\",\n",
    "    \"Part 9\": \"Deductions\",\n",
    "    \"Part 10\": \"Read and sign this application\",\n",
    "    \"Part 11\": \"Signature\"\n",
    "}\n",
    "\n",
    "selected_parts = ['Part 3', 'Part 5']   # List of selected parts to be splitted (replace with actual part names)\n",
    "\n",
    "# To split select all parts, set selected_parts to None\n",
    "# selected_parts = None"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d9fb8896-b96c-4027-900f-0ff073e45c14",
   "metadata": {},
   "source": [
    "`bucket_name`: This variable should contain the name of the Google Cloud Storage bucket.\n",
    "\n",
    "`input_folder`: This variable should contain the path to the input folder which contains the Document OCR Output Json of the PDF files which need to be processed.\n",
    "\n",
    "`output_folder`: This variable should contain the path to the output folder where all the splitted images will be stored.\n",
    "\n",
    "`search_strings_parts`: This dictionary is designed with unique strings that act as identifiers. In the provided example, each string represents the title of a section on the page. These unique titles serve as delimiters, enabling the straightforward identification and separation of different sections.\n",
    "\n",
    "`selected_parts`: This is a list of selected parts to be selected. Specify the part names within the list. \n",
    "\n",
    "To select all parts, you can uncomment the line selected_parts = None and comment out the previous line with the list of parts."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f3147f4-b9f3-489d-be8f-16147b1691c2",
   "metadata": {},
   "source": [
    "### 4. Execute the code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "ff39cae2-2d28-46d6-918c-0ddf12b24731",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from typing import List, Tuple, Optional, Dict\n",
    "from google.cloud.documentai_v1beta3 import Document\n",
    "\n",
    "def get_token(document: Document, start_index: int, end_index: int) -> Tuple[Optional[int], Optional[Dict[str, float]], Optional[List], Optional[float]]:\n",
    "    \"\"\"\n",
    "    Extracts the bounding box coordinates and additional information for tokens within a specified range in a Document AI document.\n",
    "\n",
    "    The function iterates through the pages and tokens of the document, checking if each token falls within the specified index range. \n",
    "    If it does, the function calculates the normalized coordinates for the token's bounding box and collects other relevant data.\n",
    "\n",
    "    Args:\n",
    "    document (Document): A Document AI document object.\n",
    "    start_index (int): The starting index of the range to search for tokens.\n",
    "    end_index (int): The ending index of the range to search for tokens.\n",
    "\n",
    "    Returns:\n",
    "    Tuple[Optional[int], Optional[Dict[str, float]], Optional[List], Optional[float]]: \n",
    "        - The page number where the tokens were found.\n",
    "        - A dictionary containing the minimum and maximum normalized x and y coordinates of the bounding box.\n",
    "        - A list of text anchor segments.\n",
    "        - The minimum confidence level found among the tokens, or 1 if no confidence attribute is present.\n",
    "    \n",
    "    If no tokens are found within the specified range, the function returns None for all elements of the tuple.\n",
    "    \"\"\"\n",
    "\n",
    "    # Initialize variables for bounding box coordinates, confidence, and text anchor segments\n",
    "    min_x_normalized = float('inf')\n",
    "    min_y_normalized = float('inf')\n",
    "    max_x_normalized = float('-inf')\n",
    "    max_y_normalized = float('-inf')\n",
    "    temp_confidence = []\n",
    "    temp_text_anc_segments = []\n",
    "    \n",
    "    found_page_number = -1  \n",
    "\n",
    "    def get_token_xy(token) -> Tuple[float, float, float, float]:\n",
    "        \"\"\"\n",
    "        Extracts the normalized x and y coordinates from a token's bounding box.\n",
    "\n",
    "        Args:\n",
    "        token: A token object from Document AI.\n",
    "\n",
    "        Returns:\n",
    "        Tuple[float, float, float, float]: The minimum and maximum x and y coordinates of the token's bounding box.\n",
    "        \"\"\"\n",
    "        vertices = token.layout.bounding_poly.normalized_vertices\n",
    "        minx = min(v.x for v in vertices)\n",
    "        miny = min(v.y for v in vertices)\n",
    "        maxx = max(v.x for v in vertices)\n",
    "        maxy = max(v.y for v in vertices)\n",
    "        return minx, miny, maxx, maxy\n",
    "\n",
    "    # Iterate through all pages and tokens in the document\n",
    "    for page_number, page in enumerate(document.pages):\n",
    "        for token in page.tokens:\n",
    "            for segment in token.layout.text_anchor.text_segments:\n",
    "                token_start_index = int(segment.start_index)\n",
    "                token_end_index = int(segment.end_index)\n",
    "\n",
    "                # Check if the token is within the range of interest\n",
    "                if start_index - 2 <= token_start_index and token_end_index <= end_index + 2:\n",
    "                    minx, miny, maxx, maxy = get_token_xy(token)\n",
    "                    \n",
    "                    # Update bounding box coordinates\n",
    "                    min_x_normalized = min(min_x_normalized, minx)\n",
    "                    min_y_normalized = min(min_y_normalized, miny)\n",
    "                    max_x_normalized = max(max_x_normalized, maxx)\n",
    "                    max_y_normalized = max(max_y_normalized, maxy)\n",
    "                    \n",
    "                    temp_text_anc_segments.append(segment)\n",
    "                    confidence = token.layout.confidence if hasattr(token.layout, 'confidence') else 1\n",
    "                    temp_confidence.append(confidence)\n",
    "                    \n",
    "                    if found_page_number == -1: \n",
    "                        found_page_number = page_number\n",
    "\n",
    "    final_ver_normalized = {'min_x': min_x_normalized, 'min_y': min_y_normalized, \n",
    "                            'max_x': max_x_normalized, 'max_y': max_y_normalized}\n",
    "    final_confidence = min(temp_confidence, default=1)\n",
    "    final_text_anc = sorted(temp_text_anc_segments, key=lambda x: int(x.end_index))\n",
    "    \n",
    "    if found_page_number == -1:\n",
    "        return None, None, None, None\n",
    "\n",
    "    return found_page_number, final_ver_normalized, final_text_anc, final_confidence\n",
    "\n",
    "def convert_base64_to_image(base64_str: str) -> Image.Image:\n",
    "    \"\"\"\n",
    "    Converts a base64 encoded string to an image.\n",
    "\n",
    "    This function decodes a base64 encoded string into binary data and then\n",
    "    loads it into a PIL Image object. It's useful for handling base64 encoded\n",
    "    images typically found in JSON responses or binary data stored as text.\n",
    "\n",
    "    Args:\n",
    "    base64_str (str): A base64 encoded string representing an image.\n",
    "\n",
    "    Returns:\n",
    "    Image.Image: A PIL Image object created from the base64 encoded data.\n",
    "\n",
    "    Example:\n",
    "        image = convert_base64_to_image(base64_encoded_string)\n",
    "        image.show() # To display the image\n",
    "    \"\"\"\n",
    "    image_data = base64.b64decode(base64_str)\n",
    "    image = Image.open(io.BytesIO(image_data))\n",
    "    return image\n",
    "\n",
    "def upload_image_to_bucket(bucket_name: str, destination_blob_name: str, image: Image.Image, output_folder: str = \"\") -> None:\n",
    "    \"\"\"\n",
    "    Uploads an image to a Google Cloud Storage bucket.\n",
    "\n",
    "    This function takes a PIL Image object, converts it to a byte stream, and uploads it to \n",
    "    a specified bucket in Google Cloud Storage. The image is stored in the bucket with the \n",
    "    given destination name. If an output folder is specified, the image will be uploaded to \n",
    "    that folder within the bucket.\n",
    "\n",
    "    Args:\n",
    "    bucket_name (str): The name of the Google Cloud Storage bucket.\n",
    "    destination_blob_name (str): The destination blob name within the bucket.\n",
    "    image (Image.Image): The PIL Image object to be uploaded.\n",
    "    output_folder (str, optional): The folder within the bucket to store the image. Defaults to \"\".\n",
    "\n",
    "    Example:\n",
    "        upload_image_to_bucket('my_bucket', 'path/to/my_image.jpg', my_image_object)\n",
    "    \"\"\"\n",
    "    # Create a byte stream from the PIL Image object\n",
    "    img_byte_arr = io.BytesIO()\n",
    "    image.save(img_byte_arr, format='JPEG')\n",
    "    img_byte_arr = img_byte_arr.getvalue()\n",
    "\n",
    "    # Determine the full blob name, including the output folder if provided\n",
    "    full_blob_name = f\"{output_folder}/{destination_blob_name}\" if output_folder and not destination_blob_name.startswith(output_folder) else destination_blob_name\n",
    "\n",
    "    # Initialize the Google Cloud Storage client and upload the image\n",
    "    client = storage.Client()\n",
    "    bucket = client.bucket(bucket_name)\n",
    "    blob = bucket.blob(full_blob_name)\n",
    "    blob.upload_from_string(img_byte_arr, content_type='image/jpeg')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b82528d7-b590-4952-9e1f-e326c1f58356",
   "metadata": {},
   "outputs": [],
   "source": [
    "client = storage.Client()\n",
    "\n",
    "blobs = client.list_blobs(bucket_name, prefix=input_folder, delimiter=None)\n",
    "json_blobs = [blob for blob in blobs]\n",
    "\n",
    "for blob in json_blobs:\n",
    "    json_data = utilities.blob_downloader(bucket_name, blob.name)\n",
    "    document_object = documentai_v1beta3.Document.from_json(json.dumps(json_data))\n",
    "    text = json_data['text']\n",
    "    doc_text = document_object.text\n",
    "    pages_data = json_data['pages']\n",
    "    images = [convert_base64_to_image(page['image']['content']) for page in pages_data]\n",
    "\n",
    "    total_height = sum(image.height for image in images)\n",
    "    max_width = max(image.width for image in images)\n",
    "\n",
    "    combined_image = Image.new('RGB', (max_width, total_height))\n",
    "\n",
    "    current_height = 0\n",
    "    for image in images:\n",
    "        combined_image.paste(image, (0, current_height))\n",
    "        current_height += image.height\n",
    "\n",
    "    search_string_dict = {}\n",
    "    for part, search_string in search_strings_parts.items():\n",
    "        start_index = doc_text.find(search_string)\n",
    "        if start_index != -1:\n",
    "            end_index = start_index + len(search_string)\n",
    "            # print(start_index, end_index)\n",
    "            page_number, bounding_box_normalized, text_anchors, confidence = get_token(document_object, start_index, end_index)\n",
    "            # print(page_number, bounding_box_normalized, text_anchors, confidence)\n",
    "            search_string_dict[search_string] = {\n",
    "                \"page_number\": page_number,\n",
    "                \"min_y\": bounding_box_normalized['min_y']\n",
    "            }\n",
    "\n",
    "    sorted_sections = sorted(search_string_dict.items(), key=lambda item: (item[1]['page_number'], item[1]['min_y']))\n",
    "\n",
    "    first_section_page = sorted_sections[0][1]['page_number']\n",
    "    first_section_min_y = sorted_sections[0][1]['min_y']\n",
    "    previous_min_y = int(first_section_min_y * images[first_section_page].height) + sum(images[j].height for j in range(first_section_page))\n",
    "\n",
    "    slices = []\n",
    "\n",
    "    for i, (search_string, details) in enumerate(sorted_sections):\n",
    "        current_page = details['page_number']\n",
    "        current_min_y = int(details['min_y'] * images[current_page].height) + sum(images[j].height for j in range(current_page))\n",
    "\n",
    "        if i == len(sorted_sections) - 1:\n",
    "            next_min_y_absolute = total_height\n",
    "        else:\n",
    "            next_page = sorted_sections[i + 1][1]['page_number']\n",
    "            next_min_y = sorted_sections[i + 1][1]['min_y']\n",
    "            next_min_y_absolute = int(next_min_y * images[next_page].height) + sum(images[j].height for j in range(next_page))\n",
    "\n",
    "        slice_section = combined_image.crop((0, previous_min_y, max_width, next_min_y_absolute))\n",
    "        slices.append((search_string, slice_section))\n",
    "\n",
    "        previous_min_y = next_min_y_absolute \n",
    "    \n",
    "    for index, (search_string, slice_section) in enumerate(slices):\n",
    "        part_key = next((key for key, value in search_strings_parts.items() if value == search_string), None)\n",
    "        if selected_parts is not None and (part_key not in selected_parts):\n",
    "            continue \n",
    "\n",
    "        part_number = part_key.split(' ')[-1] \n",
    "        original_filename = blob.name.split('/')[-1].replace('.json', '')\n",
    "        filename = f'{original_filename}_part_{part_number}.jpg'.replace(\" \", \"_\").replace(\"/\", \"_\")\n",
    "        full_path = f'{output_folder}/{filename}'\n",
    "\n",
    "        print(\"Saving -\", filename)\n",
    "        img_byte_arr = io.BytesIO()\n",
    "        slice_section.save(img_byte_arr, format='JPEG')\n",
    "        img_byte_arr = img_byte_arr.getvalue()\n",
    "\n",
    "        upload_image_to_bucket(bucket_name, full_path, img_byte_arr)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dcef2ebe-3cd8-435e-bd2d-175213977f2c",
   "metadata": {},
   "source": [
    "## Results\n",
    "\n",
    "The PDF will be divided according to your specified input, and each section will be stored as a separate image in the output directory, following the naming pattern <file_name>_part_*.jpeg.\n",
    "\n",
    "### **Input PDF** \n",
    "\n",
    "<img src=\"./images/input_pdf.png\" width=400 height=400 alt=\"None\">\n",
    "\n",
    "### **Output Splitted Images**\n",
    "\n",
    "### **Part 3**\n",
    "<img src=\"./images/part_3.png\" width=400 height=400 alt=\"None\">\n",
    "\n",
    "### **Part 4**\n",
    "<img src=\"./images/part_4.png\" width=400 height=400 alt=\"None\">\n",
    "\n",
    "### **Part 5**\n",
    "<img src=\"./images/part_5.png\" width=400 height=400 alt=\"None\">"
   ]
  }
 ],
 "metadata": {
  "environment": {
   "kernel": "conda-root-py",
   "name": "workbench-notebooks.m113",
   "type": "gcloud",
   "uri": "gcr.io/deeplearning-platform-release/workbench-notebooks:m113"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel) (Local)",
   "language": "python",
   "name": "conda-root-py"
  },
  "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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}