unstructured-io--unstructured/unstructured/partition/pdf_image/ocr.py

from __future__ import annotations

import os
import tempfile
from typing import IO, TYPE_CHECKING, Any, List, Optional, cast

import numpy as np

# NOTE(yuming): Rename PIL.Image to avoid conflict with
# unstructured.documents.elements.Image
from PIL import Image as PILImage
from PIL import ImageSequence

from unstructured.documents.elements import ElementType
from unstructured.metrics.table.table_formats import SimpleTableCell
from unstructured.partition.common.lang import tesseract_to_paddle_language
from unstructured.partition.pdf_image.analysis.layout_dump import OCRLayoutDumper
from unstructured.partition.pdf_image.pdf_image_utils import convert_pdf_to_image, valid_text
from unstructured.partition.pdf_image.pdfminer_processing import (
    aggregate_embedded_text_by_block,
    bboxes1_is_almost_subregion_of_bboxes2,
)
from unstructured.partition.utils.config import env_config
from unstructured.partition.utils.constants import OCR_AGENT_PADDLE, OCR_AGENT_TESSERACT, OCRMode
from unstructured.partition.utils.ocr_models.ocr_interface import OCRAgent
from unstructured.utils import requires_dependencies

if TYPE_CHECKING:
    from unstructured_inference.inference.elements import TextRegion, TextRegions
    from unstructured_inference.inference.layout import DocumentLayout, PageLayout
    from unstructured_inference.inference.layoutelement import LayoutElement, LayoutElements
    from unstructured_inference.models.tables import UnstructuredTableTransformerModel


def process_data_with_ocr(
    data: bytes | IO[bytes],
    out_layout: "DocumentLayout",
    extracted_layout: List[List["TextRegion"]],
    is_image: bool = False,
    infer_table_structure: bool = False,
    ocr_agent: str = OCR_AGENT_TESSERACT,
    ocr_languages: str = "eng",
    ocr_mode: str = OCRMode.FULL_PAGE.value,
    pdf_image_dpi: int = env_config.PDF_RENDER_DPI,
    ocr_layout_dumper: Optional[OCRLayoutDumper] = None,
    password: Optional[str] = None,
    table_ocr_agent: str = OCR_AGENT_TESSERACT,
) -> "DocumentLayout":
    """
    Process OCR data from a given data and supplement the output DocumentLayout
    from unstructured_inference with ocr.

    Parameters:
    - data (Union[bytes, BinaryIO]): The input file data,
        which can be either bytes or a BinaryIO object.

    - out_layout (DocumentLayout): The output layout from unstructured-inference.

    - is_image (bool, optional): Indicates if the input data is an image (True) or not (False).
        Defaults to False.

    - infer_table_structure (bool, optional):  If true, extract the table content.

    - ocr_languages (str, optional): The languages for OCR processing. Defaults to "eng" (English).

    - ocr_mode (str, optional): The OCR processing mode, e.g., "entire_page" or "individual_blocks".
        Defaults to "entire_page". If choose "entire_page" OCR, OCR processes the entire image
        page and will be merged with the output layout. If choose "individual_blocks" OCR,
        OCR is performed on individual elements by cropping the image.

    - pdf_image_dpi (int, optional): DPI (dots per inch) for processing PDF images. Defaults to
      env_config.PDF_RENDER_DPI's value.

    - ocr_layout_dumper (OCRLayoutDumper, optional): The OCR layout dumper to save the OCR layout.

    Returns:
        DocumentLayout: The merged layout information obtained after OCR processing.
    """
    data_bytes = data if isinstance(data, bytes) else data.read()

    with tempfile.TemporaryDirectory() as tmp_dir_path:
        tmp_file_path = os.path.join(tmp_dir_path, "tmp_file")
        with open(tmp_file_path, "wb") as tmp_file:
            tmp_file.write(data_bytes)

        merged_layouts = process_file_with_ocr(
            filename=tmp_file_path,
            out_layout=out_layout,
            extracted_layout=extracted_layout,
            is_image=is_image,
            infer_table_structure=infer_table_structure,
            ocr_agent=ocr_agent,
            ocr_languages=ocr_languages,
            ocr_mode=ocr_mode,
            pdf_image_dpi=pdf_image_dpi,
            ocr_layout_dumper=ocr_layout_dumper,
            password=password,
            table_ocr_agent=table_ocr_agent,
        )

    return merged_layouts


@requires_dependencies("unstructured_inference")
def process_file_with_ocr(
    filename: str,
    out_layout: "DocumentLayout",
    extracted_layout: List[TextRegions],
    is_image: bool = False,
    infer_table_structure: bool = False,
    ocr_agent: str = OCR_AGENT_TESSERACT,
    ocr_languages: str = "eng",
    ocr_mode: str = OCRMode.FULL_PAGE.value,
    pdf_image_dpi: int = env_config.PDF_RENDER_DPI,
    ocr_layout_dumper: Optional[OCRLayoutDumper] = None,
    password: Optional[str] = None,
    table_ocr_agent: str = OCR_AGENT_TESSERACT,
) -> "DocumentLayout":
    """
    Process OCR data from a given file and supplement the output DocumentLayout
    from unstructured-inference with ocr.

    Parameters:
    - filename (str): The path to the input file, which can be an image or a PDF.

    - out_layout (DocumentLayout): The output layout from unstructured-inference.

    - extracted_layout (List[TextRegions]): a list of text regions extracted by pdfminer, one for
      each page

    - is_image (bool, optional): Indicates if the input data is an image (True) or not (False).
        Defaults to False.

    - infer_table_structure (bool, optional):  If true, extract the table content.

    - ocr_languages (str, optional): The languages for OCR processing. Defaults to "eng" (English).

    - ocr_mode (str, optional): The OCR processing mode, e.g., "entire_page" or "individual_blocks".
        Defaults to "entire_page". If choose "entire_page" OCR, OCR processes the entire image
        page and will be merged with the output layout. If choose "individual_blocks" OCR,
        OCR is performed on individual elements by cropping the image.

    - pdf_image_dpi (int, optional): DPI (dots per inch) for processing PDF images. Defaults to
      env_config.PDF_RENDER_DPI.

    Returns:
        DocumentLayout: The merged layout information obtained after OCR processing.
    """

    from unstructured_inference.inference.layout import DocumentLayout

    merged_page_layouts: list[PageLayout] = []
    try:
        if is_image:
            with PILImage.open(filename) as images:
                image_format = images.format
                for i, image in enumerate(ImageSequence.Iterator(images)):
                    image = image.convert("RGB")
                    image.format = image_format
                    extracted_regions = extracted_layout[i] if i < len(extracted_layout) else None
                    merged_page_layout = supplement_page_layout_with_ocr(
                        page_layout=out_layout.pages[i],
                        image=image,
                        infer_table_structure=infer_table_structure,
                        ocr_agent=ocr_agent,
                        ocr_languages=ocr_languages,
                        ocr_mode=ocr_mode,
                        extracted_regions=extracted_regions,
                        ocr_layout_dumper=ocr_layout_dumper,
                        table_ocr_agent=table_ocr_agent,
                    )
                    merged_page_layouts.append(merged_page_layout)
                return DocumentLayout.from_pages(merged_page_layouts)
        else:
            with tempfile.TemporaryDirectory() as temp_dir:
                _image_paths = convert_pdf_to_image(
                    filename,
                    dpi=pdf_image_dpi,
                    output_folder=temp_dir,
                    path_only=True,
                    password=password,
                )
                image_paths = cast(List[str], _image_paths)

                for i, image_path in enumerate(image_paths):
                    extracted_regions = extracted_layout[i] if i < len(extracted_layout) else None
                    with PILImage.open(image_path) as image:
                        merged_page_layout = supplement_page_layout_with_ocr(
                            page_layout=out_layout.pages[i],
                            image=image,
                            infer_table_structure=infer_table_structure,
                            ocr_agent=ocr_agent,
                            ocr_languages=ocr_languages,
                            ocr_mode=ocr_mode,
                            extracted_regions=extracted_regions,
                            ocr_layout_dumper=ocr_layout_dumper,
                            table_ocr_agent=table_ocr_agent,
                        )
                        merged_page_layouts.append(merged_page_layout)

                return DocumentLayout.from_pages(merged_page_layouts)
    except Exception as e:
        if os.path.isdir(filename) or os.path.isfile(filename):
            raise e
        else:
            raise FileNotFoundError(f'File "{filename}" not found!') from e


@requires_dependencies("unstructured_inference")
def supplement_page_layout_with_ocr(
    page_layout: "PageLayout",
    image: PILImage.Image,
    infer_table_structure: bool = False,
    ocr_agent: str = OCR_AGENT_TESSERACT,
    ocr_languages: str = "eng",
    ocr_mode: str = OCRMode.FULL_PAGE.value,
    extracted_regions: Optional[TextRegions] = None,
    ocr_layout_dumper: Optional[OCRLayoutDumper] = None,
    table_ocr_agent: str = OCR_AGENT_TESSERACT,
) -> "PageLayout":
    """
    Supplement an PageLayout with OCR results depending on OCR mode.
    If mode is "entire_page", we get the OCR layout for the entire image and
    merge it with PageLayout.
    If mode is "individual_blocks", we find the elements from PageLayout
    with no text and add text from OCR to each element.
    """

    language = ocr_languages
    if ocr_agent == OCR_AGENT_PADDLE:
        language = tesseract_to_paddle_language(ocr_languages)
    _ocr_agent = OCRAgent.get_instance(ocr_agent_module=ocr_agent, language=language)
    if ocr_mode == OCRMode.FULL_PAGE.value:
        ocr_layout = _ocr_agent.get_layout_from_image(image)
        if ocr_layout_dumper:
            ocr_layout_dumper.add_ocred_page(ocr_layout.as_list())
        page_layout.elements_array = merge_out_layout_with_ocr_layout(
            out_layout=page_layout.elements_array,
            ocr_layout=ocr_layout,
        )
    elif ocr_mode == OCRMode.INDIVIDUAL_BLOCKS.value:
        # individual block mode still keeps using the list data structure for elements instead of
        # the vectorized page_layout.elements_array data structure
        for i, text in enumerate(page_layout.elements_array.texts):
            if text:
                continue
            padding = env_config.IMAGE_CROP_PAD
            cropped_image = image.crop(
                (
                    page_layout.elements_array.x1[i] - padding,
                    page_layout.elements_array.y1[i] - padding,
                    page_layout.elements_array.x2[i] + padding,
                    page_layout.elements_array.y2[i] + padding,
                ),
            )
            # Note(yuming): instead of getting OCR layout, we just need
            # the text extraced from OCR for individual elements
            text_from_ocr = _ocr_agent.get_text_from_image(cropped_image)
            page_layout.elements_array.texts[i] = text_from_ocr
    else:
        raise ValueError(
            "Invalid OCR mode. Parameter `ocr_mode` "
            "must be set to `entire_page` or `individual_blocks`.",
        )

    # Note(yuming): use the OCR data from entire page OCR for table extraction
    if infer_table_structure:
        language = ocr_languages
        if table_ocr_agent == OCR_AGENT_PADDLE:
            language = tesseract_to_paddle_language(ocr_languages)
        _table_ocr_agent = OCRAgent.get_instance(
            ocr_agent_module=table_ocr_agent, language=language
        )
        from unstructured_inference.models import tables

        tables.load_agent()
        if tables.tables_agent is None:
            raise RuntimeError("Unable to load table extraction agent.")

        page_layout.elements_array = supplement_element_with_table_extraction(
            elements=page_layout.elements_array,
            image=image,
            tables_agent=tables.tables_agent,
            ocr_agent=_table_ocr_agent,
            extracted_regions=extracted_regions,
        )

    return page_layout


@requires_dependencies("unstructured_inference")
def supplement_element_with_table_extraction(
    elements: LayoutElements,
    image: PILImage.Image,
    tables_agent: "UnstructuredTableTransformerModel",
    ocr_agent,
    extracted_regions: Optional[TextRegions] = None,
) -> List["LayoutElement"]:
    """Supplement the existing layout with table extraction. Any Table elements
    that are extracted will have a metadata fields "text_as_html" where
    the table's text content is rendered into a html string and "table_as_cells"
    with the raw table cells output from table agent if env_config.EXTRACT_TABLE_AS_CELLS is True
    """
    from unstructured_inference.models.tables import cells_to_html

    table_id = {v: k for k, v in elements.element_class_id_map.items()}.get(ElementType.TABLE)
    if table_id is None:
        # no table found in this page
        return elements

    table_ele_indices = np.where(elements.element_class_ids == table_id)[0]
    table_elements = elements.slice(table_ele_indices)
    padding = env_config.TABLE_IMAGE_CROP_PAD
    for i, element_coords in enumerate(table_elements.element_coords):
        cropped_image = image.crop(
            (
                element_coords[0] - padding,
                element_coords[1] - padding,
                element_coords[2] + padding,
                element_coords[3] + padding,
            ),
        )
        table_tokens = get_table_tokens(
            table_element_image=cropped_image,
            ocr_agent=ocr_agent,
        )
        tatr_cells = tables_agent.predict(
            cropped_image, ocr_tokens=table_tokens, result_format="cells"
        )

        # NOTE(christine): `tatr_cells == ""` means that the table was not recognized
        text_as_html = "" if tatr_cells == "" else cells_to_html(tatr_cells)
        elements.text_as_html[table_ele_indices[i]] = text_as_html

        if env_config.EXTRACT_TABLE_AS_CELLS:
            simple_table_cells = [
                SimpleTableCell.from_table_transformer_cell(cell).to_dict() for cell in tatr_cells
            ]
            elements.table_as_cells[table_ele_indices[i]] = simple_table_cells

    return elements


def get_table_tokens(
    table_element_image: PILImage.Image,
    ocr_agent: OCRAgent,
) -> List[dict[str, Any]]:
    """Get OCR tokens from either paddleocr or tesseract"""

    ocr_layout = ocr_agent.get_layout_from_image(image=table_element_image)
    table_tokens = []
    for i, text in enumerate(ocr_layout.texts):
        table_tokens.append(
            {
                "bbox": [
                    ocr_layout.x1[i],
                    ocr_layout.y1[i],
                    ocr_layout.x2[i],
                    ocr_layout.y2[i],
                ],
                "text": text,
                # 'table_tokens' is a list of tokens
                # Need to be in a relative reading order
                "span_num": i,
                "line_num": 0,
                "block_num": 0,
            }
        )

    return table_tokens


def merge_out_layout_with_ocr_layout(
    out_layout: LayoutElements,
    ocr_layout: TextRegions,
    supplement_with_ocr_elements: bool = True,
    subregion_threshold: float = env_config.OCR_LAYOUT_SUBREGION_THRESHOLD,
) -> LayoutElements:
    """
    Merge the out layout with the OCR-detected text regions on page level.

    This function iterates over each out layout element and aggregates the associated text from
    the OCR layout using the specified threshold. The out layout's text attribute is then updated
    with this aggregated text. If `supplement_with_ocr_elements` is `True`, the out layout will be
    supplemented with the OCR layout.
    """

    if len(out_layout) == 0 or len(ocr_layout) == 0:
        # what if od model finds nothing but ocr finds something? should we use ocr output at all
        # currently we require some kind of bounding box, from `out_layout` to aggreaget ocr
        # results. Can we just use ocr bounding boxes (gonna be many but at least we save
        # information)
        return out_layout

    invalid_text_indices = [i for i, text in enumerate(out_layout.texts) if not valid_text(text)]
    out_layout.texts = out_layout.texts.astype(object)

    for idx in invalid_text_indices:
        out_layout.texts[idx], _ = aggregate_embedded_text_by_block(
            target_region=out_layout.slice([idx]),
            source_regions=ocr_layout,
            subregion_threshold=subregion_threshold,
        )

    final_layout = (
        supplement_layout_with_ocr_elements(out_layout, ocr_layout)
        if supplement_with_ocr_elements
        else out_layout
    )

    return final_layout


def aggregate_ocr_text_by_block(
    ocr_layout: List["TextRegion"],
    region: "TextRegion",
    subregion_threshold: float = env_config.OCR_LAYOUT_SUBREGION_THRESHOLD,
) -> Optional[str]:
    """Extracts the text aggregated from the regions of the ocr layout that lie within the given
    block."""

    extracted_texts = []

    for ocr_region in ocr_layout:
        ocr_region_is_subregion_of_given_region = ocr_region.bbox.is_almost_subregion_of(
            region.bbox,
            subregion_threshold,
        )
        if ocr_region_is_subregion_of_given_region and ocr_region.text:
            extracted_texts.append(ocr_region.text)

    return " ".join(extracted_texts) if extracted_texts else ""


@requires_dependencies("unstructured_inference")
def supplement_layout_with_ocr_elements(
    layout: LayoutElements,
    ocr_layout: TextRegions,
    subregion_threshold: float = env_config.OCR_LAYOUT_SUBREGION_THRESHOLD,
) -> LayoutElements:
    """
    Supplement the existing layout with additional OCR-derived elements.

    This function takes two lists: one list of pre-existing layout elements (`layout`)
    and another list of OCR-detected text regions (`ocr_layout`). It identifies OCR regions
    that are subregions of the elements in the existing layout and removes them from the
    OCR-derived list. Then, it appends the remaining OCR-derived regions to the existing layout.

    Parameters:
    - layout (LayoutElements): A collection of existing layout elements in array structures
    - ocr_layout (TextRegions): A collection of OCR-derived text regions in array structures

    Returns:
    - List[LayoutElement]: The final combined layout consisting of both the original layout
                           elements and the new OCR-derived elements.

    Note:
    - The function relies on `is_almost_subregion_of()` method to determine if an OCR region
      is a subregion of an existing layout element.
    - It also relies on `build_layout_elements_from_ocr_regions()` to convert OCR regions to
     layout elements.
    - The env_config `OCR_LAYOUT_SUBREGION_THRESHOLD` is used to specify the subregion matching
     threshold.
    """

    from unstructured_inference.inference.layoutelement import LayoutElements

    from unstructured.partition.pdf_image.inference_utils import (
        build_layout_elements_from_ocr_regions,
    )

    if len(layout) == 0:
        if len(ocr_layout) == 0:
            return layout
        else:
            ocr_regions_to_add = ocr_layout
    else:
        mask = ~bboxes1_is_almost_subregion_of_bboxes2(
            ocr_layout.element_coords, layout.element_coords, subregion_threshold
        ).sum(axis=1).astype(bool)

        # add ocr regions that are not covered by layout
        ocr_regions_to_add = ocr_layout.slice(mask)

    if len(ocr_regions_to_add):
        ocr_elements_to_add = build_layout_elements_from_ocr_regions(ocr_regions_to_add)
        final_layout = LayoutElements.concatenate([layout, ocr_elements_to_add])
    else:
        final_layout = layout

    return final_layout