chore: import upstream snapshot with attribution

unstructured/partition/utils/ocr_models/google_vision_ocr.py

@@ -0,0 +1,116 @@
+from __future__ import annotations
+
+from io import BytesIO
+from typing import TYPE_CHECKING, Optional
+
+from google.cloud.vision import Image, ImageAnnotatorClient, ImageContext, Paragraph, TextAnnotation
+
+from unstructured.logger import logger, trace_logger
+from unstructured.partition.utils.config import env_config
+from unstructured.partition.utils.constants import Source
+from unstructured.partition.utils.ocr_models.ocr_interface import OCRAgent
+
+if TYPE_CHECKING:
+    from PIL import Image as PILImage
+    from unstructured_inference.inference.elements import TextRegion, TextRegions
+    from unstructured_inference.inference.layoutelement import LayoutElements
+
+
+class OCRAgentGoogleVision(OCRAgent):
+    """OCR service implementation for Google Vision API."""
+
+    def __init__(self, language: Optional[str] = None) -> None:
+        self.language = language
+        client_options = {}
+        api_endpoint = env_config.GOOGLEVISION_API_ENDPOINT
+        if api_endpoint:
+            logger.info(f"Using Google Vision OCR with endpoint {api_endpoint}")
+            client_options["api_endpoint"] = api_endpoint
+        else:
+            logger.info("Using Google Vision OCR with default endpoint")
+        self.client = ImageAnnotatorClient(client_options=client_options)
+
+    def is_text_sorted(self) -> bool:
+        return True
+
+    def get_text_from_image(self, image: PILImage.Image) -> str:
+        image_context = ImageContext(language_hints=[self.language]) if self.language else None
+        with BytesIO() as buffer:
+            image.save(buffer, format="PNG")
+            response = self.client.document_text_detection(
+                image=Image(content=buffer.getvalue()), image_context=image_context
+            )
+        document = response.full_text_annotation
+        assert isinstance(document, TextAnnotation)
+        return document.text
+
+    def get_layout_from_image(self, image: PILImage.Image) -> TextRegions:
+        trace_logger.detail("Processing entire page OCR with Google Vision API...")
+        image_context = ImageContext(language_hints=[self.language]) if self.language else None
+        with BytesIO() as buffer:
+            image.save(buffer, format="PNG")
+            response = self.client.document_text_detection(
+                image=Image(content=buffer.getvalue()), image_context=image_context
+            )
+        document = response.full_text_annotation
+        assert isinstance(document, TextAnnotation)
+        regions = self._parse_regions(document)
+        return regions
+
+    def get_layout_elements_from_image(self, image: PILImage.Image) -> LayoutElements:
+        from unstructured.partition.pdf_image.inference_utils import (
+            build_layout_elements_from_ocr_regions,
+        )
+
+        ocr_regions = self.get_layout_from_image(
+            image,
+        )
+        ocr_text = self.get_text_from_image(
+            image,
+        )
+        return build_layout_elements_from_ocr_regions(
+            ocr_regions=ocr_regions,
+            ocr_text=ocr_text,
+            group_by_ocr_text=False,
+        )
+
+    def _parse_regions(self, ocr_data: TextAnnotation) -> TextRegions:
+        from unstructured_inference.inference.elements import TextRegions
+
+        from unstructured.partition.pdf_image.inference_utils import build_text_region_from_coords
+
+        text_regions: list[TextRegion] = []
+        for page_idx, page in enumerate(ocr_data.pages):
+            for block in page.blocks:
+                for paragraph in block.paragraphs:
+                    vertices = paragraph.bounding_box.vertices
+                    x1, y1 = vertices[0].x, vertices[0].y
+                    x2, y2 = vertices[2].x, vertices[2].y
+                    text_region = build_text_region_from_coords(
+                        x1,
+                        y1,
+                        x2,
+                        y2,
+                        text=self._get_text_from_paragraph(paragraph),
+                        source=Source.OCR_GOOGLEVISION,
+                    )
+                    text_regions.append(text_region)
+        return TextRegions.from_list(text_regions)
+
+    def _get_text_from_paragraph(self, paragraph: Paragraph) -> str:
+        breaks = TextAnnotation.DetectedBreak.BreakType
+        para = ""
+        line = ""
+        for word in paragraph.words:
+            for symbol in word.symbols:
+                line += symbol.text
+                if symbol.property.detected_break.type_ == breaks.SPACE:
+                    line += " "
+                if symbol.property.detected_break.type_ == breaks.EOL_SURE_SPACE:
+                    line += " "
+                    para += line
+                    line = ""
+                if symbol.property.detected_break.type_ == breaks.LINE_BREAK:
+                    para += line
+                    line = ""
+        return para

unstructured/partition/utils/ocr_models/ocr_interface.py

@@ -0,0 +1,96 @@
+from __future__ import annotations
+
+import functools
+import importlib
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
+
+from unstructured.logger import logger
+from unstructured.partition.utils.config import env_config
+from unstructured.partition.utils.constants import (
+    OCR_AGENT_MODULES_WHITELIST,
+    OCR_AGENT_PADDLE,
+    OCR_AGENT_PADDLE_OLD,
+    OCR_AGENT_TESSERACT,
+    OCR_AGENT_TESSERACT_OLD,
+)
+
+if TYPE_CHECKING:
+    from PIL import Image as PILImage
+    from unstructured_inference.inference.elements import TextRegions
+    from unstructured_inference.inference.layoutelement import LayoutElements
+
+
+class OCRAgent(ABC):
+    """Defines the interface for an Optical Character Recognition (OCR) service."""
+
+    @classmethod
+    def get_agent(cls, language: str) -> OCRAgent:
+        """Get the configured OCRAgent instance.
+
+        The OCR package used by the agent is determined by the `OCR_AGENT` environment variable.
+        """
+        ocr_agent_cls_qname = cls._get_ocr_agent_cls_qname()
+        return cls.get_instance(ocr_agent_cls_qname, language)
+
+    @staticmethod
+    @functools.lru_cache(maxsize=env_config.OCR_AGENT_CACHE_SIZE)
+    def get_instance(ocr_agent_module: str, language: str) -> "OCRAgent":
+        module_name, class_name = ocr_agent_module.rsplit(".", 1)
+        if module_name not in OCR_AGENT_MODULES_WHITELIST:
+            raise ValueError(
+                f"Environment variable OCR_AGENT module name {module_name} must be set to a "
+                f"whitelisted module part of {OCR_AGENT_MODULES_WHITELIST}."
+            )
+
+        try:
+            module = importlib.import_module(module_name)
+            loaded_class = getattr(module, class_name)
+            return loaded_class(language)
+        except (ImportError, AttributeError) as e:
+            logger.error(f"Failed to get OCRAgent instance: {e}")
+            raise RuntimeError(
+                "Could not get the OCRAgent instance. Please check the OCR package and the "
+                "OCR_AGENT environment variable."
+            )
+
+    @abstractmethod
+    def get_layout_elements_from_image(self, image: PILImage.Image) -> LayoutElements:
+        pass
+
+    @abstractmethod
+    def get_layout_from_image(self, image: PILImage.Image) -> TextRegions:
+        pass
+
+    @abstractmethod
+    def get_text_from_image(self, image: PILImage.Image) -> str:
+        pass
+
+    @abstractmethod
+    def is_text_sorted(self) -> bool:
+        pass
+
+    @staticmethod
+    def _get_ocr_agent_cls_qname() -> str:
+        """Get the fully-qualified class name of the configured OCR agent.
+
+        The qualified name (qname) looks like:
+            "unstructured.partition.utils.ocr_models.tesseract_ocr.OCRAgentTesseract"
+
+        The qname provides the full module address and class name of the OCR agent.
+        """
+        ocr_agent_qname = env_config.OCR_AGENT
+
+        # -- map legacy method of setting OCR agent by key-name to full qname --
+        qnames_by_keyname = {
+            OCR_AGENT_TESSERACT_OLD: OCR_AGENT_TESSERACT,
+            OCR_AGENT_PADDLE_OLD: OCR_AGENT_PADDLE,
+        }
+        if qname_mapped_from_keyname := qnames_by_keyname.get(ocr_agent_qname.lower()):
+            logger.warning(
+                f"OCR agent name {ocr_agent_qname} is outdated and will be removed in a future"
+                f" release; please use {qname_mapped_from_keyname} instead"
+            )
+            return qname_mapped_from_keyname
+
+        return ocr_agent_qname

unstructured/partition/utils/ocr_models/paddle_ocr.py

@@ -0,0 +1,146 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+from PIL import Image as PILImage
+
+from unstructured.documents.elements import ElementType
+from unstructured.logger import logger, trace_logger
+from unstructured.partition.utils.constants import Source
+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.layoutelement import LayoutElements
+
+
+class OCRAgentPaddle(OCRAgent):
+    """OCR service implementation for PaddleOCR."""
+
+    def __init__(self, language: str = "en"):
+        self.agent = self.load_agent(language)
+
+    def load_agent(self, language: str):
+        """Loads the PaddleOCR agent as a global variable to ensure that we only load it once."""
+
+        import paddle
+        from unstructured_paddleocr import PaddleOCR
+
+        # Disable signal handlers at C++ level upon failing
+        # ref: https://www.paddlepaddle.org.cn/documentation/docs/en/api/paddle/
+        #      disable_signal_handler_en.html#disable-signal-handler
+        paddle.disable_signal_handler()
+        # Use paddlepaddle-gpu if there is gpu device available
+        gpu_available = paddle.device.cuda.device_count() > 0
+        if gpu_available:
+            logger.info(f"Loading paddle with GPU on language={language}...")
+        else:
+            logger.info(f"Loading paddle with CPU on language={language}...")
+        try:
+            # Enable MKL-DNN for paddle to speed up OCR if OS supports it
+            # ref: https://paddle-inference.readthedocs.io/en/master/
+            #      api_reference/cxx_api_doc/Config/CPUConfig.html
+            paddle_ocr = PaddleOCR(
+                use_angle_cls=True,
+                use_gpu=gpu_available,
+                lang=language,
+                enable_mkldnn=True,
+                show_log=False,
+                rec_batch_num=1,
+            )
+        except AttributeError:
+            paddle_ocr = PaddleOCR(
+                use_angle_cls=True,
+                use_gpu=gpu_available,
+                lang=language,
+                enable_mkldnn=False,
+                show_log=False,
+                rec_batch_num=1,
+            )
+        return paddle_ocr
+
+    def get_text_from_image(self, image: PILImage.Image) -> str:
+        ocr_regions = self.get_layout_from_image(image)
+        return "\n\n".join(ocr_regions.texts)
+
+    def is_text_sorted(self):
+        return False
+
+    def get_layout_from_image(self, image: PILImage.Image) -> TextRegions:
+        """Get the OCR regions from image as a list of text regions with paddle."""
+
+        trace_logger.detail("Processing entire page OCR with paddle...")
+
+        # TODO(yuming): pass in language parameter once we
+        # have the mapping for paddle lang code
+        # see CORE-2034
+        ocr_data = self.agent.ocr(np.array(image), cls=True)
+        ocr_regions = self.parse_data(ocr_data)
+
+        return ocr_regions
+
+    @requires_dependencies("unstructured_inference")
+    def get_layout_elements_from_image(self, image: PILImage.Image) -> LayoutElements:
+        ocr_regions = self.get_layout_from_image(image)
+
+        # NOTE(christine): For paddle, there is no difference in `ocr_layout` and `ocr_text` in
+        # terms of grouping because we get ocr_text from `ocr_layout, so the first two grouping
+        # and merging steps are not necessary.
+        return LayoutElements(
+            element_coords=ocr_regions.element_coords,
+            texts=ocr_regions.texts,
+            element_class_ids=np.zeros(ocr_regions.texts.shape),
+            element_class_id_map={0: ElementType.UNCATEGORIZED_TEXT},
+        )
+
+    @requires_dependencies("unstructured_inference")
+    def parse_data(self, ocr_data: list[Any]) -> TextRegions:
+        """Parse the OCR result data to extract a list of TextRegion objects from paddle.
+
+        The function processes the OCR result dictionary, looking for bounding
+        box information and associated text to create instances of the TextRegion
+        class, which are then appended to a list.
+
+        Parameters:
+        - ocr_data (list): A list containing the OCR result data
+
+        Returns:
+        - TextRegions:
+            TextRegions object, containing data from all text regions in numpy arrays; each row
+            represents a detected text region within the OCR-ed image.
+
+        Note:
+        - An empty string or a None value for the 'text' key in the input
+          dictionary will result in its associated bounding box being ignored.
+        """
+
+        from unstructured_inference.inference.elements import TextRegions
+
+        from unstructured.partition.pdf_image.inference_utils import build_text_region_from_coords
+
+        text_regions: list[TextRegion] = []
+        for idx in range(len(ocr_data)):
+            res = ocr_data[idx]
+            if not res:
+                continue
+
+            for line in res:
+                x1 = min([i[0] for i in line[0]])
+                y1 = min([i[1] for i in line[0]])
+                x2 = max([i[0] for i in line[0]])
+                y2 = max([i[1] for i in line[0]])
+                text = line[1][0]
+                if not text:
+                    continue
+                cleaned_text = text.strip()
+                if cleaned_text:
+                    text_region = build_text_region_from_coords(
+                        x1, y1, x2, y2, text=cleaned_text, source=Source.OCR_PADDLE
+                    )
+                    text_regions.append(text_region)
+
+        # FIXME (yao): find out if paddle supports a vectorized output format so we can skip the
+        # step of parsing a list
+        return TextRegions.from_list(text_regions)

unstructured/partition/utils/ocr_models/tesseract_ocr.py

@@ -0,0 +1,260 @@
+from __future__ import annotations
+
+import os
+import re
+from typing import TYPE_CHECKING
+
+import cv2
+import numpy as np
+import pandas as pd
+import unstructured_pytesseract
+from lxml import etree
+from PIL import Image as PILImage
+
+from unstructured.logger import trace_logger
+from unstructured.partition.utils.config import env_config
+from unstructured.partition.utils.constants import (
+    IMAGE_COLOR_DEPTH,
+    TESSERACT_MAX_SIZE,
+    TESSERACT_TEXT_HEIGHT,
+    Source,
+)
+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 TextRegions
+    from unstructured_inference.inference.layoutelement import LayoutElements
+
+_RE_X_CONF = re.compile(r"x_conf (\d+\.\d+)")
+
+# -- force tesseract to be single threaded, otherwise we see major performance problems --
+if "OMP_THREAD_LIMIT" not in os.environ:
+    os.environ["OMP_THREAD_LIMIT"] = "1"
+
+
+class OCRAgentTesseract(OCRAgent):
+    """OCR service implementation for Tesseract."""
+
+    hocr_namespace = {"h": "http://www.w3.org/1999/xhtml"}
+
+    def __init__(self, language: str = "eng"):
+        self.language = language
+
+    def is_text_sorted(self):
+        return True
+
+    def get_text_from_image(self, image: PILImage.Image) -> str:
+        return unstructured_pytesseract.image_to_string(np.array(image), lang=self.language)
+
+    def get_layout_from_image(self, image: PILImage.Image) -> TextRegions:
+        """Get the OCR regions from image as a list of text regions with tesseract."""
+
+        trace_logger.detail("Processing entire page OCR with tesseract...")
+        zoom = 1
+        ocr_df: pd.DataFrame = self.image_to_data_with_character_confidence_filter(
+            np.array(image),
+            lang=self.language,
+            character_confidence_threshold=env_config.TESSERACT_CHARACTER_CONFIDENCE_THRESHOLD,
+        )
+        ocr_df = ocr_df.dropna()
+
+        # tesseract performance degrades when the text height is out of the preferred zone so we
+        # zoom the image (in or out depending on estimated text height) for optimum OCR results
+        # but this needs to be evaluated based on actual use case as the optimum scaling also
+        # depend on type of characters (font, language, etc); be careful about this
+        # functionality
+        text_height = ocr_df[TESSERACT_TEXT_HEIGHT].quantile(
+            env_config.TESSERACT_TEXT_HEIGHT_QUANTILE
+        )
+        if (
+            text_height < env_config.TESSERACT_MIN_TEXT_HEIGHT
+            or text_height > env_config.TESSERACT_MAX_TEXT_HEIGHT
+        ):
+            max_zoom = max(
+                0,
+                np.round(np.sqrt(TESSERACT_MAX_SIZE / np.prod(image.size) / IMAGE_COLOR_DEPTH), 1),
+            )
+            # rounding avoids unnecessary precision and potential numerical issues associated
+            # with numbers very close to 1 inside cv2 image processing
+            zoom = min(
+                np.round(env_config.TESSERACT_OPTIMUM_TEXT_HEIGHT / text_height, 1),
+                max_zoom,
+            )
+            ocr_df = self.image_to_data_with_character_confidence_filter(
+                np.array(zoom_image(image, zoom)),
+                lang=self.language,
+                character_confidence_threshold=env_config.TESSERACT_CHARACTER_CONFIDENCE_THRESHOLD,
+            )
+            ocr_df = ocr_df.dropna()
+        ocr_regions = self.parse_data(ocr_df, zoom=zoom)
+
+        return ocr_regions
+
+    def image_to_data_with_character_confidence_filter(
+        self,
+        image: np.ndarray,
+        lang: str = "eng",
+        config: str = "",
+        character_confidence_threshold: float = 0.0,
+    ) -> pd.DataFrame:
+        hocr: str = unstructured_pytesseract.image_to_pdf_or_hocr(
+            image,
+            lang=lang,
+            config="-c hocr_char_boxes=1 " + config,
+            extension="hocr",
+        )
+        ocr_df = self.hocr_to_dataframe(hocr, character_confidence_threshold)
+        return ocr_df
+
+    def hocr_to_dataframe(
+        self, hocr: str, character_confidence_threshold: float = 0.0
+    ) -> pd.DataFrame:
+        df_entries = []
+
+        if not hocr:
+            return pd.DataFrame(df_entries, columns=["left", "top", "width", "height", "text"])
+
+        root = etree.fromstring(hocr)
+        word_spans = root.findall('.//h:span[@class="ocrx_word"]', self.hocr_namespace)
+
+        for word_span in word_spans:
+            word_title = word_span.get("title", "")
+            bbox_match = re.search(r"bbox (\d+) (\d+) (\d+) (\d+)", word_title)
+
+            text = self.extract_word_from_hocr(
+                word=word_span, character_confidence_threshold=character_confidence_threshold
+            )
+            if text and bbox_match:
+                word_bbox = list(map(int, bbox_match.groups()))
+                left, top, right, bottom = word_bbox
+                df_entries.append(
+                    {
+                        "left": left,
+                        "top": top,
+                        "right": right,
+                        "bottom": bottom,
+                        "text": text,
+                    }
+                )
+        ocr_df = pd.DataFrame(df_entries, columns=["left", "top", "right", "bottom", "text"])
+
+        ocr_df["width"] = ocr_df["right"] - ocr_df["left"]
+        ocr_df["height"] = ocr_df["bottom"] - ocr_df["top"]
+
+        ocr_df = ocr_df.drop(columns=["right", "bottom"])
+        return ocr_df
+
+    def extract_word_from_hocr(
+        self, word: etree.Element, character_confidence_threshold: float = 0.0
+    ) -> str:
+        """Extracts a word from an hOCR word tag, filtering out characters with low confidence."""
+
+        character_spans = word.findall('.//h:span[@class="ocrx_cinfo"]', self.hocr_namespace)
+        if len(character_spans) == 0:
+            return ""
+
+        chars = []
+        for character_span in character_spans:
+            char = character_span.text
+
+            char_title = character_span.get("title", "")
+            conf_match = _RE_X_CONF.search(char_title)
+
+            if not (char and conf_match):
+                continue
+
+            character_probability = float(conf_match.group(1)) / 100
+
+            if character_probability >= character_confidence_threshold:
+                chars.append(char)
+
+        return "".join(chars)
+
+    @requires_dependencies("unstructured_inference")
+    def get_layout_elements_from_image(self, image: PILImage.Image) -> LayoutElements:
+        from unstructured.partition.pdf_image.inference_utils import (
+            build_layout_elements_from_ocr_regions,
+        )
+
+        ocr_regions = self.get_layout_from_image(image)
+
+        # NOTE(christine): For tesseract, the ocr_text returned by
+        # `unstructured_pytesseract.image_to_string()` doesn't contain bounding box data but is
+        # well grouped. Conversely, the ocr_layout returned by parsing
+        # `unstructured_pytesseract.image_to_data()` contains bounding box data but is not well
+        # grouped. Therefore, we need to first group the `ocr_layout` by `ocr_text` and then merge
+        # the text regions in each group to create a list of layout elements.
+
+        ocr_text = self.get_text_from_image(image)
+
+        return build_layout_elements_from_ocr_regions(
+            ocr_regions=ocr_regions,
+            ocr_text=ocr_text,
+            group_by_ocr_text=True,
+        )
+
+    @requires_dependencies("unstructured_inference")
+    def parse_data(self, ocr_data: pd.DataFrame, zoom: float = 1) -> TextRegions:
+        """Parse the OCR result data to extract a list of TextRegion objects from tesseract.
+
+        The function processes the OCR result data frame, looking for bounding
+        box information and associated text to create instances of the TextRegion
+        class, which are then appended to a list.
+
+        Parameters:
+        - ocr_data (pd.DataFrame):
+            A Pandas DataFrame containing the OCR result data.
+            It should have columns like 'text', 'left', 'top', 'width', and 'height'.
+
+        - zoom (float, optional):
+            A zoom factor to scale the coordinates of the bounding boxes from image scaling.
+            Default is 1.
+
+        Returns:
+        - TextRegions:
+            TextRegions object, containing data from all text regions in numpy arrays; each row
+            represents a detected text region within the OCR-ed image.
+
+        Note:
+        - An empty string or a None value for the 'text' key in the input
+          data frame will result in its associated bounding box being ignored.
+        """
+
+        from unstructured_inference.inference.elements import TextRegions
+
+        if zoom <= 0:
+            zoom = 1
+
+        texts = ocr_data.text.apply(
+            lambda text: str(text) if not isinstance(text, str) else text.strip()
+        ).values
+        mask = texts != ""
+        element_coords = ocr_data[["left", "top", "width", "height"]].values
+        element_coords[:, 2] += element_coords[:, 0]
+        element_coords[:, 3] += element_coords[:, 1]
+        element_coords = element_coords.astype(float) / zoom
+        return TextRegions(
+            element_coords=element_coords[mask],
+            texts=texts[mask],
+            sources=np.array([Source.OCR_TESSERACT] * mask.sum()),
+        )
+
+
+def zoom_image(image: PILImage.Image, zoom: float = 1) -> PILImage.Image:
+    """scale an image based on the zoom factor using cv2; the scaled image is post processed by
+    dilation then erosion to improve edge sharpness for OCR tasks"""
+    if zoom <= 0:
+        # no zoom but still does dilation and erosion
+        zoom = 1
+    new_image = cv2.resize(
+        cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR),
+        None,
+        fx=zoom,
+        fy=zoom,
+        interpolation=cv2.INTER_CUBIC,
+    )
+
+    # Skip dilation and erosion for 1x1 kernel as they are no-ops
+
+    return PILImage.fromarray(new_image)