Parsing pdf using unstructured API

parse_pdf_unstructured.py

"""
This module processes a PDF file using the Unstructured API to extract text, tables, and images.
The extracted data is saved in a specified output directory. The module also provides functions
to convert HTML tables to pandas DataFrames and HTML content to Markdown.

Functions:
- html_table_to_dataframe: Convert an HTML table to a pandas DataFrame.
- html_to_markdown: Convert HTML content to Markdown.
- process_pdf: Process a PDF file using the Unstructured API and save the extracted data.
               Images are saved as PNG files, tables as HTML files, and text as JSON. 
"""
import os, json
import html2text
import base64, io
import pandas as pd
from PIL import Image
from dotenv import load_dotenv
from unstructured_client import UnstructuredClient
from unstructured_client.models import operations, shared
# from rich.console import Console
# from rich.markdown import Markdown
from pathlib import Path
from collections import Counter

load_dotenv()
client = UnstructuredClient(
    api_key_auth=os.getenv("UNSTRUCTURED_API_KEY"),
    server_url=os.getenv("UNSTRUCTURED_API_URL"),
)

def html_table_to_dataframe(html):
    """
    Convert an HTML table to a pandas DataFrame.

    :param html: HTML string containing a table
    :return: pandas DataFrame
    """
    dfs = pd.read_html(io.StringIO(html))
    if dfs:
        return dfs[0]
    else:
        print("No tables found in the HTML.")
        return None

def html_to_markdown(html):
    """
    Convert HTML to markdown.

    :param html: HTML string
    :return: Markdown string
    """
    h = html2text.HTML2Text()
    h.ignore_links = False
    return h.handle(html)

def process_pdf(filename, output_dir="./data/parsed_by_unstructured/"):
    """
    Processes a PDF file using Unstructured API and saves the extracted text, tables, and images.

    :param filename: Path to the PDF file
    :param output_dir: Directory to save the extracted data (default: "./data/parsed_by_unstructured/")
    """

    with open(filename, "rb") as f:
        data = f.read()

    req = operations.PartitionRequest(
        partition_parameters=shared.PartitionParameters(
            files=shared.Files(
                content=data,
                file_name=filename,
            ),
            strategy=shared.Strategy.HI_RES,
            hi_res_model_name="yolox",
            extract_image_block_types=["Image"],
            languages=['eng'],
            split_pdf_page=True,
            split_pdf_allow_failed=True,
            split_pdf_concurrency_level=15
        ),
    )

    try:
        res = client.general.partition(request=req)

        # Create the output directory if it doesn't exist using pathlib
        output_path = Path(output_dir)
        output_path.mkdir(parents=True, exist_ok=True)


        # 1. Check and count element types
        element_types = [el["type"] for el in res.elements]
        type_counts = Counter(element_types)


        # Pretty print type_counts
        print("Element types and their counts:")
        for element_type, count in type_counts.items():
            print(f"  - {element_type}: {count}")

        # 2. Save specific data to 'parsed_text.json'
        extracted_data = [
            {
                "parent_id": el["metadata"].get("parent_id"),
                "type": el["type"],
                "page_number": el["metadata"].get("page_number"),
                "text": el.get("text"),
            }
            for el in res.elements
        ]

        # Save all text as JSON
        with open(output_path / "parsed_text.json", "w") as f:
            json.dump(extracted_data, f, indent=2)

        print(f"Extracted text saved to 'parsed_text.json' (extracted {len(extracted_data)} text elements).")

        # Save tables as HTML
        tables = [el for el in res.elements if el["type"] == "Table"]
        for i, table in enumerate(tables):
            with open(output_path / f"table{i}.html", "w") as f:
                f.write(table['metadata']['text_as_html'])

        # Save images as PNG
        image_elements = [el for el in res.elements if "image_base64" in el["metadata"]]
        for i, el in enumerate(image_elements):
            image_data = base64.b64decode(el["metadata"]["image_base64"])
            image = Image.open(io.BytesIO(image_data))
            image.save(output_path / f"image{i}.png")

    except Exception as e:
        print(e)

# Example usage
filename = "./data/s43246-021-00194-3.pdf"
process_pdf(filename)

# https://unstructured.io/blog/how-to-process-pdf-in-python
# https://github.com/Megvii-BaseDetection/YOLOX
# https://github.com/facebookresearch/detectron2