pdf-craft 0.0.1__py3-none-any.whl

pdf_craft/__init__.py

@@ -0,0 +1,3 @@
+from .pdf import *
+from .markdown import *
+from doc_page_extractor import PaddleLang

pdf_craft/markdown.py

@@ -0,0 +1,76 @@
+import io
+import os
+
+from hashlib import sha256
+from typing import Iterable
+from .pdf import Text, TextKind, TextBlock, AssetBlock, Block
+
+class MarkDownWriter:
+  def __init__(self, md_path: str, assets_path: str, encoding: str | None):
+    self._assets_path: str = assets_path
+    self._abs_assets_path: str = os.path.abspath(os.path.join(md_path, "..", assets_path))
+    self._file: io.TextIOWrapper = open(md_path, "w", encoding=encoding)
+    self._texts_buffer: list[Text] = []
+
+  def __enter__(self) -> "MarkDownWriter":
+    return self
+
+  def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+    self.close()
+
+  def flush(self) -> None:
+    self._file.flush()
+
+  def close(self) -> None:
+    self._close_texts_buffer()
+    self._file.close()
+
+  def write(self, block: Block) -> None:
+    if isinstance(block, TextBlock):
+      if block.kind == TextKind.TITLE:
+        self._close_texts_buffer()
+        self._file.write("# ")
+        self._write_text_contents(block.texts)
+        self._file.write("\n\n")
+      elif block.kind == TextKind.PLAIN_TEXT:
+        self._write_plain_text(block)
+    elif isinstance(block, AssetBlock):
+      self._close_texts_buffer()
+      self._write_image(block)
+
+  def _close_texts_buffer(self):
+    if len(self._texts_buffer) > 0:
+      self._write_text_contents(self._texts_buffer)
+      self._file.write("\n\n")
+      self._texts_buffer.clear()
+
+  def _write_plain_text(self, block: TextBlock) -> None:
+    if block.has_paragraph_indentation:
+      self._close_texts_buffer()
+    self._texts_buffer.extend(block.texts)
+    if not block.last_line_touch_end:
+      self._close_texts_buffer()
+
+  def _write_image(self, block: AssetBlock) -> None:
+    os.makedirs(self._abs_assets_path, exist_ok=True)
+    hash = sha256()
+    hash.update(block.image.tobytes())
+    file_name = f"{hash.hexdigest()}.png"
+    file_path = os.path.join(self._abs_assets_path, file_name)
+    relative_path = os.path.join(self._assets_path, file_name)
+
+    if not os.path.exists(file_path):
+      block.image.save(file_path, "PNG")
+
+    self._file.write("![")
+    self._write_text_contents(block.texts, "]")
+    self._file.write(f"]({relative_path})")
+    self._file.write("\n\n")
+
+  def _write_text_contents(self, texts: Iterable[Text], ban_symbol: str | None = None) -> None:
+    for text in texts:
+      content = text.content.strip()
+      content = content.replace("\n", " ")
+      if ban_symbol is not None:
+        content.replace(ban_symbol, "")
+      self._file.write(content)

pdf_craft/pdf/__init__.py

@@ -0,0 +1 @@
+from .extractor import Block, Text, TextKind, TextBlock, AssetBlock, PDFPageExtractor

pdf_craft/pdf/document.py

@@ -0,0 +1,93 @@
+import os
+import fitz
+
+from typing import Generator, Literal
+from PIL.Image import frombytes, Image
+from doc_page_extractor import plot, Layout, PaddleLang, DocExtractor, ExtractedResult
+from .section import Section
+
+
+class DocumentExtractor:
+  def __init__(
+      self,
+      device: Literal["cpu", "cuda"],
+      model_dir_path: str,
+      debug_dir_path: str | None = None,
+    ):
+    self._debug_dir_path: str | None = debug_dir_path
+    self._doc_extractor = DocExtractor(
+      device=device,
+      model_dir_path=model_dir_path,
+      order_by_layoutreader=False,
+    )
+
+  def extract(self, pdf: str | fitz.Document, lang: PaddleLang) -> Generator[tuple[ExtractedResult, list[Layout]], None, None]:
+    for result, section in self._extract_results_and_sections(pdf, lang):
+      framework_layouts = section.framework()
+      yield result, [
+        layout for layout in result.layouts
+        if layout not in framework_layouts
+      ]
+
+  def _extract_results_and_sections(self, pdf: str | fitz.Document, lang: PaddleLang):
+    max_len = 2 # section can be viewed up to 2 pages back
+    queue: list[tuple[ExtractedResult, Section]] = []
+
+    for result in self._extract_page_result(pdf, lang):
+      section = Section(result.layouts)
+      for i, (_, pre_section) in enumerate(queue):
+        offset = len(queue) - i
+        pre_section.link_next(section, offset)
+
+      queue.append((result, section))
+      if len(queue) > max_len:
+        yield queue.pop(0)
+
+    for result, section in queue:
+      yield result, section
+
+  def _extract_page_result(self, pdf: str | fitz.Document, lang: PaddleLang):
+    if self._debug_dir_path is not None:
+      os.makedirs(self._debug_dir_path, exist_ok=True)
+
+    document: fitz.Document
+    should_close = False
+    if isinstance(pdf, str):
+      document = fitz.open(pdf)
+      should_close = True
+    else:
+      document = pdf
+
+    try:
+      for i, page in enumerate(document.pages()):
+        dpi = 300 # for scanned book pages
+        image = self._page_screenshot_image(page, dpi)
+        result = self._doc_extractor.extract(
+          image=image,
+          lang=lang,
+          adjust_points=False,
+        )
+        if self._debug_dir_path is not None:
+          self._generate_plot(image, i, result, self._debug_dir_path)
+        yield result
+    finally:
+      if should_close:
+        document.close()
+
+  def _page_screenshot_image(self, page: fitz.Page, dpi: int):
+    default_dpi = 72
+    matrix = fitz.Matrix(dpi / default_dpi, dpi / default_dpi)
+    pixmap = page.get_pixmap(matrix=matrix)
+    return frombytes("RGB", (pixmap.width, pixmap.height), pixmap.samples)
+
+  def _generate_plot(self, image: Image, index: int, result: ExtractedResult, plot_path: str):
+    plot_image: Image
+    if result.adjusted_image is None:
+      plot_image = image.copy()
+    else:
+      plot_image = result.adjusted_image
+
+    plot(plot_image, result.layouts)
+    os.makedirs(plot_path, exist_ok=True)
+    image_path = os.path.join(plot_path, f"plot_{index + 1}.png")
+    plot_image.save(image_path)

pdf_craft/pdf/extractor.py

@@ -0,0 +1,163 @@
+from typing import Literal
+from dataclasses import dataclass
+from enum import Enum
+from typing import Iterable, Generator
+from PIL.Image import Image
+from fitz import Document
+from doc_page_extractor import clip, PaddleLang, Rectangle, Layout, LayoutClass, OCRFragment, ExtractedResult
+from .document import DocumentExtractor
+
+
+class TextKind(Enum):
+  TITLE = 0
+  PLAIN_TEXT = 1
+  ABANDON = 2
+
+@dataclass
+class Text:
+  content: str
+  rank: float
+  rect: Rectangle
+
+@dataclass
+class TextBlock:
+  rect: Rectangle
+  kind: TextKind
+  texts: list[Text]
+  has_paragraph_indentation: bool = False
+  last_line_touch_end: bool = False
+
+@dataclass
+class AssetBlock:
+  rect: Rectangle
+  image: Image
+  texts: list[Text]
+
+Block = TextBlock | AssetBlock
+
+class PDFPageExtractor:
+  def __init__(
+      self,
+      device: Literal["cpu", "cuda"],
+      model_dir_path: str,
+      debug_dir_path: str | None = None,
+    ):
+    self._doc_extractor: DocumentExtractor = DocumentExtractor(
+      device=device,
+      model_dir_path=model_dir_path,
+      debug_dir_path=debug_dir_path,
+    )
+
+  def extract(self, pdf: str | Document, lang: PaddleLang) -> Generator[list[Block], None, None]:
+    for result, layouts in self._doc_extractor.extract(pdf, lang):
+      blocks = self._convert_to_blocks(result, layouts)
+      page_range = self._texts_range(blocks)
+
+      for block in blocks:
+        if not isinstance(block, TextBlock) or \
+           block.kind == TextKind.ABANDON:
+          continue
+
+        if len(block.texts) == 1:
+          mean_line_height, x1, x2 = page_range
+        else:
+          mean_line_height, x1, x2 = self._texts_range((block,))
+
+        first_text = block.texts[0]
+        last_text = block.texts[-1]
+        first_delta_x = (first_text.rect.lt[0] + first_text.rect.lb[0]) / 2 - x1
+        last_delta_x = x2 - (last_text.rect.rt[0] + last_text.rect.rb[0]) / 2
+        block.has_paragraph_indentation = first_delta_x > mean_line_height
+        block.last_line_touch_end = last_delta_x < mean_line_height
+
+      yield blocks
+
+  def _convert_to_blocks(self, result: ExtractedResult, layouts: list[Layout]) -> list[Block]:
+    store: list[tuple[LayoutClass, Block]] = []
+    def previous_block(cls: LayoutClass) -> Block | None:
+      for i in range(len(store) - 1, -1, -1):
+        block_cls, block = store[i]
+        if cls == block_cls:
+          return block
+        if cls != LayoutClass.ABANDON:
+          return None
+      return None
+
+    for layout in layouts:
+      cls = layout.cls
+      if cls == LayoutClass.TITLE:
+        store.append((cls, TextBlock(
+          rect=layout.rect,
+          kind=TextKind.TITLE,
+          texts=self._convert_to_text(layout.fragments),
+        )))
+      elif cls == LayoutClass.PLAIN_TEXT:
+        store.append((cls, TextBlock(
+          rect=layout.rect,
+          kind=TextKind.PLAIN_TEXT,
+          texts=self._convert_to_text(layout.fragments),
+        )))
+      elif cls == LayoutClass.ABANDON:
+        store.append((cls, TextBlock(
+          rect=layout.rect,
+          kind=TextKind.ABANDON,
+          texts=self._convert_to_text(layout.fragments),
+        )))
+      elif cls == LayoutClass.FIGURE or \
+           cls == LayoutClass.TABLE or \
+           cls == LayoutClass.ISOLATE_FORMULA:
+        store.append((cls, AssetBlock(
+          rect=layout.rect,
+          texts=[],
+          image=clip(result, layout),
+        )))
+      elif cls == LayoutClass.FIGURE_CAPTION:
+        block = previous_block(LayoutClass.FIGURE)
+        if block is not None:
+          assert isinstance(block, AssetBlock)
+          block.texts.extend(self._convert_to_text(layout.fragments))
+      elif cls == LayoutClass.TABLE_CAPTION or \
+           cls == LayoutClass.TABLE_FOOTNOTE:
+        block = previous_block(LayoutClass.TABLE)
+        if block is not None:
+          assert isinstance(block, AssetBlock)
+          block.texts.extend(self._convert_to_text(layout.fragments))
+      elif cls == LayoutClass.FORMULA_CAPTION:
+        block = previous_block(LayoutClass.ISOLATE_FORMULA)
+        if block is not None:
+          assert isinstance(block, AssetBlock)
+          block.texts.extend(self._convert_to_text(layout.fragments))
+
+    return [block for _, block in store]
+
+  def _texts_range(self, blocks: Iterable[Block]) -> tuple[float, float, float]:
+    sum_lines_height: float = 0.0
+    texts_count: int = 0
+    x1: float = float("inf")
+    x2: float = float("-inf")
+
+    for block in blocks:
+      if not isinstance(block, TextBlock):
+        continue
+      if block.kind == TextKind.ABANDON:
+        continue
+      for text in block.texts:
+        sum_lines_height += text.rect.size[1]
+        texts_count += 1
+        for x, _ in text.rect:
+          x1 = min(x1, x)
+          x2 = max(x2, x)
+
+    if texts_count == 0:
+      return 0.0, 0.0, 0.0
+    return sum_lines_height / texts_count, x1, x2
+
+  def _convert_to_text(self, fragments: list[OCRFragment]) -> list[Text]:
+    return [
+      Text(
+        content=f.text,
+        rank=f.rank,
+        rect=f.rect,
+      )
+      for f in fragments
+    ]

pdf_craft/pdf/section.py

@@ -0,0 +1,175 @@
+from __future__ import annotations
+
+from typing import Iterable, Sequence, Callable
+from doc_page_extractor import Point, Rectangle, Layout, OCRFragment
+from .text_matcher import check_texts_matching_rate
+from .utils import rate, intersection_area_size
+
+class _Shape:
+  def __init__(self, layout: Layout):
+    self.layout: Layout = layout
+    self.pre: list[Layout | None] = [None, None]
+    self.nex: list[Layout | None] = [None, None]
+
+  @property
+  def distance2(self) -> float:
+    x, y = self.layout.rect.lt
+    return x*x + y*y
+
+class Section:
+  def __init__(self, layouts: Iterable[Layout]) -> list[Layout]:
+    self._shapes: list[_Shape] = [_Shape(layout) for layout in layouts]
+
+  def framework(self) -> list[Layout]:
+    pre_shapes = list(self._side_framework(lambda shape: shape.pre))
+    nex_shapes = list(self._side_framework(lambda shape: shape.nex))
+    framework_shapes: list[_Shape]
+
+    if len(pre_shapes) < len(nex_shapes):
+      framework_shapes = nex_shapes
+    else:
+      framework_shapes = pre_shapes
+
+    return [shape.layout for shape in framework_shapes]
+
+  def _side_framework(self, get_side: Callable[[_Shape], list[Layout | None]]):
+    for shape in self._shapes:
+      matched_shapes = get_side(shape)
+      # The logic here should be understood as the processing of book page scans.
+      # When the book is opened, the even-numbered pages will be on the left and the odd-numbered pages will be on the right.
+      # Usually when editing and typesetting, even-numbered pages will be matched with one typesetting format,
+      # and odd-numbered pages will be matched with another typesetting format, in order to achieve aesthetics.
+      # Therefore, typesetting format (framework) matching usually takes into account the situation of matching every other page.
+      if matched_shapes[0] is not None or matched_shapes[1] is not None:
+        yield shape
+
+  def link_next(self, next: Section, offset: int) -> None:
+    assert offset in (1, 2), f"invalid offset {offset}"
+    matched_shapes_matrix: list[list[_Shape]] = []
+    for shape in self._shapes:
+      matched_shapes_matrix.append([
+        # pylint: disable=W0212
+        next_shape for next_shape in next._shapes
+        if self._is_shape_contents_matches(shape, next_shape)
+      ])
+
+    origin_shapes = self._find_origin_shapes(matched_shapes_matrix)
+    if origin_shapes is not None:
+      origins = (origin_shapes[0].layout.rect.lt, origin_shapes[1].layout.rect.lt)
+      for shape1, shape2 in self._iter_matched_shapes(origins, matched_shapes_matrix):
+        shape1.nex[offset - 1] = shape2
+        shape2.pre[offset - 1] = shape1
+
+  def _is_shape_contents_matches(self, shape1: _Shape, shape2: _Shape) -> bool:
+    size_match_rate = 0.95
+    layout1 = shape1.layout
+    layout2 = shape2.layout
+    size1 = layout1.rect.size
+    size2 = layout2.rect.size
+    if rate(size1[0], size2[0]) < size_match_rate or \
+       rate(size1[1], size2[1]) < size_match_rate:
+      return False
+
+    matched_count: int = 0
+
+    for fragment1 in layout1.fragments:
+      for fragment2 in layout2.fragments:
+        if self._is_fragments_matches(layout1, layout2, fragment1, fragment2):
+          matched_count += 1
+          break
+
+    fragments_count = max(len(layout1.fragments), len(layout2.fragments))
+    if fragments_count == 0:
+      return True
+
+    return self._check_group_matches(
+       matched_count / fragments_count,
+      fragments_count,
+      (0.0, 0.45, 0.45, 0.6, 0.8, 0.95),
+    )
+
+  def _is_fragments_matches(self, layout1: Layout, layout2: Layout, fragment1: OCRFragment, fragment2: OCRFragment) -> bool:
+    size_match_rate = 0.85
+    rect1 = self._relative_rect(layout1.rect.lt, fragment1.rect)
+    rect2 = self._relative_rect(layout2.rect.lt, fragment2.rect)
+    if self._intersection_rate(rect1, rect2) < size_match_rate:
+      return False
+
+    text_rate, text_length = check_texts_matching_rate(fragment1.text, fragment2.text)
+    return self._check_group_matches(
+      text_rate,
+      text_length,
+      (0.0, 0.0, 0.5, 0.55, 0.6, 0.75, 0.8, 0.95),
+    )
+
+  def _find_origin_shapes(self, matched_shapes: list[list[_Shape]]):
+    origin_shape1: _Shape | None = None
+    origin_shape2: _Shape | None = None
+    origin_matched_shape: list[_Shape] | None = None
+    min_distance2: float = float("inf")
+
+    for shape, matched_shapes in zip(self._shapes, matched_shapes):
+      if len(matched_shapes) == 0:
+        continue
+      distance2 = shape.distance2
+      if distance2 < min_distance2:
+        origin_shape1 = shape
+        origin_matched_shape = matched_shapes
+        min_distance2 = distance2
+
+    if origin_shape1 is None:
+      return None
+
+    assert origin_matched_shape is not None
+    min_distance2 = float("inf")
+
+    for shape in origin_matched_shape:
+      distance2 = shape.distance2
+      if distance2 < min_distance2:
+        origin_shape2 = shape
+        min_distance2 = distance2
+
+    assert origin_shape2 is not None
+    return origin_shape1, origin_shape2
+
+  def _iter_matched_shapes(self, origins: tuple[Point, Point], matched_shapes_matrix: list[list[_Shape]]):
+    for shape1, matched_shapes in zip(self._shapes, matched_shapes_matrix):
+      if len(matched_shapes) == 0:
+        continue
+
+      rect1 = self._relative_rect(origins[0], shape1.layout.rect)
+      max_area_rate: float = float("-inf")
+      matched_shape2: _Shape | None = None
+
+      for shape2 in matched_shapes:
+        rect2 = self._relative_rect(origins[1], shape2.layout.rect)
+        size_rate = self._intersection_rate(rect1, rect2)
+        if size_rate > max_area_rate:
+          max_area_rate = size_rate
+          matched_shape2 = shape2
+
+      if max_area_rate >= 0.85:
+        assert matched_shape2 is not None
+        yield shape1, matched_shape2
+
+  def _check_group_matches(self, calculated_rate: float, group_count: int, rates_list: Sequence[float]) -> bool:
+    if group_count < len(rates_list):
+      return calculated_rate >= rates_list[group_count]
+    else:
+      return calculated_rate >= rates_list[-1]
+
+  def _relative_rect(self, origin: Point, rect: Rectangle) -> Rectangle:
+    return Rectangle(
+      lt=(rect.lt[0] - origin[0], rect.lt[1] - origin[1]),
+      rt=(rect.rt[0] - origin[0], rect.rt[1] - origin[1]),
+      lb=(rect.lb[0] - origin[0], rect.lb[1] - origin[1]),
+      rb=(rect.rb[0] - origin[0], rect.rb[1] - origin[1])
+    )
+
+  def _intersection_rate(self, rect1: Rectangle, rect2: Rectangle) -> float:
+    width1, height1 = rect1.size
+    width2, height2 = rect2.size
+    width, height = intersection_area_size(rect1, rect2)
+    width_rate = width / max(width1, width2)
+    height_rate = height / max(height1, height2)
+    return 0.5 * (width_rate + height_rate)

pdf_craft/pdf/text_matcher.py

@@ -0,0 +1,93 @@
+import re
+import io
+
+from enum import Enum
+from unicodedata import category
+from alphabet_detector import AlphabetDetector
+
+
+def check_texts_matching_rate(text1: str, text2: str) -> tuple[float, int]:
+  words1: list[str] = list(split_into_words(text1))
+  words2: list[str] = list(split_into_words(text2))
+
+  if len(words1) > len(words2):
+    words1, words2 = words2, words1
+
+  marked_taken_list: list[bool] = [False] * len(words2)
+  matched_indexes: list[int] = []
+  for word1 in words1:
+    matched_index: int = -1
+    for i, word2 in enumerate(words2):
+      if not marked_taken_list[i] and word1 == word2:
+        marked_taken_list[i] = True
+        matched_index = i
+        break
+    matched_indexes.append(matched_index)
+
+  not_matched_count: int = len(words2) - len(words1)
+  for i, matched_index in enumerate(matched_indexes):
+    if matched_index == -1 or i > matched_index:
+      not_matched_count += 1
+
+  return 1.0 - not_matched_count / len(words2), len(words2)
+
+class _Phase(Enum):
+  Init = 0,
+  Letter = 1,
+  Character = 2,
+  Number = 3,
+  Space = 4,
+
+def split_into_words(text: str):
+  space_pattern = re.compile(r"\s")
+  number_pattern = re.compile(r"\d")
+  number_signs_pattern = re.compile(r"[\.,']")
+  word_buffer = io.StringIO()
+  phase: _Phase = _Phase.Init
+
+  for char in text:
+    if _is_letter(char):
+      if phase == _Phase.Number:
+        yield word_buffer.getvalue()
+        word_buffer = io.StringIO()
+      word_buffer.write(char)
+      phase = _Phase.Letter
+
+    elif number_pattern.match(char):
+      if phase == _Phase.Letter:
+        yield word_buffer.getvalue()
+        word_buffer = io.StringIO()
+      word_buffer.write(char)
+      phase = _Phase.Letter
+
+    elif phase == _Phase.Number and \
+         number_signs_pattern.match(char):
+      word_buffer.write(char)
+
+    else:
+      if phase == _Phase.Letter or \
+        phase == _Phase.Number:
+        yield word_buffer.getvalue()
+        word_buffer = io.StringIO()
+
+      if space_pattern.match(char):
+        phase = _Phase.Space
+      else: # others (like Chinese character)
+        yield char
+        phase = _Phase.Character
+
+  if phase == _Phase.Letter:
+    yield word_buffer.getvalue()
+
+ad = AlphabetDetector()
+
+# check is Latin, Cyrillic, Greek, or Hebrew letter
+def _is_letter(char: str):
+  if not category(char).startswith("L"):
+    return False
+
+  # AlphabetDetector unable to process punctuation
+  return ad.is_latin(char) or \
+         ad.is_cyrillic(char) or \
+         ad.is_greek(char) or \
+         ad.is_hebrew(char)

pdf_craft/pdf/utils.py

@@ -0,0 +1,29 @@
+from doc_page_extractor import Rectangle
+from shapely.geometry import Polygon
+
+def rate(value1: float, value2: float) -> float:
+  if value1 > value2:
+    value1, value2 = value2, value1
+  if value2 == 0.0:
+    return 1.0
+  else:
+    return value1 / value2
+
+def intersection_area_size(rect1: Rectangle, rect2: Rectangle) -> tuple[float, float]:
+  poly1 = Polygon(rect1)
+  poly2 = Polygon(rect2)
+  intersection: Polygon = poly1.intersection(poly2)
+  if intersection.is_empty:
+    return 0.0, 0.0
+
+  x1: float = float("inf")
+  y1: float = float("inf")
+  x2: float = float("-inf")
+  y2: float = float("-inf")
+
+  for x, y in intersection.exterior.coords:
+    x1 = min(x1, x)
+    y1 = min(y1, y)
+    x2 = max(x2, x)
+    y2 = max(y2, y)
+  return x2 - x1, y2 - y1