doc-page-extractor 0.2.4__cp310-cp310-macosx_15_0_arm64.whl

doc_page_extractor/__init__.py

@@ -0,0 +1,16 @@
+from .extractor import DocExtractor
+from .clipper import clip, clip_from_image
+from .plot import plot
+from .rectangle import Point, Rectangle
+from .model import Model, HuggingfaceModel
+from .types import (
+  ExtractedResult,
+  OCRFragment,
+  LayoutClass,
+  TableLayoutParsedFormat,
+  Layout,
+  BaseLayout,
+  PlainLayout,
+  FormulaLayout,
+  TableLayout,
+)

doc_page_extractor/clipper.py

@@ -0,0 +1,119 @@
+import numpy as np
+
+from math import pi, ceil, sin, cos, sqrt
+from PIL.Image import Image, Transform
+from .types import Layout, ExtractedResult
+from .rectangle import Rectangle
+from .rotation import calculate_rotation_with_rect, normal_vertical_rotation
+
+
+def clip(
+    extracted_result: ExtractedResult,
+    layout: Layout,
+    wrapped_width: float = 0.0,
+    wrapped_height: float = 0.0,
+  ) -> Image:
+  image: Image | None
+  if extracted_result.adjusted_image is None:
+    image = extracted_result.extracted_image
+  else:
+    image = extracted_result.adjusted_image
+  assert image is not None, "Image must not be None"
+  return clip_from_image(
+    image, layout.rect,
+    wrapped_width, wrapped_height,
+  )
+
+def clip_from_image(
+    image: Image,
+    rect: Rectangle,
+    wrapped_width: float = 0.0,
+    wrapped_height: float = 0.0,
+  ) -> Image:
+  horizontal_rotation, vertical_rotation = calculate_rotation_with_rect(rect)
+  image = image.copy()
+  matrix_move = np.array(_get_move_matrix(rect.lt[0], rect.lt[1])).reshape(3, 3)
+  matrix_rotate = np.array(_get_rotate_matrix(-horizontal_rotation)).reshape(3, 3)
+  matrix = np.dot(matrix_move, matrix_rotate)
+
+  y_axis_rotation = normal_vertical_rotation(vertical_rotation - horizontal_rotation)
+
+  if abs(y_axis_rotation - 0.25 * pi) > 0.0:
+    x = cos(y_axis_rotation)
+    y = sin(y_axis_rotation)
+    matrix_shear = np.array(_get_shear_matrix(x, y)).reshape(3, 3)
+    matrix = np.dot(matrix, matrix_shear)
+
+  width, height, max_width, max_height = _size_and_wrapper(rect)
+  max_width += wrapped_width
+  max_height += wrapped_height
+
+  if max_width != width or max_height != height:
+    dx = (max_width - width) / 2.0
+    dy = (max_height - height) / 2.0
+    matrix_move = np.array(_get_move_matrix(-dx, -dy)).reshape(3, 3)
+    matrix = np.dot(matrix, matrix_move)
+
+  return image.transform(
+    size=(ceil(max_width), ceil(max_height)),
+    method=Transform.AFFINE,
+    data=_to_pillow_matrix(matrix),
+  )
+
+def _size_and_wrapper(rect: Rectangle):
+  widths: list[float] = []
+  heights: list[float] = []
+
+  for i, (p1, p2) in enumerate(rect.segments):
+    dx = p2[0] - p1[0]
+    dy = p2[1] - p1[1]
+    distance = sqrt(dx*dx + dy*dy)
+    if i % 2 == 0:
+      heights.append(distance)
+    else:
+      widths.append(distance)
+
+  if len(widths) == 0 and len(heights) == 0:
+    return 0.0, 0.0, 0.0, 0.0
+
+  width: float = sum(widths) / len(widths)
+  height: float = sum(heights) / len(heights)
+  max_width: float = width
+  max_height: float = height
+
+  for width in widths:
+    if width > max_width:
+      max_width = width
+
+  for height in heights:
+    if height > max_height:
+      max_height = height
+
+  return width, height, max_width, max_height
+
+def _to_pillow_matrix(matrix):
+  return (
+    matrix[0][0], matrix[0][1], matrix[0][2],
+    matrix[1][0], matrix[1][1], matrix[1][2],
+  )
+
+def _get_move_matrix(dx: float, dy: float):
+  return (
+    1.0, 0.0, dx,
+    0.0, 1.0, dy,
+    0.0, 0.0, 1.0,
+  )
+
+def _get_rotate_matrix(rotation: float):
+  return (
+    cos(rotation),  sin(rotation),  0.0,
+    -sin(rotation), cos(rotation),  0.0,
+    0.0,            0.0,            1.0
+  )
+
+def _get_shear_matrix(x0: float, y0: float):
+  return (
+    1.0, 0.0, 0.0,
+    x0,  y0,  0.0,
+    0.0, 0.0, 1.0,
+  )

doc_page_extractor/downloader.py

@@ -0,0 +1,16 @@
+import os
+import requests
+from pathlib import Path
+
+
+def download(url: str, file_path: Path):
+  response = requests.get(url, stream=True, timeout=60)
+  if response.status_code != 200:
+    raise FileNotFoundError(f"Failed to download file from {url}: {response.status_code}")
+  try:
+    with open(file_path, "wb") as file:
+      file.write(response.content)
+  except Exception as e:
+    if os.path.exists(file_path):
+      os.remove(file_path)
+    raise e

doc_page_extractor/extractor.py

@@ -0,0 +1,218 @@
+import torch
+
+from os import PathLike
+from typing import cast, Any, Literal, Generator
+from PIL.Image import Image
+from doclayout_yolo import YOLOv10
+
+from .model import Model, HuggingfaceModel
+from .ocr import OCR
+from .ocr_corrector import correct_fragments
+from .raw_optimizer import RawOptimizer
+from .rectangle import intersection_area, Rectangle
+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,
+  Layout,
+  LayoutClass,
+  PlainLayout,
+  TableLayout,
+  FormulaLayout,
+  TableLayoutParsedFormat
+)
+
+
+class DocExtractor:
+  def __init__(
+        self,
+        model_cache_dir: PathLike | None = None,
+        device: Literal["cpu", "cuda"] = "cpu",
+        model: Model | None = None,
+      ) -> None:
+
+    if model is None:
+      if model_cache_dir is None:
+        raise ValueError("You must provide a model_cache_dir or a model instance.")
+      model = HuggingfaceModel(model_cache_dir)
+
+    if device == "cuda" and not torch.cuda.is_available():
+      device = "cpu"
+      print("CUDA is not available. Using CPU instead.")
+
+    self._device: Literal["cpu", "cuda"] = device
+    self._model: Model = model
+    self._yolo: YOLOv10 | None = None
+    self._ocr: OCR = OCR(device, model)
+    self._table: Table = Table(device, model)
+    self._latex: LaTeX = LaTeX(device, model)
+    self._layout_order: LayoutOrder = LayoutOrder(device, model)
+
+  def prepare_models(self):
+    self._model.get_onnx_ocr_path()
+    self._model.get_yolo_path()
+    self._model.get_layoutreader_path()
+    self._model.get_struct_eqtable_path()
+    self._model.get_latex_path()
+
+  def extract(
+      self,
+      image: Image,
+      extract_formula: bool,
+      extract_table_format: TableLayoutParsedFormat | None = None,
+      ocr_for_each_layouts: bool = False,
+      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 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, extract_formula=extract_formula, extract_table_format=extract_table_format)
+
+    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=cast(Any, 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,
+      extract_formula: bool,
+      extract_table_format: TableLayoutParsedFormat | None,
+  ):
+    for layout in layouts:
+      if isinstance(layout, FormulaLayout) and extract_formula:
+        image = clip_from_image(raw_optimizer.image, layout.rect)
+        layout.latex = self._latex.extract(image)
+      elif isinstance(layout, TableLayout) and extract_table_format is not None:
+        image = clip_from_image(raw_optimizer.image, layout.rect)
+        parsed = self._table.predict(image, extract_table_format)
+        if parsed is not None:
+          layout.parsed = (parsed, 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[tuple[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:
+      model_path = self._model.get_yolo_path()
+      self._yolo = YOLOv10(str(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
+    )
+

doc_page_extractor/latex.py

@@ -0,0 +1,33 @@
+import os
+import torch
+
+from munch import Munch
+from pix2tex.cli import LatexOCR
+from PIL.Image import Image
+from typing import Literal
+from .utils import expand_image
+from .model import Model
+
+
+class LaTeX:
+  def __init__(self, device: Literal["cpu", "cuda"], model: Model) -> None:
+    self._model: Model = model
+    self._latex_model: LatexOCR | None = None
+    self._device: Literal["cpu", "cuda"] = device
+
+  def extract(self, image: Image) -> str | None:
+    image = expand_image(image, 0.1) # 添加边缘提高识别准确率
+    model = self._get_model()
+    with torch.no_grad():
+      return model(image)
+
+  def _get_model(self) -> LatexOCR:
+    if self._latex_model is None:
+      model_path = self._model.get_latex_path()
+      self._latex_model = LatexOCR(Munch({
+        "config": os.path.join("settings", "config.yaml"),
+        "checkpoint": os.path.join(model_path, "checkpoints", "weights.pth"),
+        "no_cuda": self._device == "cpu",
+        "no_resize": False,
+      }))
+    return self._latex_model

doc_page_extractor/layout_order.py

@@ -0,0 +1,239 @@
+import torch
+
+from typing import Generator, Literal
+from dataclasses import dataclass
+from transformers import LayoutLMv3ForTokenClassification
+
+from .types import Layout, LayoutClass
+from .model import Model
+from .layoutreader import prepare_inputs, boxes2inputs, parse_logits
+
+
+@dataclass
+class _BBox:
+  layout_index: int
+  fragment_index: int
+  virtual: bool
+  order: int
+  value: tuple[float, float, float, float]
+
+class LayoutOrder:
+  def __init__(self, device: Literal["cpu", "cuda"], model: Model):
+    self._model: Model = model
+    self._order_model: LayoutLMv3ForTokenClassification | None = None
+    self._device: Literal["cpu", "cuda"] = device
+
+  def _get_model(self) -> LayoutLMv3ForTokenClassification:
+    if self._order_model is None:
+      model_path = self._model.get_layoutreader_path()
+      self._order_model = LayoutLMv3ForTokenClassification.from_pretrained(
+        pretrained_model_name_or_path=model_path,
+        local_files_only=True,
+      ).to(device=self._device)
+    return self._order_model
+
+  def sort(self, layouts: list[Layout], size: tuple[int, int]) -> list[Layout]:
+    width, height = size
+    if width == 0 or height == 0:
+      return layouts
+
+    bbox_list = self._order_and_get_bbox_list(
+      layouts=layouts,
+      width=width,
+      height=height,
+    )
+    if bbox_list is None:
+      return layouts
+
+    return self._sort_layouts_and_fragments(layouts, bbox_list)
+
+  def _order_and_get_bbox_list(
+      self,
+      layouts: list[Layout],
+      width: int,
+      height: int,
+    ) -> list[_BBox] | None:
+
+    line_height = self._line_height(layouts)
+    bbox_list: list[_BBox] = []
+
+    for i, layout in enumerate(layouts):
+      if layout.cls == LayoutClass.PLAIN_TEXT and \
+         len(layout.fragments) > 0:
+        for j, fragment in enumerate(layout.fragments):
+          bbox_list.append(_BBox(
+            layout_index=i,
+            fragment_index=j,
+            virtual=False,
+            order=0,
+            value=fragment.rect.wrapper,
+          ))
+      else:
+        bbox_list.extend(
+          self._generate_virtual_lines(
+            layout=layout,
+            layout_index=i,
+            line_height=line_height,
+            width=width,
+            height=height,
+          ),
+        )
+
+    if len(bbox_list) > 200:
+      # https://github.com/opendatalab/MinerU/blob/980f5c8cd70f22f8c0c9b7b40eaff6f4804e6524/magic_pdf/pdf_parse_union_core_v2.py#L522
+      return None
+
+    layoutreader_size = 1000.0
+    x_scale = layoutreader_size / float(width)
+    y_scale = layoutreader_size / float(height)
+
+    for bbox in bbox_list:
+      x0, y0, x1, y1 = self._squeeze(bbox, width, height)
+      x0 = round(x0 * x_scale)
+      y0 = round(y0 * y_scale)
+      x1 = round(x1 * x_scale)
+      y1 = round(y1 * y_scale)
+      bbox.value = (x0, y0, x1, y1)
+
+    bbox_list.sort(key=lambda b: b.value) # 必须排序，乱序传入 layoutreader 会令它无法识别正确顺序
+    model = self._get_model()
+
+    with torch.no_grad():
+      inputs = boxes2inputs([list(bbox.value) for bbox in bbox_list])
+      inputs = prepare_inputs(inputs, model)
+      logits = model(**inputs).logits.cpu().squeeze(0)
+      orders = parse_logits(logits, len(bbox_list))
+
+    sorted_bbox_list = [bbox_list[i] for i in orders]
+    for i, bbox in enumerate(sorted_bbox_list):
+      bbox.order = i
+
+    return sorted_bbox_list
+
+  def _sort_layouts_and_fragments(self, layouts: list[Layout], bbox_list: list[_BBox]):
+    layout_bbox_list: list[list[_BBox]] = [[] for _ in range(len(layouts))]
+    for bbox in bbox_list:
+      layout_bbox_list[bbox.layout_index].append(bbox)
+
+    layouts_with_median_order: list[tuple[Layout, float]] = []
+    for layout_index, bbox_list in enumerate(layout_bbox_list):
+      layout = layouts[layout_index]
+      orders = [b.order for b in bbox_list] # virtual bbox 保证了 orders 不可能为空
+      median_order = self._median(orders)
+      layouts_with_median_order.append((layout, median_order))
+
+    for layout, bbox_list in zip(layouts, layout_bbox_list):
+      for bbox in bbox_list:
+        if not bbox.virtual:
+          layout.fragments[bbox.fragment_index].order = bbox.order
+      if all(not bbox.virtual for bbox in bbox_list):
+        layout.fragments.sort(key=lambda f: f.order)
+
+    layouts_with_median_order.sort(key=lambda x: x[1])
+    layouts = [layout for layout, _ in layouts_with_median_order]
+    next_fragment_order: int = 0
+
+    for layout in layouts:
+      for fragment in layout.fragments:
+        fragment.order = next_fragment_order
+        next_fragment_order += 1
+
+    return layouts
+
+  def _line_height(self, layouts: list[Layout]) -> float:
+    line_height: float = 0.0
+    count: int = 0
+    for layout in layouts:
+      for fragment in layout.fragments:
+        _, height = fragment.rect.size
+        line_height += height
+        count += 1
+    if count == 0:
+      return 10.0
+    return line_height / float(count)
+
+  def _generate_virtual_lines(
+        self,
+        layout: Layout,
+        layout_index: int,
+        line_height: float,
+        width: int,
+        height: int,
+      ) -> Generator[_BBox, None, None]:
+
+    # https://github.com/opendatalab/MinerU/blob/980f5c8cd70f22f8c0c9b7b40eaff6f4804e6524/magic_pdf/pdf_parse_union_core_v2.py#L451-L490
+    x0, y0, x1, y1 = layout.rect.wrapper
+    layout_height = y1 - y0
+    layout_weight = x1 - x0
+    lines = int(layout_height / line_height)
+
+    if layout_height <= line_height * 2:
+      yield _BBox(
+        layout_index=layout_index,
+        fragment_index=0,
+        virtual=True,
+        order=0,
+        value=(x0, y0, x1, y1),
+      )
+      return
+
+    elif layout_height <= height * 0.25 or \
+         width * 0.5 <= layout_weight or \
+         width * 0.25 < layout_weight:
+      if layout_weight > width * 0.4:
+        lines = 3
+      elif layout_weight <= width * 0.25:
+        if layout_height / layout_weight > 1.2:  # 细长的不分
+          yield _BBox(
+            layout_index=layout_index,
+            fragment_index=0,
+            virtual=True,
+            order=0,
+            value=(x0, y0, x1, y1),
+          )
+          return
+        else:  # 不细长的还是分成两行
+          lines = 2
+
+    lines = max(1, lines)
+    line_height = (y1 - y0) / lines
+    current_y = y0
+
+    for i in range(lines):
+      yield _BBox(
+        layout_index=layout_index,
+        fragment_index=i,
+        virtual=True,
+        order=0,
+        value=(x0, current_y, x1, current_y + line_height),
+      )
+      current_y += line_height
+
+  def _median(self, numbers: list[int]) -> float:
+    sorted_numbers = sorted(numbers)
+    n = len(sorted_numbers)
+
+    # 判断是奇数还是偶数个元素
+    if n % 2 == 1:
+      # 奇数情况，直接取中间的数
+      return float(sorted_numbers[n // 2])
+    else:
+      # 偶数情况，取中间两个数的平均值
+      mid1 = sorted_numbers[n // 2 - 1]
+      mid2 = sorted_numbers[n // 2]
+      return float((mid1 + mid2) / 2)
+
+  def _squeeze(self, bbox: _BBox, width: int, height: int) -> tuple[float, float, float, float]:
+    x0, y0, x1, y1 = bbox.value
+    x0 = self._squeeze_value(x0, width)
+    x1 = self._squeeze_value(x1, width)
+    y0 = self._squeeze_value(y0, height)
+    y1 = self._squeeze_value(y1, height)
+    return x0, y0, x1, y1
+
+  def _squeeze_value(self, position: float, size: int) -> float:
+    if position < 0:
+      position = 0.0
+    if position > size:
+      position = float(size)
+    return position