docling 2.69.0__py3-none-any.whl

docling/utils/model_downloader.py

@@ -0,0 +1,168 @@
+import logging
+from pathlib import Path
+from typing import Optional
+
+from docling.datamodel.layout_model_specs import DOCLING_LAYOUT_V2
+from docling.datamodel.pipeline_options import (
+    LayoutOptions,
+    granite_picture_description,
+    smolvlm_picture_description,
+)
+from docling.datamodel.settings import settings
+from docling.datamodel.vlm_model_specs import (
+    GRANITEDOCLING_MLX,
+    GRANITEDOCLING_TRANSFORMERS,
+    SMOLDOCLING_MLX,
+    SMOLDOCLING_TRANSFORMERS,
+)
+from docling.models.stages.code_formula.code_formula_model import CodeFormulaModel
+from docling.models.stages.layout.layout_model import LayoutModel
+from docling.models.stages.ocr.easyocr_model import EasyOcrModel
+from docling.models.stages.ocr.rapid_ocr_model import RapidOcrModel
+from docling.models.stages.picture_classifier.document_picture_classifier import (
+    DocumentPictureClassifier,
+    DocumentPictureClassifierOptions,
+)
+from docling.models.stages.picture_description.picture_description_vlm_model import (
+    PictureDescriptionVlmModel,
+)
+from docling.models.stages.table_structure.table_structure_model import (
+    TableStructureModel,
+)
+from docling.models.utils.hf_model_download import download_hf_model
+
+_log = logging.getLogger(__name__)
+
+
+def download_models(
+    output_dir: Optional[Path] = None,
+    *,
+    force: bool = False,
+    progress: bool = False,
+    with_layout: bool = True,
+    with_tableformer: bool = True,
+    with_code_formula: bool = True,
+    with_picture_classifier: bool = True,
+    with_smolvlm: bool = False,
+    with_granitedocling: bool = False,
+    with_granitedocling_mlx: bool = False,
+    with_smoldocling: bool = False,
+    with_smoldocling_mlx: bool = False,
+    with_granite_vision: bool = False,
+    with_rapidocr: bool = True,
+    with_easyocr: bool = False,
+):
+    if output_dir is None:
+        output_dir = settings.cache_dir / "models"
+
+    # Make sure the folder exists
+    output_dir.mkdir(exist_ok=True, parents=True)
+
+    if with_layout:
+        _log.info("Downloading layout model...")
+        LayoutModel.download_models(
+            local_dir=output_dir / LayoutOptions().model_spec.model_repo_folder,
+            force=force,
+            progress=progress,
+        )
+
+    if with_tableformer:
+        _log.info("Downloading tableformer model...")
+        TableStructureModel.download_models(
+            local_dir=output_dir / TableStructureModel._model_repo_folder,
+            force=force,
+            progress=progress,
+        )
+
+    if with_picture_classifier:
+        _log.info("Downloading picture classifier model...")
+        pic_opts = DocumentPictureClassifierOptions()
+        DocumentPictureClassifier.download_models(
+            repo_id=pic_opts.repo_id,
+            revision=pic_opts.revision,
+            local_dir=output_dir / pic_opts.repo_cache_folder,
+            force=force,
+            progress=progress,
+        )
+
+    if with_code_formula:
+        _log.info("Downloading code formula model...")
+        CodeFormulaModel.download_models(
+            local_dir=output_dir / CodeFormulaModel._model_repo_folder,
+            force=force,
+            progress=progress,
+        )
+
+    if with_smolvlm:
+        _log.info("Downloading SmolVlm model...")
+        download_hf_model(
+            repo_id=smolvlm_picture_description.repo_id,
+            local_dir=output_dir / smolvlm_picture_description.repo_cache_folder,
+            force=force,
+            progress=progress,
+        )
+
+    if with_granitedocling:
+        _log.info("Downloading GraniteDocling model...")
+        download_hf_model(
+            repo_id=GRANITEDOCLING_TRANSFORMERS.repo_id,
+            local_dir=output_dir / GRANITEDOCLING_TRANSFORMERS.repo_cache_folder,
+            force=force,
+            progress=progress,
+        )
+
+    if with_granitedocling_mlx:
+        _log.info("Downloading GraniteDocling MLX model...")
+        download_hf_model(
+            repo_id=GRANITEDOCLING_MLX.repo_id,
+            local_dir=output_dir / GRANITEDOCLING_MLX.repo_cache_folder,
+            force=force,
+            progress=progress,
+        )
+
+    if with_smoldocling:
+        _log.info("Downloading SmolDocling model...")
+        download_hf_model(
+            repo_id=SMOLDOCLING_TRANSFORMERS.repo_id,
+            local_dir=output_dir / SMOLDOCLING_TRANSFORMERS.repo_cache_folder,
+            force=force,
+            progress=progress,
+        )
+
+    if with_smoldocling_mlx:
+        _log.info("Downloading SmolDocling MLX model...")
+        download_hf_model(
+            repo_id=SMOLDOCLING_MLX.repo_id,
+            local_dir=output_dir / SMOLDOCLING_MLX.repo_cache_folder,
+            force=force,
+            progress=progress,
+        )
+
+    if with_granite_vision:
+        _log.info("Downloading Granite Vision model...")
+        download_hf_model(
+            repo_id=granite_picture_description.repo_id,
+            local_dir=output_dir / granite_picture_description.repo_cache_folder,
+            force=force,
+            progress=progress,
+        )
+
+    if with_rapidocr:
+        for backend in ("torch", "onnxruntime"):
+            _log.info(f"Downloading rapidocr {backend} models...")
+            RapidOcrModel.download_models(
+                backend=backend,
+                local_dir=output_dir / RapidOcrModel._model_repo_folder,
+                force=force,
+                progress=progress,
+            )
+
+    if with_easyocr:
+        _log.info("Downloading easyocr models...")
+        EasyOcrModel.download_models(
+            local_dir=output_dir / EasyOcrModel._model_repo_folder,
+            force=force,
+            progress=progress,
+        )
+
+    return output_dir

docling/utils/ocr_utils.py

@@ -0,0 +1,69 @@
+from typing import Optional, Tuple
+
+from docling_core.types.doc import BoundingBox, CoordOrigin
+from docling_core.types.doc.page import BoundingRectangle
+
+from docling.utils.orientation import CLIPPED_ORIENTATIONS, rotate_bounding_box
+
+
+def map_tesseract_script(script: str) -> str:
+    r""" """
+    if script == "Katakana" or script == "Hiragana":
+        script = "Japanese"
+    elif script == "Han":
+        script = "HanS"
+    elif script == "Korean":
+        script = "Hangul"
+    return script
+
+
+def parse_tesseract_orientation(orientation: str) -> int:
+    # Tesseract orientation is [0, 90, 180, 270] clockwise, bounding rectangle angles
+    # are [0, 360[ counterclockwise
+    parsed = int(orientation)
+    if parsed not in CLIPPED_ORIENTATIONS:
+        msg = (
+            f"invalid tesseract document orientation {orientation}, "
+            f"expected orientation: {sorted(CLIPPED_ORIENTATIONS)}"
+        )
+        raise ValueError(msg)
+    parsed = -parsed
+    parsed %= 360
+    return parsed
+
+
+def tesseract_box_to_bounding_rectangle(
+    bbox: BoundingBox,
+    *,
+    original_offset: Optional[BoundingBox] = None,
+    scale: float,
+    orientation: int,
+    im_size: Tuple[int, int],
+) -> BoundingRectangle:
+    # box is in the top, left, height, width format, top left coordinates
+    rect = rotate_bounding_box(bbox, angle=orientation, im_size=im_size)
+    rect = BoundingRectangle(
+        r_x0=rect.r_x0 / scale,
+        r_y0=rect.r_y0 / scale,
+        r_x1=rect.r_x1 / scale,
+        r_y1=rect.r_y1 / scale,
+        r_x2=rect.r_x2 / scale,
+        r_y2=rect.r_y2 / scale,
+        r_x3=rect.r_x3 / scale,
+        r_y3=rect.r_y3 / scale,
+        coord_origin=CoordOrigin.TOPLEFT,
+    )
+    if original_offset is not None:
+        if original_offset.coord_origin is not CoordOrigin.TOPLEFT:
+            msg = f"expected coordinate origin to be {CoordOrigin.TOPLEFT.value}"
+            raise ValueError(msg)
+        if original_offset is not None:
+            rect.r_x0 += original_offset.l
+            rect.r_x1 += original_offset.l
+            rect.r_x2 += original_offset.l
+            rect.r_x3 += original_offset.l
+            rect.r_y0 += original_offset.t
+            rect.r_y1 += original_offset.t
+            rect.r_y2 += original_offset.t
+            rect.r_y3 += original_offset.t
+    return rect

docling/utils/orientation.py

@@ -0,0 +1,65 @@
+from typing import Tuple
+
+from docling_core.types.doc import BoundingBox, CoordOrigin
+from docling_core.types.doc.page import BoundingRectangle
+
+CLIPPED_ORIENTATIONS = [0, 90, 180, 270]
+
+
+def rotate_bounding_box(
+    bbox: BoundingBox, angle: int, im_size: Tuple[int, int]
+) -> BoundingRectangle:
+    # The box is left top width height in TOPLEFT coordinates
+    # Bounding rectangle start with r_0 at the bottom left whatever the
+    # coordinate system. Then other corners are found rotating counterclockwise
+    bbox = bbox.to_top_left_origin(im_size[1])
+    left, top, width, height = bbox.l, bbox.t, bbox.width, bbox.height
+    im_w, im_h = im_size
+    angle = angle % 360
+    if angle == 0:
+        return BoundingRectangle.from_bounding_box(bbox)
+    elif angle == 90:
+        r_x0 = top + height
+        r_y0 = im_w - left
+        r_x1 = r_x0
+        r_y1 = r_y0 - width
+        r_x2 = r_x1 - height
+        r_y2 = r_y1
+        r_x3 = r_x2
+        r_y3 = r_y0
+    elif angle == 180:
+        r_x0 = im_w - left
+        r_y0 = im_h - (top + height)
+        r_x1 = r_x0 - width
+        r_y1 = r_y0
+        r_x2 = r_x1
+        r_y2 = r_y1 + height
+        r_x3 = r_x0
+        r_y3 = r_y2
+    elif angle == 270:
+        r_x0 = im_h - (top + height)
+        r_y0 = left
+        r_x1 = r_x0
+        r_y1 = r_y0 + width
+        r_x2 = r_x1 + height
+        r_y2 = r_y1
+        r_x3 = r_x2
+        r_y3 = r_y0
+    else:
+        msg = (
+            f"invalid orientation {angle}, expected values in:"
+            f" {sorted(CLIPPED_ORIENTATIONS)}"
+        )
+        raise ValueError(msg)
+    rectangle = BoundingRectangle(
+        r_x0=r_x0,
+        r_y0=r_y0,
+        r_x1=r_x1,
+        r_y1=r_y1,
+        r_x2=r_x2,
+        r_y2=r_y2,
+        r_x3=r_x3,
+        r_y3=r_y3,
+        coord_origin=CoordOrigin.TOPLEFT,
+    )
+    return rectangle

docling/utils/profiling.py

@@ -0,0 +1,65 @@
+import time
+from datetime import datetime
+from enum import Enum
+from typing import TYPE_CHECKING, List
+
+import numpy as np
+from pydantic import BaseModel
+
+from docling.datamodel.settings import settings
+
+if TYPE_CHECKING:
+    from docling.datamodel.document import ConversionResult
+
+
+class ProfilingScope(str, Enum):
+    PAGE = "page"
+    DOCUMENT = "document"
+
+
+class ProfilingItem(BaseModel):
+    scope: ProfilingScope
+    count: int = 0
+    times: List[float] = []
+    start_timestamps: List[datetime] = []
+
+    def total(self) -> float:
+        return np.sum(self.times)  # type: ignore
+
+    def avg(self) -> float:
+        return np.average(self.times)  # type: ignore
+
+    def std(self) -> float:
+        return np.std(self.times)  # type: ignore
+
+    def mean(self) -> float:
+        return np.mean(self.times)  # type: ignore
+
+    def percentile(self, perc: float) -> float:
+        return np.percentile(self.times, perc)  # type: ignore
+
+
+class TimeRecorder:
+    def __init__(
+        self,
+        conv_res: "ConversionResult",
+        key: str,
+        scope: ProfilingScope = ProfilingScope.PAGE,
+    ):
+        if settings.debug.profile_pipeline_timings:
+            if key not in conv_res.timings.keys():
+                conv_res.timings[key] = ProfilingItem(scope=scope)
+            self.conv_res = conv_res
+            self.key = key
+
+    def __enter__(self):
+        if settings.debug.profile_pipeline_timings:
+            self.start = time.monotonic()
+            self.conv_res.timings[self.key].start_timestamps.append(datetime.utcnow())
+        return self
+
+    def __exit__(self, *args):
+        if settings.debug.profile_pipeline_timings:
+            elapsed = time.monotonic() - self.start
+            self.conv_res.timings[self.key].times.append(elapsed)
+            self.conv_res.timings[self.key].count += 1

docling/utils/utils.py

@@ -0,0 +1,65 @@
+import hashlib
+from io import BytesIO
+from itertools import islice
+from pathlib import Path
+from typing import List, Union
+
+import requests
+from tqdm import tqdm
+
+
+def chunkify(iterator, chunk_size):
+    """Yield successive chunks of chunk_size from the iterable."""
+    if isinstance(iterator, List):
+        iterator = iter(iterator)
+    for first in iterator:  # Take the first element from the iterator
+        yield [first, *list(islice(iterator, chunk_size - 1))]
+
+
+def create_file_hash(path_or_stream: Union[BytesIO, Path]) -> str:
+    """Create a stable page_hash of the path_or_stream of a file"""
+
+    block_size = 65536
+    hasher = hashlib.sha256(usedforsecurity=False)
+
+    def _hash_buf(binary_stream):
+        buf = binary_stream.read(block_size)  # read and page_hash in chunks
+        while len(buf) > 0:
+            hasher.update(buf)
+            buf = binary_stream.read(block_size)
+
+    if isinstance(path_or_stream, Path):
+        with path_or_stream.open("rb") as afile:
+            _hash_buf(afile)
+    elif isinstance(path_or_stream, BytesIO):
+        _hash_buf(path_or_stream)
+
+    return hasher.hexdigest()
+
+
+def create_hash(string: str):
+    hasher = hashlib.sha256(usedforsecurity=False)
+    hasher.update(string.encode("utf-8"))
+
+    return hasher.hexdigest()
+
+
+def download_url_with_progress(url: str, progress: bool = False) -> BytesIO:
+    buf = BytesIO()
+    with requests.get(url, stream=True, allow_redirects=True) as response:
+        total_size = int(response.headers.get("content-length", 0))
+        progress_bar = tqdm(
+            total=total_size,
+            unit="B",
+            unit_scale=True,
+            unit_divisor=1024,
+            disable=(not progress),
+        )
+
+        for chunk in response.iter_content(10 * 1024):
+            buf.write(chunk)
+            progress_bar.update(len(chunk))
+        progress_bar.close()
+
+    buf.seek(0)
+    return buf

docling/utils/visualization.py

@@ -0,0 +1,85 @@
+from docling_core.types.doc import DocItemLabel
+from PIL import Image, ImageDraw, ImageFont
+from PIL.ImageFont import FreeTypeFont
+
+from docling.datamodel.base_models import Cluster
+
+
+def draw_clusters(
+    image: Image.Image, clusters: list[Cluster], scale_x: float, scale_y: float
+) -> None:
+    """
+    Draw clusters on an image
+    """
+    draw = ImageDraw.Draw(image, "RGBA")
+    # Create a smaller font for the labels
+    font: ImageFont.ImageFont | FreeTypeFont
+    try:
+        font = ImageFont.truetype("arial.ttf", 12)
+    except OSError:
+        # Fallback to default font if arial is not available
+        font = ImageFont.load_default()
+    for c_tl in clusters:
+        all_clusters = [c_tl, *c_tl.children]
+        for c in all_clusters:
+            # Draw cells first (underneath)
+            cell_color = (0, 0, 0, 40)  # Transparent black for cells
+            for tc in c.cells:
+                cx0, cy0, cx1, cy1 = tc.rect.to_bounding_box().as_tuple()
+                cx0 *= scale_x
+                cx1 *= scale_x
+                cy0 *= scale_x
+                cy1 *= scale_y
+
+                draw.rectangle(
+                    [(cx0, cy0), (cx1, cy1)],
+                    outline=None,
+                    fill=cell_color,
+                )
+            # Draw cluster rectangle
+            x0, y0, x1, y1 = c.bbox.as_tuple()
+            x0 *= scale_x
+            x1 *= scale_x
+            y0 *= scale_x
+            y1 *= scale_y
+
+            if y1 <= y0:
+                y1, y0 = y0, y1
+            if x1 <= x0:
+                x1, x0 = x0, x1
+
+            cluster_fill_color = (*list(DocItemLabel.get_color(c.label)), 70)
+            cluster_outline_color = (
+                *list(DocItemLabel.get_color(c.label)),
+                255,
+            )
+            draw.rectangle(
+                [(x0, y0), (x1, y1)],
+                outline=cluster_outline_color,
+                fill=cluster_fill_color,
+            )
+            # Add label name and confidence
+            label_text = f"{c.label.name} ({c.confidence:.2f})"
+            # Create semi-transparent background for text
+            text_bbox = draw.textbbox((x0, y0), label_text, font=font)
+            text_bg_padding = 2
+            draw.rectangle(
+                [
+                    (
+                        text_bbox[0] - text_bg_padding,
+                        text_bbox[1] - text_bg_padding,
+                    ),
+                    (
+                        text_bbox[2] + text_bg_padding,
+                        text_bbox[3] + text_bg_padding,
+                    ),
+                ],
+                fill=(255, 255, 255, 180),  # Semi-transparent white
+            )
+            # Draw text
+            draw.text(
+                (x0, y0),
+                label_text,
+                fill=(0, 0, 0, 255),  # Solid black
+                font=font,
+            )