doc-page-extractor 0.1.1__py3-none-any.whl → 1.0.2__py3-none-any.whl

doc_page_extractor/__init__.py

@@ -1,15 +1,6 @@
-from .extractor import DocExtractor
-from .clipper import clip, clip_from_image
+from .extractor import Layout, PageExtractor
+from .model import DeepSeekOCRSize
 from .plot import plot
-from .rectangle import Point, Rectangle
-from .types import (
-  ExtractedResult,
-  OCRFragment,
-  LayoutClass,
-  TableLayoutParsedFormat,
-  Layout,
-  BaseLayout,
-  PlainLayout,
-  FormulaLayout,
-  TableLayout,
-)
+
+__version__ = "1.0.0"
+__all__ = ["DeepSeekOCRSize", "Layout", "PageExtractor", "plot"]

doc_page_extractor/check_env.py

@@ -0,0 +1,40 @@
+import warnings
+
+import torch
+
+_env_checked = False
+
+
+def check_env() -> None:
+    global _env_checked
+    if _env_checked:
+        return
+    _env_checked = True
+
+    if torch.cuda.is_available():
+        return
+    warnings.warn(
+        """
+━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
+  CUDA is not available!
+━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
+
+  This package requires CUDA to run, but torch.cuda.is_available() returned False.
+
+  Possible causes:
+  1. You installed CPU-only PyTorch. Reinstall with CUDA support:
+     pip uninstall torch torchvision
+     pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121
+
+  2. Your NVIDIA GPU driver is outdated. Update it from:
+     https://www.nvidia.com/download/index.aspx
+
+  3. You don't have a CUDA-compatible GPU.
+
+  To verify your setup, run: nvidia-smi
+
+━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
+        """.strip(),
+        RuntimeWarning,
+        stacklevel=2,
+    )

doc_page_extractor/extractor.py

@@ -1,213 +1,88 @@
-import os
-
-from typing import Literal, Generator
+import tempfile
+from dataclasses import dataclass
+from os import PathLike
 from pathlib import Path
-from PIL.Image import Image
-from doclayout_yolo import YOLOv10
-
-from .ocr import OCR
-from .ocr_corrector import correct_fragments
-from .raw_optimizer import RawOptimizer
-from .rectangle import intersection_area, Rectangle
-from .downloader import download
-from .table import Table
-from .latex import LaTeX
-from .layout_order import LayoutOrder
-from .overlap import merge_fragments_as_line, remove_overlap_layouts
-from .clipper import clip_from_image
-from .types import (
-  ExtractedResult,
-  OCRFragment,
-  TableLayoutParsedFormat,
-  Layout,
-  LayoutClass,
-  PlainLayout,
-  TableLayout,
-  FormulaLayout,
-)
-
-
-class DocExtractor:
-  def __init__(
-      self,
-      model_dir_path: str,
-      device: Literal["cpu", "cuda"] = "cpu",
-      ocr_for_each_layouts: bool = True,
-      extract_formula: bool = True,
-      extract_table_format: TableLayoutParsedFormat | None = None,
-    ):
-    self._model_dir_path: str = model_dir_path
-    self._device: Literal["cpu", "cuda"] = device
-    self._ocr_for_each_layouts: bool = ocr_for_each_layouts
-    self._extract_formula: bool = extract_formula
-    self._extract_table_format: TableLayoutParsedFormat | None = extract_table_format
-    self._yolo: YOLOv10 | None = None
-    self._ocr: OCR = OCR(
-      device=device,
-      model_dir_path=os.path.join(model_dir_path, "onnx_ocr"),
-    )
-    self._table: Table = Table(
-      device=device,
-      model_path=os.path.join(model_dir_path, "struct_eqtable"),
-    )
-    self._latex: LaTeX = LaTeX(
-      model_path=os.path.join(model_dir_path, "latex"),
-    )
-    self._layout_order: LayoutOrder = LayoutOrder(
-      model_path=os.path.join(model_dir_path, "layoutreader"),
-    )
-
-  def extract(
-      self,
-      image: Image,
-      adjust_points: bool = False,
-    ) -> ExtractedResult:
-
-    raw_optimizer = RawOptimizer(image, adjust_points)
-    fragments = list(self._ocr.search_fragments(raw_optimizer.image_np))
-    raw_optimizer.receive_raw_fragments(fragments)
-    layouts = list(self._yolo_extract_layouts(raw_optimizer.image))
-    layouts = self._layouts_matched_by_fragments(fragments, layouts)
-    layouts = remove_overlap_layouts(layouts)
-
-    if self._ocr_for_each_layouts:
-      self._correct_fragments_by_ocr_layouts(raw_optimizer.image, layouts)
-
-    layouts = self._layout_order.sort(layouts, raw_optimizer.image.size)
-    layouts = [layout for layout in layouts if self._should_keep_layout(layout)]
-
-    self._parse_table_and_formula_layouts(layouts, raw_optimizer)
-
-    for layout in layouts:
-      layout.fragments = merge_fragments_as_line(layout.fragments)
-
-    raw_optimizer.receive_raw_layouts(layouts)
-
-    return ExtractedResult(
-      rotation=raw_optimizer.rotation,
-      layouts=layouts,
-      extracted_image=image,
-      adjusted_image=raw_optimizer.adjusted_image,
-    )
-
-  def _yolo_extract_layouts(self, source: Image) -> Generator[Layout, None, None]:
-    # about source parameter to see:
-    # https://github.com/opendatalab/DocLayout-YOLO/blob/7c4be36bc61f11b67cf4a44ee47f3c41e9800a91/doclayout_yolo/data/build.py#L157-L175
-    det_res = self._get_yolo().predict(
-      source=source,
-      imgsz=1024,
-      conf=0.2,
-      device=self._device    # Device to use (e.g., "cuda" or "cpu")
-    )
-    boxes = det_res[0].__dict__["boxes"]
-
-    for cls_id, rect in zip(boxes.cls, boxes.xyxy):
-      cls_id = cls_id.item()
-      cls=LayoutClass(round(cls_id))
-
-      x1, y1, x2, y2 = rect
-      x1 = x1.item()
-      y1 = y1.item()
-      x2 = x2.item()
-      y2 = y2.item()
-      rect = Rectangle(
-        lt=(x1, y1),
-        rt=(x2, y1),
-        lb=(x1, y2),
-        rb=(x2, y2),
-      )
-      if rect.is_valid:
-        if cls == LayoutClass.TABLE:
-          yield TableLayout(cls=cls, rect=rect, fragments=[], parsed=None)
-        elif cls == LayoutClass.ISOLATE_FORMULA:
-          yield FormulaLayout(cls=cls, rect=rect, fragments=[], latex=None)
-        else:
-          yield PlainLayout(cls=cls, rect=rect, fragments=[])
-
-  def _layouts_matched_by_fragments(self, fragments: list[OCRFragment], layouts: list[Layout]):
-    layouts_group = self._split_layouts_by_group(layouts)
-    for fragment in fragments:
-      for sub_layouts in layouts_group:
-        layout = self._find_matched_layout(fragment, sub_layouts)
-        if layout is not None:
-          layout.fragments.append(fragment)
-          break
-    return layouts
-
-  def _correct_fragments_by_ocr_layouts(self, source: Image, layouts: list[Layout]):
-    for layout in layouts:
-      correct_fragments(self._ocr, source, layout)
-
-  def _parse_table_and_formula_layouts(self, layouts: list[Layout], raw_optimizer: RawOptimizer):
-    for layout in layouts:
-      if isinstance(layout, FormulaLayout) and self._extract_formula:
-        image = clip_from_image(raw_optimizer.image, layout.rect)
-        layout.latex = self._latex.extract(image)
-      elif isinstance(layout, TableLayout) and self._extract_table_format is not None:
-        image = clip_from_image(raw_optimizer.image, layout.rect)
-        parsed = self._table.predict(image, self._extract_table_format)
-        if parsed is not None:
-          layout.parsed = (parsed, self._extract_table_format)
-
-  def _split_layouts_by_group(self, layouts: list[Layout]):
-    texts_layouts: list[Layout] = []
-    abandon_layouts: list[Layout] = []
-
-    for layout in layouts:
-      cls = layout.cls
-      if cls == LayoutClass.TITLE or \
-         cls == LayoutClass.PLAIN_TEXT or \
-         cls == LayoutClass.FIGURE_CAPTION or \
-         cls == LayoutClass.TABLE_CAPTION or \
-         cls == LayoutClass.TABLE_FOOTNOTE or \
-         cls == LayoutClass.FORMULA_CAPTION:
-        texts_layouts.append(layout)
-      elif cls == LayoutClass.ABANDON:
-        abandon_layouts.append(layout)
-
-    return texts_layouts, abandon_layouts
-
-  def _find_matched_layout(self, fragment: OCRFragment, layouts: list[Layout]) -> Layout | None:
-    fragment_area = fragment.rect.area
-    primary_layouts: list[(Layout, float)] = []
-
-    if fragment_area == 0.0:
-      return None
-
-    for layout in layouts:
-      area = intersection_area(fragment.rect, layout.rect)
-      if area / fragment_area > 0.85:
-        primary_layouts.append((layout, layout.rect.area))
-
-    min_area: float = float("inf")
-    min_layout: Layout | None = None
-
-    for layout, area in primary_layouts:
-      if area < min_area:
-        min_area = area
-        min_layout = layout
-
-    return min_layout
-
-  def _get_yolo(self) -> YOLOv10:
-    if self._yolo is None:
-      base_path = os.path.join(self._model_dir_path, "yolo")
-      os.makedirs(base_path, exist_ok=True)
-      yolo_model_url = "https://huggingface.co/opendatalab/PDF-Extract-Kit-1.0/resolve/main/models/Layout/YOLO/doclayout_yolo_ft.pt"
-      yolo_model_name = "doclayout_yolo_ft.pt"
-      yolo_model_path = Path(os.path.join(base_path, yolo_model_name))
-      if not yolo_model_path.exists():
-        download(yolo_model_url, yolo_model_path)
-      self._yolo = YOLOv10(str(yolo_model_path))
-    return self._yolo
-
-  def _should_keep_layout(self, layout: Layout) -> bool:
-    if len(layout.fragments) > 0:
-      return True
-    cls = layout.cls
-    return (
-      cls == LayoutClass.FIGURE or
-      cls == LayoutClass.TABLE or
-      cls == LayoutClass.ISOLATE_FORMULA
-    )
-
+from typing import Generator, cast
+
+from PIL import Image
+
+from .check_env import check_env
+from .model import DeepSeekOCRModel, DeepSeekOCRSize
+from .parser import ParsedItemKind, parse_ocr_response
+from .redacter import background_color, redact
+
+
+@dataclass
+class Layout:
+    ref: str
+    det: tuple[int, int, int, int]
+    text: str | None
+
+
+class PageExtractor:
+    def __init__(
+        self,
+        model_path: PathLike | None = None,
+        local_only: bool = False,
+    ) -> None:
+        self._model: DeepSeekOCRModel = DeepSeekOCRModel(
+            model_path=Path(model_path) if model_path else None,
+            local_only=local_only,
+        )
+
+    def download_models(self) -> None:
+        self._model.download()
+
+    def load_models(self) -> None:
+        self._model.load()
+
+    def extract(
+        self, image: Image.Image, size: DeepSeekOCRSize, stages: int = 1
+    ) -> Generator[tuple[Image.Image, list[Layout]], None, None]:
+        check_env()
+        assert stages >= 1, "stages must be at least 1"
+        with tempfile.TemporaryDirectory() as temp_path:
+            fill_color: tuple[int, int, int] | None = None
+            for i in range(stages):
+                response = self._model.generate(
+                    image=image,
+                    prompt="<image>\n<|grounding|>Convert the document to markdown.",
+                    temp_path=temp_path,
+                    size=size,
+                )
+                layouts: list[Layout] = []
+                for ref, det, text in self._parse_response(image, response):
+                    layouts.append(Layout(ref, det, text))
+                yield image, layouts
+                if i < stages - 1:
+                    if fill_color is None:
+                        fill_color = background_color(image)
+                    image = redact(
+                        image=image.copy(),
+                        fill_color=fill_color,
+                        rectangles=(layout.det for layout in layouts),
+                    )
+
+    def _parse_response(
+        self, image: Image.Image, response: str
+    ) -> Generator[tuple[str, tuple[int, int, int, int], str | None], None, None]:
+        width, height = image.size
+        det: tuple[int, int, int, int] | None = None
+        ref: str | None = None
+
+        for kind, content in parse_ocr_response(response, width, height):
+            if kind == ParsedItemKind.TEXT:
+                if det is not None and ref is not None:
+                    yield ref, det, cast(str, content)
+                    det = None
+                    ref = None
+            if det is not None and ref is not None:
+                yield ref, det, None
+                det = None
+                ref = None
+            elif kind == ParsedItemKind.DET:
+                det = cast(tuple[int, int, int, int], content)
+            elif kind == ParsedItemKind.REF:
+                ref = cast(str, content)
+        if det is not None and ref is not None:
+            yield ref, det, None

doc_page_extractor/model.py

@@ -0,0 +1,97 @@
+import os
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Literal
+
+import torch
+from huggingface_hub import snapshot_download
+from PIL import Image
+from transformers import AutoModel, AutoTokenizer
+
+DeepSeekOCRSize = Literal["tiny", "small", "base", "large", "gundam"]
+
+
+@dataclass
+class _SizeConfig:
+    base_size: int
+    image_size: int
+    crop_mode: bool
+
+
+_SIZE_CONFIGS: dict[DeepSeekOCRSize, _SizeConfig] = {
+    "tiny": _SizeConfig(base_size=512, image_size=512, crop_mode=False),
+    "small": _SizeConfig(base_size=640, image_size=640, crop_mode=False),
+    "base": _SizeConfig(base_size=1024, image_size=1024, crop_mode=False),
+    "large": _SizeConfig(base_size=1280, image_size=1280, crop_mode=False),
+    "gundam": _SizeConfig(base_size=1024, image_size=640, crop_mode=True),
+}
+
+_ATTN_IMPLEMENTATION: str
+try:
+    import flash_attn  # type: ignore # pylint: disable=unused-import
+
+    _ATTN_IMPLEMENTATION = "flash_attention_2"
+except ImportError:
+    _ATTN_IMPLEMENTATION = "eager"
+
+_Models = tuple[Any, Any]
+
+
+class DeepSeekOCRModel:
+    def __init__(self, model_path: Path | None, local_only: bool) -> None:
+        self._model_name = "deepseek-ai/DeepSeek-OCR"
+        self._cache_dir = str(model_path) if model_path else None
+        self._local_only = local_only
+        self._models: _Models | None = None
+
+    def download(self) -> None:
+        snapshot_download(
+            repo_id=self._model_name,
+            repo_type="model",
+            cache_dir=self._cache_dir,
+        )
+
+    def load(self) -> None:
+        self._ensure_models()
+
+    def _ensure_models(self) -> _Models:
+        if self._models is None:
+            tokenizer = AutoTokenizer.from_pretrained(
+                self._model_name,
+                trust_remote_code=True,
+                cache_dir=self._cache_dir,
+                local_files_only=self._local_only,
+            )
+            model = AutoModel.from_pretrained(
+                pretrained_model_name_or_path=self._model_name,
+                _attn_implementation=_ATTN_IMPLEMENTATION,
+                trust_remote_code=True,
+                use_safetensors=True,
+                cache_dir=self._cache_dir,
+                local_files_only=self._local_only,
+            )
+            model = model.cuda().to(torch.bfloat16)
+            self._models = (tokenizer, model)
+
+        return self._models
+
+    def generate(
+        self, image: Image.Image, prompt: str, temp_path: str, size: DeepSeekOCRSize
+    ) -> str:
+        tokenizer, model = self._ensure_models()
+        config = _SIZE_CONFIGS[size]
+        temp_image_path = os.path.join(temp_path, "temp_image.png")
+        image.save(temp_image_path)
+        text_result = model.infer(
+            tokenizer,
+            prompt=prompt,
+            image_file=temp_image_path,
+            output_path=temp_path,
+            base_size=config.base_size,
+            image_size=config.image_size,
+            crop_mode=config.crop_mode,
+            save_results=True,
+            test_compress=True,
+            eval_mode=True,
+        )
+        return text_result

doc_page_extractor/parser.py

@@ -0,0 +1,51 @@
+import re
+from enum import Enum, auto
+from typing import Generator
+
+_TAG_PATTERN = re.compile(r"<\|(det|ref)\|>(.+?)<\|/\1\|>")
+_DET_COORDS_PATTERN = re.compile(r"\[\[(\d+),\s*(\d+),\s*(\d+),\s*(\d+)\]\]")
+
+
+class ParsedItemKind(Enum):
+    DET = auto()
+    REF = auto()
+    TEXT = auto()
+
+
+ParsedItem = (
+    tuple[ParsedItemKind.DET, tuple[int, int, int, int]]
+    | tuple[ParsedItemKind.REF, str]
+    | tuple[ParsedItemKind.TEXT, str]
+)
+
+
+def parse_ocr_response(
+    response: str, width: int, height: int
+) -> Generator[ParsedItem, None, None]:
+    last_end: int = 0
+    for matched in _TAG_PATTERN.finditer(response):
+        if matched.start() > last_end:
+            plain_text = response[last_end : matched.start()]
+            if plain_text:
+                yield ParsedItemKind.TEXT, plain_text
+        tag_type = matched.group(1)
+        content = matched.group(2)
+        if tag_type == "det":
+            coords_match = _DET_COORDS_PATTERN.search(content)
+            if coords_match:
+                x1_norm, y1_norm, x2_norm, y2_norm = [
+                    int(c) for c in coords_match.groups()
+                ]
+                x1 = round(x1_norm / 1000 * width)
+                y1 = round(y1_norm / 1000 * height)
+                x2 = round(x2_norm / 1000 * width)
+                y2 = round(y2_norm / 1000 * height)
+                yield ParsedItemKind.DET, (x1, y1, x2, y2)
+        elif tag_type == "ref":
+            yield ParsedItemKind.REF, content
+        last_end = matched.end()
+
+    if last_end < len(response):
+        plain_text = response[last_end:]
+        if plain_text:
+            yield ParsedItemKind.TEXT, plain_text

doc_page_extractor/plot.py

@@ -1,91 +1,64 @@
-from typing import Iterable
+from typing import Iterable, cast
+
 from PIL import ImageDraw
-from PIL.ImageFont import load_default, FreeTypeFont
 from PIL.Image import Image
-from .types import Layout, LayoutClass
-from .rectangle import Point
+from PIL.ImageFont import FreeTypeFont, load_default
+
+from .extractor import Layout
 
-_FRAGMENT_COLOR = (0x49, 0xCF, 0xCB) # Light Green
+_FRAGMENT_COLOR = (0x49, 0xCF, 0xCB)  # Light Green
 _Color = tuple[int, int, int]
 
-def plot(image: Image, layouts: Iterable[Layout]) -> None:
-  layout_font = load_default(size=35)
-  fragment_font = load_default(size=25)
-  draw = ImageDraw.Draw(image, mode="RGBA")
 
-  def _draw_number(position: Point, number: int, font: FreeTypeFont, bold: bool, color: _Color) -> None:
-    nonlocal draw
-    x, y = position
-    text = str(object=number)
-    width = len(text) * font.size
-    offset = round(font.size * 0.15)
+def plot(image: Image, layouts: Iterable[Layout]) -> Image:
+    layout_font = cast(FreeTypeFont, load_default(size=35))
+    draw = ImageDraw.Draw(image, mode="RGBA")
 
-    for dx, dy in _generate_delta(bold):
-      draw.text(
-        xy=(x + dx - width - offset, y + dy),
-        text=text,
-        font=font,
-        fill=color,
-      )
+    def _draw_text(
+        position: tuple[int, int],
+        text: str,
+        font: FreeTypeFont,
+        bold: bool,
+        color: _Color,
+    ) -> None:
+        nonlocal draw
+        x, y = position
+        bbox = font.getbbox(text)
+        text_width = bbox[2] - bbox[0]
+        offset = round(font.size * 0.15)
 
-  for layout in layouts:
-    draw.polygon(
-      xy=[p for p in layout.rect],
-      outline=_layout_color(layout),
-      width=5,
-    )
+        for dx, dy in _generate_delta(bold):
+            draw.text(
+                xy=(x + dx - text_width - offset, y + dy),
+                text=text,
+                font=font,
+                fill=color,
+            )
 
-  for layout in layouts:
-    for fragment in layout.fragments:
-      draw.polygon(
-        xy=[p for p in fragment.rect],
-        outline=_FRAGMENT_COLOR,
-        width=3,
-      )
-      _draw_number(
-        position=fragment.rect.lt,
-        number=fragment.order + 1,
-        font=fragment_font,
-        bold=False,
-        color=_FRAGMENT_COLOR,
-      )
+    for layout in layouts:
+        x1, y1, x2, y2 = layout.det
+        draw.polygon(
+            xy=[(x1, y1), (x2, y1), (x2, y2), (x1, y2)],
+            outline=_FRAGMENT_COLOR,
+            width=5,
+        )
 
-  for i, layout in enumerate(layouts):
-    _draw_number(
-      position=layout.rect.lt,
-      number=i + 1,
-      font=layout_font,
-      bold=True,
-      color=_layout_color(layout),
-    )
+    for i, layout in enumerate(layouts):
+        x1, y1, _, _ = layout.det
+        _draw_text(
+            position=(x1, y1),
+            text=f"{i + 1}. {layout.ref.strip()}",
+            font=layout_font,
+            bold=True,
+            color=_FRAGMENT_COLOR,
+        )
+    return image
 
-def _generate_delta(bold: bool):
-  if bold:
-    for dx in range(-1, 2):
-      for dy in range(-1, 2):
-        yield dx, dy
-  else:
-    yield 0, 0
 
-def _layout_color(layout: Layout) -> _Color:
-  cls = layout.cls
-  if cls == LayoutClass.TITLE:
-    return (0x0A, 0x12, 0x2C) # Dark
-  elif cls == LayoutClass.PLAIN_TEXT:
-    return (0x3C, 0x67, 0x90) # Blue
-  elif cls == LayoutClass.ABANDON:
-    return (0xC0, 0xBB, 0xA9) # Gray
-  elif cls == LayoutClass.FIGURE:
-    return (0x5B, 0x91, 0x3C) # Dark Green
-  elif cls == LayoutClass.FIGURE_CAPTION:
-    return (0x77, 0xB3, 0x54) # Green
-  elif cls == LayoutClass.TABLE:
-    return (0x44, 0x17, 0x52) # Dark Purple
-  elif cls == LayoutClass.TABLE_CAPTION:
-    return (0x81, 0x75, 0xA0) # Purple
-  elif cls == LayoutClass.TABLE_FOOTNOTE:
-    return (0xEF, 0xB6, 0xC9) # Pink Purple
-  elif cls == LayoutClass.ISOLATE_FORMULA:
-    return (0xFA, 0x38, 0x27) # Red
-  elif cls == LayoutClass.FORMULA_CAPTION:
-    return (0xFF, 0x9D, 0x24) # Orange
+def _generate_delta(bold: bool):
+    if bold:
+        for dx in range(-1, 2):
+            for dy in range(-1, 2):
+                yield dx, dy
+    else:
+        yield 0, 0