kreuzberg 2.1.2__py3-none-any.whl → 3.0.1__py3-none-any.whl

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/_playa.py

@@ -0,0 +1,276 @@
+from __future__ import annotations
+
+from datetime import datetime
+from typing import TYPE_CHECKING, Any, cast
+
+from playa import asobj, parse
+from playa.utils import decode_text
+
+from kreuzberg.exceptions import ParsingError
+
+if TYPE_CHECKING:
+    from playa.document import Document
+
+    from kreuzberg._types import Metadata
+
+
+GRAY_COMPONENTS = 1
+RGB_COMPONENTS = 3
+CMYK_COMPONENTS = 4
+UTF16BE_BOM = b"\xfe\xff"
+UTF16BE_ENCODING = "utf-16be"
+MIN_DATE_LENGTH = 8
+FULL_DATE_LENGTH = 14
+BOM_CHAR = "\ufeff"
+
+
+async def extract_pdf_metadata(pdf_content: bytes) -> Metadata:
+    """Extract metadata from a PDF document.
+
+    Args:
+        pdf_content: The bytes of the PDF document.
+
+    Raises:
+        ParsingError: If the PDF metadata could not be extracted.
+
+    Returns:
+        A dictionary of metadata extracted from the PDF.
+    """
+    try:
+        document = parse(pdf_content, max_workers=1)
+        metadata: Metadata = {}
+
+        for raw_info in document.info:
+            pdf_info = {k.lower(): v for k, v in asobj(raw_info).items()}
+            _extract_basic_metadata(pdf_info, metadata)
+            _extract_author_metadata(pdf_info, metadata)
+            _extract_keyword_metadata(pdf_info, metadata)
+            _extract_category_metadata(pdf_info, metadata)
+            _extract_date_metadata(pdf_info, metadata)
+            _extract_creator_metadata(pdf_info, metadata)
+
+        if document.pages:
+            _extract_document_dimensions(document, metadata)
+
+        if document.outline and "description" not in metadata:
+            metadata["description"] = _generate_outline_description(document)
+
+        if "summary" not in metadata:
+            metadata["summary"] = _generate_document_summary(document)
+
+        _extract_structure_information(document, metadata)
+
+        return metadata
+    except Exception as e:
+        raise ParsingError(f"Failed to extract PDF metadata: {e!s}") from e
+
+
+def _extract_basic_metadata(pdf_info: dict[str, Any], result: Metadata) -> None:
+    if "title" not in result and (title := pdf_info.get("title")):
+        result["title"] = decode_text(title)
+
+    if "subject" not in result and (subject := pdf_info.get("subject")):
+        result["subject"] = decode_text(subject)
+
+    if "publisher" not in result and (publisher := pdf_info.get("Publisher", pdf_info.get("publisher"))):
+        result["publisher"] = decode_text(publisher)
+
+    if "copyright" not in result and (copyright_info := pdf_info.get("copyright") or pdf_info.get("rights")):
+        result["copyright"] = decode_text(copyright_info)
+
+    if "comments" not in result and (comments := pdf_info.get("comments")):
+        result["comments"] = decode_text(comments)
+
+    if "identifier" not in result and (identifier := pdf_info.get("identifier") or pdf_info.get("id")):
+        result["identifier"] = decode_text(identifier)
+
+    if "license" not in result and (license_info := pdf_info.get("license")):
+        result["license"] = decode_text(license_info)
+
+    if "modified_by" not in result and (modified_by := pdf_info.get("modifiedby") or pdf_info.get("last_modified_by")):
+        result["modified_by"] = decode_text(modified_by)
+
+    if "version" not in result and (version := pdf_info.get("version")):
+        result["version"] = decode_text(version)
+
+
+def _extract_author_metadata(pdf_info: dict[str, Any], result: Metadata) -> None:
+    if author := pdf_info.get("author"):
+        if isinstance(author, (str, bytes)):
+            author_str = decode_text(author)
+            author_str = author_str.replace(" and ", ", ")
+
+            authors = []
+            for author_segment in author_str.split(";"):
+                authors.extend(
+                    [author_name.strip() for author_name in author_segment.split(",") if author_name.strip()]
+                )
+            result["authors"] = authors
+        elif isinstance(author, list):
+            result["authors"] = [decode_text(a) for a in author]
+
+
+def _extract_keyword_metadata(pdf_info: dict[str, Any], result: Metadata) -> None:
+    if keywords := pdf_info.get("keywords"):
+        if isinstance(keywords, (str, bytes)):
+            kw_str = decode_text(keywords)
+            kw_list = [k.strip() for k in kw_str.split(",")]
+            kw_list = [k.strip() for k in " ".join(kw_list).split(";")]
+            result["keywords"] = [k for k in kw_list if k]
+        elif isinstance(keywords, list):
+            result["keywords"] = [decode_text(k) for k in keywords]
+
+
+def _extract_category_metadata(pdf_info: dict[str, Any], result: Metadata) -> None:
+    if categories := pdf_info.get("categories") or pdf_info.get("category"):
+        if isinstance(categories, (str, bytes)):
+            cat_str = decode_text(categories)
+            cat_list = [c.strip() for c in cat_str.split(",")]
+            result["categories"] = [c for c in cat_list if c]
+        elif isinstance(categories, list):
+            result["categories"] = [decode_text(c) for c in categories]
+
+
+def _parse_date_string(date_str: str) -> str:
+    date_str = date_str.removeprefix("D:")
+    if len(date_str) >= MIN_DATE_LENGTH:
+        year = date_str[0:4]
+        month = date_str[4:6]
+        day = date_str[6:8]
+        time_part = ""
+        if len(date_str) >= FULL_DATE_LENGTH:
+            hour = date_str[8:10]
+            minute = date_str[10:12]
+            second = date_str[12:14]
+            time_part = f"T{hour}:{minute}:{second}"
+        return datetime.strptime(f"{year}-{month}-{day}{time_part}", "%Y%m%d%H%M%S").isoformat()  # noqa: DTZ007
+    return date_str
+
+
+def _extract_date_metadata(pdf_info: dict[str, Any], result: Metadata) -> None:
+    if created := pdf_info.get("creationdate") or pdf_info.get("createdate"):
+        try:
+            date_str = decode_text(created)
+            result["created_at"] = _parse_date_string(date_str)
+        except (ValueError, IndexError):
+            result["created_at"] = decode_text(created)
+
+    if modified := pdf_info.get("moddate") or pdf_info.get("modificationdate"):
+        try:
+            date_str = decode_text(modified)
+            result["modified_at"] = _parse_date_string(date_str)
+        except (ValueError, IndexError):
+            result["modified_at"] = decode_text(modified)
+
+
+def _extract_creator_metadata(pdf_info: dict[str, Any], result: Metadata) -> None:
+    if creator := pdf_info.get("creator"):
+        result["created_by"] = decode_text(creator)
+
+    if producer := pdf_info.get("producer"):
+        producer_str = decode_text(producer)
+        if "created_by" not in result:
+            result["created_by"] = producer_str
+        elif producer_str not in result["created_by"]:
+            result["created_by"] = f"{result['created_by']} (Producer: {producer_str})"
+
+
+def _extract_document_dimensions(document: Document, result: Metadata) -> None:
+    first_page = document.pages[0]
+    if hasattr(first_page, "width") and hasattr(first_page, "height"):
+        result["width"] = int(first_page.width)
+        result["height"] = int(first_page.height)
+
+
+def _format_outline(entries: list[Any], level: int = 0) -> list[str]:
+    outline_text: list[str] = []
+    for entry in entries:
+        if hasattr(entry, "title") and entry.title:
+            indent = "  " * level
+            outline_text.append(f"{indent}- {entry.title}")
+        if hasattr(entry, "children") and entry.children:
+            _format_outline(entry.children, level + 1)
+
+    return outline_text
+
+
+def _generate_outline_description(document: Document) -> str:
+    if outline_text := _format_outline(cast("list[Any]", document.outline)):
+        return "Table of Contents:\n" + "\n".join(outline_text)
+    return ""
+
+
+def _generate_document_summary(document: Document) -> str:
+    summary_parts = []
+
+    page_count = len(document.pages)
+    summary_parts.append(f"PDF document with {page_count} page{'s' if page_count != 1 else ''}.")
+
+    if hasattr(document, "pdf_version"):
+        summary_parts.append(f"PDF version {document.pdf_version}.")
+
+    if hasattr(document, "is_encrypted") and document.is_encrypted:
+        summary_parts.append("Document is encrypted.")
+
+        if hasattr(document, "encryption_method") and document.encryption_method:
+            summary_parts.append(f"Encryption: {document.encryption_method}.")
+
+    permissions = _collect_document_permissions(document)
+    if permissions:
+        summary_parts.append(f"Document is {', '.join(permissions)}.")
+
+    if hasattr(document, "status") and document.status:
+        status = decode_text(document.status)
+        summary_parts.append(f"Status: {status}.")
+
+    if hasattr(document, "is_pdf_a") and document.is_pdf_a:
+        if hasattr(document, "pdf_a_level") and document.pdf_a_level:
+            summary_parts.append(f"PDF/A-{document.pdf_a_level} compliant.")
+        else:
+            summary_parts.append("PDF/A compliant.")
+
+    return " ".join(summary_parts)
+
+
+def _collect_document_permissions(document: Document) -> list[str]:
+    permissions = []
+    if document.is_printable:
+        permissions.append("printable")
+    if document.is_modifiable:
+        permissions.append("modifiable")
+    if document.is_extractable:
+        permissions.append("extractable")
+    return permissions
+
+
+def _extract_structure_information(document: Document, result: Metadata) -> None:
+    """Extract language and subtitle from document structure."""
+    if document.structure:
+        languages = set()
+        subtitle = None
+
+        def extract_languages(elements: list[Any]) -> None:
+            nonlocal subtitle
+            for element in elements:
+                if hasattr(element, "language") and element.language:
+                    languages.add(element.language.lower())
+
+                if (
+                    subtitle is None
+                    and hasattr(element, "role")
+                    and element.role == "H1"
+                    and hasattr(element, "text")
+                    and element.text
+                ):
+                    subtitle = decode_text(element.text)
+
+                if hasattr(element, "children") and element.children:
+                    extract_languages(element.children)
+
+        extract_languages(cast("list[Any]", document.structure))
+
+        if languages:
+            result["languages"] = list(languages)
+
+        if subtitle and "title" in result and subtitle != result["title"]:
+            result["subtitle"] = subtitle