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)