kreuzberg 3.0.0__py3-none-any.whl → 3.1.0__py3-none-any.whl

kreuzberg/_ocr/_paddleocr.py

@@ -0,0 +1,283 @@
+from __future__ import annotations
+
+import platform
+from dataclasses import dataclass
+from importlib.util import find_spec
+from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal
+
+from PIL import Image
+
+from kreuzberg._mime_types import PLAIN_TEXT_MIME_TYPE
+from kreuzberg._ocr._base import OCRBackend
+from kreuzberg._types import ExtractionResult, Metadata
+from kreuzberg._utils._string import normalize_spaces
+from kreuzberg._utils._sync import run_sync
+from kreuzberg.exceptions import MissingDependencyError, OCRError, ValidationError
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+
+try:  # pragma: no cover
+    from typing import Unpack  # type: ignore[attr-defined]
+except ImportError:  # pragma: no cover
+    from typing_extensions import Unpack
+
+
+PADDLEOCR_SUPPORTED_LANGUAGE_CODES: Final[set[str]] = {"ch", "en", "french", "german", "japan", "korean"}
+
+
+@dataclass(unsafe_hash=True, frozen=True)
+class PaddleOCRConfig:
+    """Configuration options for PaddleOCR.
+
+    This TypedDict provides type hints and documentation for all PaddleOCR parameters.
+    """
+
+    cls_image_shape: str = "3,48,192"
+    """Image shape for classification algorithm in format 'channels,height,width'."""
+    det_algorithm: Literal["DB", "EAST", "SAST", "PSE", "FCE", "PAN", "CT", "DB++", "Layout"] = "DB"
+    """Detection algorithm."""
+    det_db_box_thresh: float = 0.5
+    """Score threshold for detected boxes. Boxes below this value are discarded."""
+    det_db_thresh: float = 0.3
+    """Binarization threshold for DB output map."""
+    det_db_unclip_ratio: float = 2.0
+    """Expansion ratio for detected text boxes."""
+    det_east_cover_thresh: float = 0.1
+    """Score threshold for EAST output boxes."""
+    det_east_nms_thresh: float = 0.2
+    """NMS threshold for EAST model output boxes."""
+    det_east_score_thresh: float = 0.8
+    """Binarization threshold for EAST output map."""
+    det_max_side_len: int = 960
+    """Maximum size of image long side. Images exceeding this will be proportionally resized."""
+    drop_score: float = 0.5
+    """Filter recognition results by confidence score. Results below this are discarded."""
+    enable_mkldnn: bool = False
+    """Whether to enable MKL-DNN acceleration (Intel CPU only)."""
+    gpu_mem: int = 8000
+    """GPU memory size (in MB) to use for initialization."""
+    language: str = "en"
+    """Language to use for OCR."""
+    max_text_length: int = 25
+    """Maximum text length that the recognition algorithm can recognize."""
+    rec: bool = True
+    """Enable text recognition when using the ocr() function."""
+    rec_algorithm: Literal[
+        "CRNN",
+        "SRN",
+        "NRTR",
+        "SAR",
+        "SEED",
+        "SVTR",
+        "SVTR_LCNet",
+        "ViTSTR",
+        "ABINet",
+        "VisionLAN",
+        "SPIN",
+        "RobustScanner",
+        "RFL",
+    ] = "CRNN"
+    """Recognition algorithm."""
+    rec_image_shape: str = "3,32,320"
+    """Image shape for recognition algorithm in format 'channels,height,width'."""
+    table: bool = True
+    """Whether to enable table recognition."""
+    use_angle_cls: bool = True
+    """Whether to use text orientation classification model."""
+    use_gpu: bool = False
+    """Whether to use GPU for inference. Requires installing the paddlepaddle-gpu package"""
+    use_space_char: bool = True
+    """Whether to recognize spaces."""
+    use_zero_copy_run: bool = False
+    """Whether to enable zero_copy_run for inference optimization."""
+
+
+class PaddleBackend(OCRBackend[PaddleOCRConfig]):
+    _paddle_ocr: ClassVar[Any] = None
+
+    async def process_image(self, image: Image.Image, **kwargs: Unpack[PaddleOCRConfig]) -> ExtractionResult:
+        """Asynchronously process an image and extract its text and metadata using PaddleOCR.
+
+        Args:
+            image: An instance of PIL.Image representing the input image.
+            **kwargs: Configuration parameters for PaddleOCR including language, detection thresholds, etc.
+
+        Returns:
+            ExtractionResult: The extraction result containing text content, mime type, and metadata.
+
+        Raises:
+            OCRError: If OCR processing fails.
+        """
+        import numpy as np
+
+        await self._init_paddle_ocr(**kwargs)
+        image_np = np.array(image)
+        try:
+            result = await run_sync(self._paddle_ocr.ocr, image_np, cls=kwargs.get("use_angle_cls", True))
+            return self._process_paddle_result(result, image)
+        except Exception as e:
+            raise OCRError(f"Failed to OCR using PaddleOCR: {e}") from e
+
+    async def process_file(self, path: Path, **kwargs: Unpack[PaddleOCRConfig]) -> ExtractionResult:
+        """Asynchronously process a file and extract its text and metadata using PaddleOCR.
+
+        Args:
+            path: A Path object representing the file to be processed.
+            **kwargs: Configuration parameters for PaddleOCR including language, detection thresholds, etc.
+
+        Returns:
+            ExtractionResult: The extraction result containing text content, mime type, and metadata.
+
+        Raises:
+            OCRError: If file loading or OCR processing fails.
+        """
+        await self._init_paddle_ocr(**kwargs)
+        try:
+            image = await run_sync(Image.open, path)
+            return await self.process_image(image, **kwargs)
+        except Exception as e:
+            raise OCRError(f"Failed to load or process image using PaddleOCR: {e}") from e
+
+    @staticmethod
+    def _process_paddle_result(result: list[Any], image: Image.Image) -> ExtractionResult:
+        """Process PaddleOCR result into an ExtractionResult with metadata.
+
+        Args:
+            result: The raw result from PaddleOCR.
+            image: The original PIL image.
+
+        Returns:
+            ExtractionResult: The extraction result containing text content, mime type, and metadata.
+        """
+        text_content = ""
+        confidence_sum = 0
+        confidence_count = 0
+
+        for page_result in result:
+            if not page_result:
+                continue
+
+            sorted_boxes = sorted(page_result, key=lambda x: x[0][0][1])
+            line_groups: list[list[Any]] = []
+            current_line: list[Any] = []
+            prev_y: float | None = None
+
+            for box in sorted_boxes:
+                box_points, (_, _) = box
+                current_y = sum(point[1] for point in box_points) / 4
+                min_box_distance = 20
+
+                if prev_y is None or abs(current_y - prev_y) > min_box_distance:
+                    if current_line:
+                        line_groups.append(current_line)
+                    current_line = [box]
+                else:
+                    current_line.append(box)
+
+                prev_y = current_y
+
+            if current_line:
+                line_groups.append(current_line)
+
+            for line in line_groups:
+                line_sorted = sorted(line, key=lambda x: x[0][0][0])
+
+                for box in line_sorted:
+                    _, (text, confidence) = box
+                    if text:
+                        text_content += text + " "
+                        confidence_sum += confidence
+                        confidence_count += 1
+
+                text_content += "\n"
+
+        width, height = image.size
+        metadata = Metadata(
+            width=width,
+            height=height,
+        )
+
+        return ExtractionResult(
+            content=normalize_spaces(text_content), mime_type=PLAIN_TEXT_MIME_TYPE, metadata=metadata, chunks=[]
+        )
+
+    @classmethod
+    def _is_mkldnn_supported(cls) -> bool:
+        """Check if the current architecture supports MKL-DNN optimization.
+
+        Returns:
+            True if MKL-DNN is supported on this architecture.
+        """
+        system = platform.system().lower()
+        processor = platform.processor().lower()
+        machine = platform.machine().lower()
+
+        if system in ("linux", "windows"):
+            return "intel" in processor or "x86" in machine or "amd64" in machine or "x86_64" in machine
+
+        if system == "darwin":
+            return machine == "x86_64"
+
+        return False
+
+    @classmethod
+    async def _init_paddle_ocr(cls, **kwargs: Unpack[PaddleOCRConfig]) -> None:
+        """Initialize PaddleOCR with the provided configuration.
+
+        Args:
+            **kwargs: Configuration parameters for PaddleOCR including language, detection thresholds, etc.
+
+        Raises:
+            MissingDependencyError: If PaddleOCR is not installed.
+            OCRError: If initialization fails.
+        """
+        if cls._paddle_ocr is not None:
+            return
+
+        try:
+            from paddleocr import PaddleOCR
+        except ImportError as e:
+            raise MissingDependencyError.create_for_package(
+                dependency_group="paddleocr", functionality="PaddleOCR as an OCR backend", package_name="paddleocr"
+            ) from e
+
+        language = cls._validate_language_code(kwargs.pop("language", "en"))
+        has_gpu_package = bool(find_spec("paddlepaddle_gpu"))
+        kwargs.setdefault("use_angle_cls", True)
+        kwargs.setdefault("use_gpu", has_gpu_package)
+        kwargs.setdefault("enable_mkldnn", cls._is_mkldnn_supported() and not has_gpu_package)
+        kwargs.setdefault("det_db_thresh", 0.3)
+        kwargs.setdefault("det_db_box_thresh", 0.5)
+        kwargs.setdefault("det_db_unclip_ratio", 1.6)
+
+        try:
+            cls._paddle_ocr = await run_sync(PaddleOCR, lang=language, show_log=False, **kwargs)
+        except Exception as e:
+            raise OCRError(f"Failed to initialize PaddleOCR: {e}") from e
+
+    @staticmethod
+    def _validate_language_code(lang_code: str) -> str:
+        """Convert a language code to PaddleOCR format.
+
+        Args:
+            lang_code: ISO language code or language name
+
+        Raises:
+            ValidationError: If the language is not supported by PaddleOCR
+
+        Returns:
+            Language code compatible with PaddleOCR
+        """
+        normalized = lang_code.lower()
+        if normalized in PADDLEOCR_SUPPORTED_LANGUAGE_CODES:
+            return normalized
+
+        raise ValidationError(
+            "The provided language code is not supported by PaddleOCR",
+            context={
+                "language_code": lang_code,
+                "supported_languages": ",".join(sorted(PADDLEOCR_SUPPORTED_LANGUAGE_CODES)),
+            },
+        )

kreuzberg/_ocr/_tesseract.py

@@ -0,0 +1,342 @@
+from __future__ import annotations
+
+import re
+import sys
+from dataclasses import dataclass
+from enum import Enum
+from typing import TYPE_CHECKING, Any, ClassVar, Final
+
+from anyio import Path as AsyncPath
+from anyio import run_process
+
+from kreuzberg._mime_types import PLAIN_TEXT_MIME_TYPE
+from kreuzberg._ocr._base import OCRBackend
+from kreuzberg._types import ExtractionResult
+from kreuzberg._utils._string import normalize_spaces
+from kreuzberg._utils._sync import run_sync
+from kreuzberg._utils._tmp import create_temp_file
+from kreuzberg.exceptions import MissingDependencyError, OCRError, ValidationError
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+    from PIL.Image import Image
+
+try:  # pragma: no cover
+    from typing import Unpack  # type: ignore[attr-defined]
+except ImportError:  # pragma: no cover
+    from typing_extensions import Unpack
+
+
+TESSERACT_SUPPORTED_LANGUAGE_CODES: Final[set[str]] = {
+    "afr",
+    "amh",
+    "ara",
+    "asm",
+    "aze",
+    "aze_cyrl",
+    "bel",
+    "ben",
+    "bod",
+    "bos",
+    "bre",
+    "bul",
+    "cat",
+    "ceb",
+    "ces",
+    "chi_sim",
+    "chi_tra",
+    "chr",
+    "cos",
+    "cym",
+    "dan",
+    "dan_frak",
+    "deu",
+    "deu_frak",
+    "deu_latf",
+    "dzo",
+    "ell",
+    "eng",
+    "enm",
+    "epo",
+    "equ",
+    "est",
+    "eus",
+    "fao",
+    "fas",
+    "fil",
+    "fin",
+    "fra",
+    "frk",
+    "frm",
+    "fry",
+    "gla",
+    "gle",
+    "glg",
+    "grc",
+    "guj",
+    "hat",
+    "heb",
+    "hin",
+    "hrv",
+    "hun",
+    "hye",
+    "iku",
+    "ind",
+    "isl",
+    "ita",
+    "ita_old",
+    "jav",
+    "jpn",
+    "kan",
+    "kat",
+    "kat_old",
+    "kaz",
+    "khm",
+    "kir",
+    "kmr",
+    "kor",
+    "kor_vert",
+    "kur",
+    "lao",
+    "lat",
+    "lav",
+    "lit",
+    "ltz",
+    "mal",
+    "mar",
+    "mkd",
+    "mlt",
+    "mon",
+    "mri",
+    "msa",
+    "mya",
+    "nep",
+    "nld",
+    "nor",
+    "oci",
+    "ori",
+    "osd",
+    "pan",
+    "pol",
+    "por",
+    "pus",
+    "que",
+    "ron",
+    "rus",
+    "san",
+    "sin",
+    "slk",
+    "slk_frak",
+    "slv",
+    "snd",
+    "spa",
+    "spa_old",
+    "sqi",
+    "srp",
+    "srp_latn",
+    "sun",
+    "swa",
+    "swe",
+    "syr",
+    "tam",
+    "tat",
+    "tel",
+    "tgk",
+    "tgl",
+    "tha",  # codespell:ignore
+    "tir",
+    "ton",
+    "tur",
+    "uig",
+    "ukr",
+    "urd",
+    "uzb",
+    "uzb_cyrl",
+    "vie",  # codespell:ignore
+    "yid",
+    "yor",
+}
+
+MINIMAL_SUPPORTED_TESSERACT_VERSION: Final[int] = 5
+
+
+class PSMMode(Enum):
+    """Enum for Tesseract Page Segmentation Modes (PSM) with human-readable values."""
+
+    OSD_ONLY = 0
+    """Orientation and script detection only."""
+    AUTO_OSD = 1
+    """Automatic page segmentation with orientation and script detection."""
+    AUTO_ONLY = 2
+    """Automatic page segmentation without OSD."""
+    AUTO = 3
+    """Fully automatic page segmentation (default)."""
+    SINGLE_COLUMN = 4
+    """Assume a single column of text."""
+    SINGLE_BLOCK_VERTICAL = 5
+    """Assume a single uniform block of vertically aligned text."""
+    SINGLE_BLOCK = 6
+    """Assume a single uniform block of text."""
+    SINGLE_LINE = 7
+    """Treat the image as a single text line."""
+    SINGLE_WORD = 8
+    """Treat the image as a single word."""
+    CIRCLE_WORD = 9
+    """Treat the image as a single word in a circle."""
+    SINGLE_CHAR = 10
+    """Treat the image as a single character."""
+
+
+@dataclass(unsafe_hash=True, frozen=True)
+class TesseractConfig:
+    """Configuration options for Tesseract OCR engine."""
+
+    classify_use_pre_adapted_templates: bool = True
+    """Whether to use pre-adapted templates during classification to improve recognition accuracy."""
+    language: str = "eng"
+    """Language code to use for OCR.
+    Examples:
+            -   'eng' for English
+            -   'deu' for German
+            -    multiple languages combined with '+', e.g. 'eng+deu')
+    """
+    language_model_ngram_on: bool = True
+    """Enable or disable the use of n-gram-based language models for improved text recognition."""
+    psm: PSMMode = PSMMode.AUTO
+    """Page segmentation mode (PSM) to guide Tesseract on how to segment the image (e.g., single block, single line)."""
+    tessedit_dont_blkrej_good_wds: bool = True
+    """If True, prevents block rejection of words identified as good, improving text output quality."""
+    tessedit_dont_rowrej_good_wds: bool = True
+    """If True, prevents row rejection of words identified as good, avoiding unnecessary omissions."""
+    tessedit_enable_dict_correction: bool = True
+    """Enable or disable dictionary-based correction for recognized text to improve word accuracy."""
+    tessedit_use_primary_params_model: bool = True
+    """If True, forces the use of the primary parameters model for text recognition."""
+    textord_space_size_is_variable: bool = True
+    """Allow variable spacing between words, useful for text with irregular spacing."""
+    thresholding_method: bool = False
+    """Enable or disable specific thresholding methods during image preprocessing for better OCR accuracy."""
+
+
+class TesseractBackend(OCRBackend[TesseractConfig]):
+    _version_checked: ClassVar[bool] = False
+
+    async def process_image(
+        self,
+        image: Image,
+        **kwargs: Unpack[TesseractConfig],
+    ) -> ExtractionResult:
+        await self._validate_tesseract_version()
+        image_path, unlink = await create_temp_file(".png")
+        await run_sync(image.save, str(image_path), format="PNG")
+        try:
+            return await self.process_file(image_path, **kwargs)
+        finally:
+            await unlink()
+
+    async def process_file(
+        self,
+        path: Path,
+        **kwargs: Unpack[TesseractConfig],
+    ) -> ExtractionResult:
+        await self._validate_tesseract_version()
+        output_path, unlink = await create_temp_file(".txt")
+        language = self._validate_language_code(kwargs.pop("language", "eng"))
+        psm = kwargs.pop("psm", PSMMode.AUTO)
+        try:
+            output_base = str(output_path).replace(".txt", "")
+            command = [
+                "tesseract",
+                str(path),
+                output_base,
+                "-l",
+                language,
+                "--psm",
+                str(psm.value),
+                "--oem",
+                "1",
+                "--loglevel",
+                "OFF",
+            ]
+            for kwarg, value in kwargs.items():
+                command.extend(["-c", f"{kwarg}={1 if value else 0}"])
+
+            env: dict[str, Any] | None = None
+            if sys.platform.startswith("linux"):
+                # we have to prevent multithreading this way otherwise we will get deadlocks ~keep
+                env = {"OMP_THREAD_LIMIT": "1"}
+
+            result = await run_process(command, env=env)
+
+            if not result.returncode == 0:
+                raise OCRError(
+                    "OCR failed with a non-0 return code.",
+                    context={"error": result.stderr.decode() if isinstance(result.stderr, bytes) else result.stderr},
+                )
+
+            output = await AsyncPath(output_path).read_text("utf-8")
+            return ExtractionResult(
+                content=normalize_spaces(output), mime_type=PLAIN_TEXT_MIME_TYPE, metadata={}, chunks=[]
+            )
+        except (RuntimeError, OSError) as e:
+            raise OCRError(f"Failed to OCR using tesseract: {e}") from e
+        finally:
+            await unlink()
+
+    @classmethod
+    async def _validate_tesseract_version(cls) -> None:
+        """Validate that Tesseract is installed and is version 5 or above.
+
+        Raises:
+            MissingDependencyError: If Tesseract is not installed or is below version 5.
+        """
+        try:
+            if cls._version_checked:
+                return
+
+            command = ["tesseract", "--version"]
+            result = await run_process(command)
+            version_match = re.search(r"tesseract\s+v?(\d+)\.\d+\.\d+", result.stdout.decode())
+            if not version_match or int(version_match.group(1)) < MINIMAL_SUPPORTED_TESSERACT_VERSION:
+                raise MissingDependencyError(
+                    "Tesseract version 5 is a required system dependency. Please install it on your system and make sure its available in $PATH."
+                )
+
+            cls._version_checked = True
+        except FileNotFoundError as e:
+            raise MissingDependencyError(
+                "Tesseract version 5 is a required system dependency. Please install it on your system and make sure its available in $PATH."
+            ) from e
+
+    @staticmethod
+    def _validate_language_code(language_code: str) -> str:
+        """Convert a language code to Tesseract format.
+
+        Args:
+            language_code: Tesseract supported language code or multiple language codes connected with '+'
+
+        Raises:
+            ValidationError: If the language is not supported by Tesseract
+
+        Returns:
+            Language code compatible with Tesseract
+        """
+        normalized = language_code.lower()
+        if normalized in TESSERACT_SUPPORTED_LANGUAGE_CODES:
+            return normalized
+
+        if "+" in normalized and all(lang in TESSERACT_SUPPORTED_LANGUAGE_CODES for lang in normalized.split("+")):
+            return normalized
+
+        raise ValidationError(
+            "The provided language code is not supported by Tesseract",
+            context={
+                "language_code": normalized
+                if "+" not in normalized
+                else ",".join(
+                    [lang for lang in normalized.split("+") if lang not in TESSERACT_SUPPORTED_LANGUAGE_CODES]
+                ),
+                "supported_languages": ",".join(sorted(TESSERACT_SUPPORTED_LANGUAGE_CODES)),
+            },
+        )

kreuzberg/_types.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import sys
 from collections.abc import Awaitable
-from dataclasses import asdict, dataclass
+from dataclasses import asdict, dataclass, field
 from typing import TYPE_CHECKING, Any, Callable, Literal, TypedDict, Union
 
 from kreuzberg._constants import DEFAULT_MAX_CHARACTERS, DEFAULT_MAX_OVERLAP
@@ -14,6 +14,10 @@ else:  # pragma: no cover
     from typing import NotRequired
 
 if TYPE_CHECKING:
+    from pandas import DataFrame
+    from PIL.Image import Image
+
+    from kreuzberg._gmft import GMFTConfig
     from kreuzberg._ocr._easyocr import EasyOCRConfig
     from kreuzberg._ocr._paddleocr import PaddleOCRConfig
     from kreuzberg._ocr._tesseract import TesseractConfig
@@ -21,6 +25,19 @@ if TYPE_CHECKING:
 OcrBackendType = Literal["tesseract", "easyocr", "paddleocr"]
 
 
+class TableData(TypedDict):
+    """Table data, returned from table extraction."""
+
+    cropped_image: Image
+    """The cropped image of the table."""
+    df: DataFrame
+    """The table data as a pandas DataFrame."""
+    page_number: int
+    """The page number of the table."""
+    text: str
+    """The table text as a markdown string."""
+
+
 class Metadata(TypedDict, total=False):
     """Base metadata common to all document types.
 
@@ -88,12 +105,14 @@ class ExtractionResult:
 
     content: str
     """The extracted content."""
-    chunks: list[str]
-    """The extracted content chunks. This is an empty list if 'chunk_content' is not set to True in the ExtractionConfig."""
     mime_type: str
     """The mime type of the extracted content. Is either text/plain or text/markdown."""
     metadata: Metadata
     """The metadata of the content."""
+    tables: list[TableData] = field(default_factory=list)
+    """Extracted tables. Is an empty list if 'extract_tables' is not set to True in the ExtractionConfig."""
+    chunks: list[str] = field(default_factory=list)
+    """The extracted content chunks. This is an empty list if 'chunk_content' is not set to True in the ExtractionConfig."""
 
 
 PostProcessingHook = Callable[[ExtractionResult], Union[ExtractionResult, Awaitable[ExtractionResult]]]
@@ -114,14 +133,22 @@ class ExtractionConfig:
     """Whether to force OCR."""
     chunk_content: bool = False
     """Whether to chunk the content into smaller chunks."""
+    extract_tables: bool = False
+    """Whether to extract tables from the content. This requires the 'gmft' dependency."""
     max_chars: int = DEFAULT_MAX_CHARACTERS
     """The size of each chunk in characters."""
     max_overlap: int = DEFAULT_MAX_OVERLAP
     """The overlap between chunks in characters."""
     ocr_backend: OcrBackendType | None = "tesseract"
-    """The OCR backend to use."""
+    """The OCR backend to use.
+
+    Notes:
+        - If set to 'None', OCR will not be performed.
+    """
     ocr_config: TesseractConfig | PaddleOCRConfig | EasyOCRConfig | None = None
     """Configuration to pass to the OCR backend."""
+    gmft_config: GMFTConfig | None = None
+    """GMFT configuration."""
     post_processing_hooks: list[PostProcessingHook] | None = None
     """Post processing hooks to call after processing is done and before the final result is returned."""
     validators: list[ValidationHook] | None = None