doc-page-extractor 0.0.2__tar.gz → 0.0.4__tar.gz

{doc_page_extractor-0.0.2 → doc_page_extractor-0.0.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: doc-page-extractor
-Version: 0.0.2
+Version: 0.0.4
 Summary: doc page extractor can identify text and format in images and return structured data.
 Home-page: https://github.com/Moskize91/doc-page-extractor
 Author: Tao Zeyu
@@ -12,7 +12,6 @@ Requires-Dist: pillow<11.0,>=10.3
 Requires-Dist: shapely<3.0,>=2.0.0
 Requires-Dist: transformers<5.0,>=4.48.0
 Requires-Dist: doclayout_yolo>=0.0.3
-Requires-Dist: paddlepaddle<3.0,>=2.6.0
 Requires-Dist: paddleocr==2.9.0
 Dynamic: author
 Dynamic: author-email

{doc_page_extractor-0.0.2 → doc_page_extractor-0.0.4}/doc_page_extractor/extractor.py

@@ -1,9 +1,7 @@
 import os
-import sys
 import torch
-import numpy as np
 
-from typing import Literal, Generator
+from typing import Literal, Iterable
 from pathlib import Path
 from PIL.Image import Image
 from transformers import LayoutLMv3ForTokenClassification
@@ -11,11 +9,13 @@ from doclayout_yolo import YOLOv10
 
 from .layoutreader import prepare_inputs, boxes2inputs, parse_logits
 from .ocr import OCR, PaddleLang
+from .ocr_corrector import correct_fragments
 from .raw_optimizer import RawOptimizer
 from .rectangle import intersection_area, Rectangle
 from .types import ExtractedResult, OCRFragment, LayoutClass, Layout
 from .downloader import download
-from .utils import ensure_dir, is_space_text
+from .overlap import regroup_lines, remove_overlap_layouts
+from .utils import ensure_dir
 
 
 class DocExtractor:
@@ -23,10 +23,12 @@ class DocExtractor:
       self,
       model_dir_path: str,
       device: Literal["cpu", "cuda"] = "cpu",
-      order_by_layoutreader: bool = True,
+      ocr_for_each_layouts: bool = True,
+      order_by_layoutreader: bool = False,
     ):
     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._order_by_layoutreader: bool = order_by_layoutreader
     self._ocr: OCR = OCR(device, os.path.join(model_dir_path, "paddle"))
     self._yolo: YOLOv10 | None = None
@@ -44,15 +46,28 @@ class DocExtractor:
     ) -> ExtractedResult:
 
     raw_optimizer = RawOptimizer(image, adjust_points)
-    fragments = list(self._search_orc_fragments(raw_optimizer.image_np, lang))
+    fragments = list(self._ocr.search_fragments(raw_optimizer.image_np, lang))
     raw_optimizer.receive_raw_fragments(fragments)
 
+    layouts = self._get_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, lang)
+
     if self._order_by_layoutreader:
       width, height = raw_optimizer.image.size
-      self._order_fragments(width, height, fragments)
+      self._order_fragments_by_ai(width, height, layouts)
+    else:
+      self._order_fragments_by_y(layouts)
 
-    layouts = self._get_layouts(raw_optimizer.image)
-    layouts = self._layouts_matched_by_fragments(fragments, layouts)
+    layouts = [layout for layout in layouts if self._should_keep_layout(layout)]
+    for layout in layouts:
+      layout.fragments = regroup_lines(layout.fragments)
+      layout.fragments.sort(key=lambda fragment: fragment.order)
+
+    layouts = self._sort_layouts(layouts)
     raw_optimizer.receive_raw_layouts(layouts)
 
     return ExtractedResult(
@@ -62,57 +77,6 @@ class DocExtractor:
       adjusted_image=raw_optimizer.adjusted_image,
     )
 
-  def _search_orc_fragments(self, image: np.ndarray, lang: PaddleLang) -> Generator[OCRFragment, None, None]:
-    index: int = 0
-    for item in self._ocr.do(lang, image):
-      for line in item:
-        react: list[list[float]] = line[0]
-        text, rank = line[1]
-        if is_space_text(text):
-          continue
-        yield OCRFragment(
-          order=index,
-          text=text,
-          rank=rank,
-          rect=Rectangle(
-            lt=(react[0][0], react[0][1]),
-            rt=(react[1][0], react[1][1]),
-            rb=(react[2][0], react[2][1]),
-            lb=(react[3][0], react[3][1]),
-          ),
-        )
-        index += 1
-
-  def _order_fragments(self, width: int, height: int, fragments: list[OCRFragment]):
-    layout_model = self._get_layout()
-    boxes: list[list[int]] = []
-    steps: float = 1000.0 # max value of layoutreader
-    x_rate: float = 1.0
-    y_rate: float = 1.0
-    x_offset: float = 0.0
-    y_offset: float = 0.0
-    if width > height:
-      y_rate = height / width
-      y_offset = (1.0 - y_rate) / 2.0
-    else:
-      x_rate = width / height
-      x_offset = (1.0 - x_rate) / 2.0
-
-    for left, top, right, bottom in self._collect_rate_boxes(fragments):
-      boxes.append([
-        round((left * x_rate + x_offset) * steps),
-        round((top * y_rate + y_offset) * steps),
-        round((right * x_rate + x_offset) * steps),
-        round((bottom * y_rate + y_offset) * steps),
-      ])
-    inputs = boxes2inputs(boxes)
-    inputs = prepare_inputs(inputs, layout_model)
-    logits = layout_model(**inputs).logits.cpu().squeeze(0)
-    orders: list[int] = parse_logits(logits, len(boxes))
-
-    for order, fragment in zip(orders, fragments):
-      fragment.order = order
-
   def _get_layouts(self, source: Image) -> list[Layout]:
     # about source parameter to see:
     # https://github.com/opendatalab/DocLayout-YOLO/blob/7c4be36bc61f11b67cf4a44ee47f3c41e9800a91/doclayout_yolo/data/build.py#L157-L175
@@ -152,14 +116,11 @@ class DocExtractor:
         if layout is not None:
           layout.fragments.append(fragment)
           break
+    return layouts
 
+  def _correct_fragments_by_ocr_layouts(self, source: Image, layouts: list[Layout], lang: PaddleLang):
     for layout in layouts:
-      layout.fragments.sort(key=lambda x: x.order)
-
-    layouts = [layout for layout in layouts if self._should_keep_layout(layout)]
-    layouts = self._sort_layouts(layouts)
-
-    return layouts
+      correct_fragments(self._ocr, source, layout, lang)
 
   def _split_layouts_by_group(self, layouts: list[Layout]):
     texts_layouts: list[Layout] = []
@@ -197,13 +158,6 @@ class DocExtractor:
 
     return min_layout
 
-  def _layout_order(self, layout: Layout) -> int:
-    fragments = layout.fragments
-    if len(fragments) == 0:
-      return sys.maxsize
-    else:
-      return fragments[0].order
-
   def _get_yolo(self) -> YOLOv10:
     if self._yolo is None:
       yolo_model_url = "https://huggingface.co/opendatalab/PDF-Extract-Kit-1.0/resolve/main/models/Layout/YOLO/doclayout_yolo_ft.pt"
@@ -214,6 +168,44 @@ class DocExtractor:
       self._yolo = YOLOv10(str(yolo_model_path))
     return self._yolo
 
+  def _order_fragments_by_y(self, layouts: list[Layout]):
+    fragments = list(self._iter_fragments(layouts))
+    fragments.sort(key=lambda f: f.rect.lt[1] + f.rect.rt[1])
+    for i, fragment in enumerate(fragments):
+      fragment.order = i
+
+  def _order_fragments_by_ai(self, width: int, height: int, layouts: list[Layout]):
+    layout_model = self._get_layout()
+    boxes: list[list[int]] = []
+    steps: float = 1000.0 # max value of layoutreader
+    x_rate: float = 1.0
+    y_rate: float = 1.0
+    x_offset: float = 0.0
+    y_offset: float = 0.0
+    if width > height:
+      y_rate = height / width
+      y_offset = (1.0 - y_rate) / 2.0
+    else:
+      x_rate = width / height
+      x_offset = (1.0 - x_rate) / 2.0
+
+    for left, top, right, bottom in self._collect_rate_boxes(
+      fragments=self._iter_fragments(layouts),
+    ):
+      boxes.append([
+        round((left * x_rate + x_offset) * steps),
+        round((top * y_rate + y_offset) * steps),
+        round((right * x_rate + x_offset) * steps),
+        round((bottom * y_rate + y_offset) * steps),
+      ])
+    inputs = boxes2inputs(boxes)
+    inputs = prepare_inputs(inputs, layout_model)
+    logits = layout_model(**inputs).logits.cpu().squeeze(0)
+    orders: list[int] = parse_logits(logits, len(boxes))
+
+    for order, fragment in zip(orders, self._iter_fragments(layouts)):
+      fragment.order = order
+
   def _get_layout(self) -> LayoutLMv3ForTokenClassification:
     if self._layout is None:
       cache_dir = ensure_dir(
@@ -237,6 +229,8 @@ class DocExtractor:
     )
 
   def _sort_layouts(self, layouts: list[Layout]) -> list[Layout]:
+    layouts.sort(key=lambda layout: layout.rect.lt[1] + layout.rect.rt[1])
+
     sorted_layouts: list[tuple[int, Layout]] = []
     empty_layouts: list[tuple[int, Layout]] = []
 
@@ -246,11 +240,9 @@ class DocExtractor:
       else:
         empty_layouts.append((i, layout))
 
-    sorted_layouts.sort(key=lambda x: x[1].fragments[0].order)
-
     # try to maintain the order of empty layouts and other layouts as much as possible
     for i, layout in empty_layouts:
-      max_less_index: int = len(layouts)
+      max_less_index: int = -1
       max_less_layout: Layout | None = None
       max_less_index_in_enumerated: int = -1
       for j, (k, sorted_layout) in enumerate(sorted_layouts):
@@ -266,7 +258,7 @@ class DocExtractor:
 
     return [layout for _, layout in sorted_layouts]
 
-  def _collect_rate_boxes(self, fragments: list[OCRFragment]):
+  def _collect_rate_boxes(self, fragments: Iterable[OCRFragment]):
     boxes = self._get_boxes(fragments)
     left = float("inf")
     top = float("inf")
@@ -290,7 +282,7 @@ class DocExtractor:
         (_bottom - top) / height,
       )
 
-  def _get_boxes(self, fragments: list[OCRFragment]):
+  def _get_boxes(self, fragments: Iterable[OCRFragment]):
     boxes: list[tuple[float, float, float, float]] = []
     for fragment in fragments:
       left: float = float("inf")
@@ -304,3 +296,7 @@ class DocExtractor:
         bottom = max(bottom, y)
       boxes.append((left, top, right, bottom))
     return boxes
+
+  def _iter_fragments(self, layouts: list[Layout]):
+    for layout in layouts:
+      yield from layout.fragments

{doc_page_extractor-0.0.2 → doc_page_extractor-0.0.4}/doc_page_extractor/ocr.py

@@ -2,9 +2,11 @@ import os
 import numpy as np
 import cv2
 
-from typing import Literal, Any
+from typing import Any, Literal, Generator
 from paddleocr import PaddleOCR
-from .utils import ensure_dir
+from .types import OCRFragment
+from .rectangle import Rectangle
+from .utils import is_space_text, ensure_dir
 
 
 # https://github.com/PaddlePaddle/PaddleOCR/blob/2c0c4beb0606819735a16083cdebf652939c781a/paddleocr.py#L108-L157
@@ -16,16 +18,33 @@ class OCR:
       self,
       device: Literal["cpu", "cuda"],
       model_dir_path: str,
-      bin: bool = True,
-      inv: bool = False,
     ):
     self._device: Literal["cpu", "cuda"] = device
     self._model_dir_path: str = model_dir_path
     self._ocr_and_lan: tuple[PaddleOCR, PaddleLang] | None = None
-    self._bin: bool = bin
-    self._inv: bool = inv
 
-  def do(self, lang: PaddleLang, image: np.ndarray) -> list[Any]:
+  def search_fragments(self, image: np.ndarray, lang: PaddleLang) -> Generator[OCRFragment, None, None]:
+    index: int = 0
+    for item in self._handle(lang, image):
+      for line in item:
+        react: list[list[float]] = line[0]
+        text, rank = line[1]
+        if is_space_text(text):
+          continue
+        yield OCRFragment(
+          order=index,
+          text=text,
+          rank=rank,
+          rect=Rectangle(
+            lt=(react[0][0], react[0][1]),
+            rt=(react[1][0], react[1][1]),
+            rb=(react[2][0], react[2][1]),
+            lb=(react[3][0], react[3][1]),
+          ),
+        )
+        index += 1
+
+  def _handle(self, lang: PaddleLang, image: np.ndarray) -> list[Any]:
     ocr = self._get_ocr(lang)
     image = self._preprocess_image(image)
     # about img parameter to see
@@ -59,10 +78,24 @@ class OCR:
 
   def _preprocess_image(self, image: np.ndarray) -> np.ndarray:
     image = self._alpha_to_color(image, (255, 255, 255))
-    if self._inv:
-      image = cv2.bitwise_not(image)
-    if self._bin:
-      image = self._binarize_img(image)
+    # image = cv2.bitwise_not(image) # inv
+    # image = self._binarize_img(image) # bin
+    image = cv2.normalize(
+      src=image,
+      dst=np.zeros((image.shape[0], image.shape[1])),
+      alpha=0,
+      beta=255,
+      norm_type=cv2.NORM_MINMAX,
+    )
+    image = cv2.fastNlMeansDenoisingColored(
+      src=image,
+      dst=None,
+      h=10,
+      hColor=10,
+      templateWindowSize=7,
+      searchWindowSize=15,
+    )
+    # image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) # image to gray
     return image
 
   def _alpha_to_color(self, image: np.ndarray, alpha_color: tuple[float, float, float]) -> np.ndarray:

doc_page_extractor-0.0.4/doc_page_extractor/ocr_corrector.py

@@ -0,0 +1,126 @@
+import numpy as np
+
+from typing import Iterable
+from shapely.geometry import Polygon
+from PIL.Image import new, Image, Resampling
+from .types import Layout, OCRFragment
+from .ocr import OCR, PaddleLang
+from .overlap import overlap_rate
+from .rectangle import Point, Rectangle
+
+
+_MIN_RATE = 0.5
+
+def correct_fragments(ocr: OCR, source: Image, layout: Layout, lang: PaddleLang):
+  x1, y1, x2, y2 = layout.rect.wrapper
+  image: Image = source.crop((
+    round(x1), round(y1),
+    round(x2), round(y2),
+  ))
+  image, dx, dy, scale = _adjust_image(image)
+  image_np = np.array(image)
+  ocr_fragments = list(ocr.search_fragments(image_np, lang))
+  corrected_fragments: list[OCRFragment] = []
+
+  for fragment in ocr_fragments:
+    _apply_fragment(fragment.rect, layout, dx, dy, scale)
+
+  matched_fragments, not_matched_fragments = _match_fragments(
+    zone_rect=layout.rect,
+    fragments1=layout.fragments,
+    fragments2=ocr_fragments,
+  )
+  for fragment1, fragment2 in matched_fragments:
+    if fragment1.rank > fragment2.rank:
+      corrected_fragments.append(fragment1)
+    else:
+      corrected_fragments.append(fragment2)
+
+  corrected_fragments.extend(not_matched_fragments)
+  layout.fragments = corrected_fragments
+
+def _adjust_image(image: Image) -> tuple[Image, int, int, float]:
+  # after testing, adding white borders to images can reduce
+  # the possibility of some text not being recognized
+  border_size: int = 50
+  adjusted_size: int = 1024 - 2 * border_size
+  width, height = image.size
+  core_width = float(max(adjusted_size, width))
+  core_height = float(max(adjusted_size, height))
+
+  scale_x = core_width / width
+  scale_y = core_height / height
+  scale = min(scale_x, scale_y)
+  adjusted_width = width * scale
+  adjusted_height = height * scale
+
+  dx = (core_width - adjusted_width) / 2.0
+  dy = (core_height - adjusted_height) / 2.0
+  dx = round(dx) + border_size
+  dy = round(dy) + border_size
+
+  if scale != 1.0:
+    width = round(width * scale)
+    height = round(height * scale)
+    image = image.resize((width, height), Resampling.BICUBIC)
+
+  width = round(core_width) + 2 * border_size
+  height = round(core_height) + 2 * border_size
+  new_image = new("RGB", (width, height), (255, 255, 255))
+  new_image.paste(image, (dx, dy))
+
+  return new_image, dx, dy, scale
+
+def _apply_fragment(rect: Rectangle, layout: Layout, dx: int, dy: int, scale: float):
+  rect.lt = _apply_point(rect.lt, layout, dx, dy, scale)
+  rect.lb = _apply_point(rect.lb, layout, dx, dy, scale)
+  rect.rb = _apply_point(rect.rb, layout, dx, dy, scale)
+  rect.rt = _apply_point(rect.rt, layout, dx, dy, scale)
+
+def _apply_point(point: Point, layout: Layout, dx: int, dy: int, scale: float) -> Point:
+  x, y = point
+  x = (x - dx) / scale + layout.rect.lt[0]
+  y = (y - dy) / scale + layout.rect.lt[1]
+  return x, y
+
+def _match_fragments(
+    zone_rect: Rectangle,
+    fragments1: Iterable[OCRFragment],
+    fragments2: Iterable[OCRFragment],
+  ) -> tuple[list[tuple[OCRFragment, OCRFragment]], list[OCRFragment]]:
+
+  zone_polygon = Polygon(zone_rect)
+  fragments2: list[OCRFragment] = list(fragments2)
+  matched_fragments: list[tuple[OCRFragment, OCRFragment]] = []
+  not_matched_fragments: list[OCRFragment] = []
+
+  for fragment1 in fragments1:
+    polygon1 = Polygon(fragment1.rect)
+    polygon1 = zone_polygon.intersection(polygon1)
+    if polygon1.is_empty:
+      continue
+
+    beast_j = -1
+    beast_rate = 0.0
+
+    for j, fragment2 in enumerate(fragments2):
+      polygon2 = Polygon(fragment2.rect)
+      rate = overlap_rate(polygon1, polygon2)
+      if rate < _MIN_RATE:
+        continue
+
+      if rate > beast_rate:
+        beast_j = j
+        beast_rate = rate
+
+    if beast_j != -1:
+      matched_fragments.append((
+        fragment1,
+        fragments2[beast_j],
+      ))
+      del fragments2[beast_j]
+    else:
+      not_matched_fragments.append(fragment1)
+
+  not_matched_fragments.extend(fragments2)
+  return matched_fragments, not_matched_fragments

doc_page_extractor-0.0.4/doc_page_extractor/overlap.py

@@ -0,0 +1,156 @@
+from typing import Generator
+from shapely.geometry import Polygon
+from .types import Layout, OCRFragment
+from .rectangle import Rectangle
+
+
+_INCLUDES_MIN_RATE = 0.99
+
+def remove_overlap_layouts(layouts: list[Layout]) -> list[Layout]:
+  ctx = _OverlapMatrixContext(layouts)
+  # the reason for repeating this multiple times is that deleting a layout
+  # may cause its parent layout to change from an originally non-deletable
+  # state to a deletable state.
+  while True:
+    removed_count = len(ctx.removed_indexes)
+    for i, layout in enumerate(layouts):
+      if i in ctx.removed_indexes or \
+         any(0.0 < rate < _INCLUDES_MIN_RATE for rate in ctx.rates_with_other(i)) or \
+         all(0.0 == rate for rate in ctx.rates_with_other(i)):
+        continue
+
+      if len(layout.fragments) == 0:
+        ctx.removed_indexes.add(i)
+      else:
+        for j in ctx.search_includes_indexes(i):
+          ctx.removed_indexes.add(j)
+          layout.fragments.extend(layouts[j].fragments)
+
+    if len(ctx.removed_indexes) == removed_count:
+      break
+
+  return [
+    layout for i, layout in enumerate(layouts)
+    if i not in ctx.removed_indexes
+  ]
+
+class _OverlapMatrixContext:
+  def __init__(self, layouts: list[Layout]):
+    length: int = len(layouts)
+    polygons: list[Polygon] = [Polygon(layout.rect) for layout in layouts]
+    self.rate_matrix: list[list[float]] = [[1.0 for _ in range(length)] for _ in range(length)]
+    self.removed_indexes: set[int] = set()
+    for i in range(length):
+      polygon1 = polygons[i]
+      rates = self.rate_matrix[i]
+      for j in range(length):
+        if i != j:
+          polygon2 = polygons[j]
+          rates[j] = overlap_rate(polygon1, polygon2)
+
+  def rates_with_other(self, index: int):
+    for i, rate in enumerate(self.rate_matrix[index]):
+      if i != index and i not in self.removed_indexes:
+        yield rate
+
+  def search_includes_indexes(self, index: int):
+    for i, rate in enumerate(self.rate_matrix[index]):
+      if i != index and \
+         i not in self.removed_indexes and \
+         rate >= _INCLUDES_MIN_RATE:
+        yield i
+
+def regroup_lines(origin_fragments: list[OCRFragment]) -> list[OCRFragment]:
+  fragments: list[OCRFragment] = []
+  for group in _split_fragments_into_groups(origin_fragments):
+    if len(group) == 1:
+      fragments.append(group[0])
+      continue
+
+    min_order: float = float("inf")
+    texts: list[str] = []
+    text_rate_weights: float = 0.0
+    proto_texts_len: int = 0
+
+    x1: float = float("inf")
+    y1: float = float("inf")
+    x2: float = float("-inf")
+    y2: float = float("-inf")
+
+    for fragment in sorted(group, key=lambda x: x.rect.lt[0] + x.rect.lb[0]):
+      proto_texts_len += len(fragment.text)
+      text_rate_weights += fragment.rank * len(fragment.text)
+      texts.append(fragment.text)
+      min_order = min(min_order, fragment.order)
+      for x, y in fragment.rect:
+        x1 = min(x1, x)
+        y1 = min(y1, y)
+        x2 = max(x2, x)
+        y2 = max(y2, y)
+
+    fragments.append(OCRFragment(
+      order=min_order,
+      text=" ".join(texts),
+      rank=text_rate_weights / proto_texts_len,
+      rect=Rectangle(
+        lt=(x1, y1),
+        rt=(x2, y1),
+        lb=(x1, y2),
+        rb=(x2, y2),
+      ),
+    ))
+  return fragments
+
+def _split_fragments_into_groups(fragments: list[OCRFragment]) -> Generator[list[OCRFragment], None, None]:
+  group: list[OCRFragment] = []
+  sum_height: float = 0.0
+  sum_median: float = 0.0
+  max_deviation_rate = 0.35
+
+  for fragment in sorted(fragments, key=lambda x: x.rect.lt[1] + x.rect.rt[1]):
+    _, y1, _, y2 = fragment.rect.wrapper
+    height = y2 - y1
+    median = (y1 + y2) / 2.0
+
+    if len(group) > 0:
+      next_mean_median = (sum_median + median) / (len(group) + 1)
+      next_mean_height = (sum_height + height) / (len(group) + 1)
+      deviation_rate = abs(median - next_mean_median) / next_mean_height
+
+      if deviation_rate > max_deviation_rate:
+        yield group
+        group = []
+        sum_height = 0.0
+        sum_median = 0.0
+
+    group.append(fragment)
+    sum_height += height
+    sum_median += median
+
+  if len(group) > 0:
+    yield group
+
+# calculating overlap ratio: The reason why area is not used is
+# that most of the measurements are of rectangles representing text lines.
+# they are very sensitive to changes in height because they are very thin and long.
+# In order to make it equally sensitive to length and width, the ratio of area is not used.
+def overlap_rate(polygon1: Polygon, polygon2: Polygon) -> float:
+  intersection: Polygon = polygon1.intersection(polygon2)
+  if intersection.is_empty:
+    return 0.0
+  else:
+    overlay_width, overlay_height = _polygon_size(intersection)
+    polygon2_width, polygon2_height = _polygon_size(polygon2)
+    return (overlay_width / polygon2_width + overlay_height / polygon2_height) / 2.0
+
+def _polygon_size(polygon: Polygon) -> tuple[float, float]:
+  x1: float = float("inf")
+  y1: float = float("inf")
+  x2: float = float("-inf")
+  y2: float = float("-inf")
+  for x, y in polygon.exterior.coords:
+    x1 = min(x1, x)
+    y1 = min(y1, y)
+    x2 = max(x2, x)
+    y2 = max(y2, y)
+  return x2 - x1, y2 - y1

{doc_page_extractor-0.0.2 → doc_page_extractor-0.0.4}/doc_page_extractor/plot.py

@@ -8,11 +8,11 @@ _FRAGMENT_COLOR = (0x49, 0xCF, 0xCB) # Light Green
 def plot(image: Image, layouts: Iterable[Layout]):
   draw = ImageDraw.Draw(image, mode="RGBA")
   for layout in layouts:
-    draw.polygon([p for p in layout.rect], outline=_layout_color(layout), width=3)
+    draw.polygon([p for p in layout.rect], outline=_layout_color(layout), width=5)
 
   for layout in layouts:
     for fragments in layout.fragments:
-      draw.polygon([p for p in fragments.rect], outline=_FRAGMENT_COLOR, width=1)
+      draw.polygon([p for p in fragments.rect], outline=_FRAGMENT_COLOR, width=3)
 
 def _layout_color(layout: Layout) -> tuple[int, int, int]:
   cls = layout.cls

{doc_page_extractor-0.0.2 → doc_page_extractor-0.0.4}/doc_page_extractor/rectangle.py

@@ -44,6 +44,19 @@ class Rectangle:
         width += distance
     return width / 2, height / 2
 
+  @property
+  def wrapper(self) -> tuple[float, float, float, float]:
+    x1: float = float("inf")
+    y1: float = float("inf")
+    x2: float = float("-inf")
+    y2: float = float("-inf")
+    for x, y in self:
+      x1 = min(x1, x)
+      y1 = min(y1, y)
+      x2 = max(x2, x)
+      y2 = max(y2, y)
+    return x1, y1, x2, y2
+
 def intersection_area(rect1: Rectangle, rect2: Rectangle) -> float:
   poly1 = Polygon(rect1)
   poly2 = Polygon(rect2)

{doc_page_extractor-0.0.2 → doc_page_extractor-0.0.4}/doc_page_extractor.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: doc-page-extractor
-Version: 0.0.2
+Version: 0.0.4
 Summary: doc page extractor can identify text and format in images and return structured data.
 Home-page: https://github.com/Moskize91/doc-page-extractor
 Author: Tao Zeyu
@@ -12,7 +12,6 @@ Requires-Dist: pillow<11.0,>=10.3
 Requires-Dist: shapely<3.0,>=2.0.0
 Requires-Dist: transformers<5.0,>=4.48.0
 Requires-Dist: doclayout_yolo>=0.0.3
-Requires-Dist: paddlepaddle<3.0,>=2.6.0
 Requires-Dist: paddleocr==2.9.0
 Dynamic: author
 Dynamic: author-email

{doc_page_extractor-0.0.2 → doc_page_extractor-0.0.4}/doc_page_extractor.egg-info/SOURCES.txt

@@ -7,6 +7,8 @@ doc_page_extractor/downloader.py
 doc_page_extractor/extractor.py
 doc_page_extractor/layoutreader.py
 doc_page_extractor/ocr.py
+doc_page_extractor/ocr_corrector.py
+doc_page_extractor/overlap.py
 doc_page_extractor/plot.py
 doc_page_extractor/raw_optimizer.py
 doc_page_extractor/rectangle.py
@@ -17,4 +19,6 @@ doc_page_extractor.egg-info/PKG-INFO
 doc_page_extractor.egg-info/SOURCES.txt
 doc_page_extractor.egg-info/dependency_links.txt
 doc_page_extractor.egg-info/requires.txt
-doc_page_extractor.egg-info/top_level.txt
+doc_page_extractor.egg-info/top_level.txt
+tests/__init__.py
+tests/test_history_bus.py

{doc_page_extractor-0.0.2 → doc_page_extractor-0.0.4}/doc_page_extractor.egg-info/requires.txt

@@ -3,5 +3,4 @@ pillow<11.0,>=10.3
 shapely<3.0,>=2.0.0
 transformers<5.0,>=4.48.0
 doclayout_yolo>=0.0.3
-paddlepaddle<3.0,>=2.6.0
 paddleocr==2.9.0

{doc_page_extractor-0.0.2 → doc_page_extractor-0.0.4}/doc_page_extractor.egg-info/top_level.txt

@@ -1 +1,2 @@
 doc_page_extractor
+tests

{doc_page_extractor-0.0.2 → doc_page_extractor-0.0.4}/setup.py

@@ -2,7 +2,7 @@ from setuptools import setup, find_packages
 
 setup(
   name="doc-page-extractor",
-  version="0.0.2",
+  version="0.0.4",
   author="Tao Zeyu",
   author_email="i@taozeyu.com",
   url="https://github.com/Moskize91/doc-page-extractor",
@@ -16,7 +16,6 @@ setup(
     "shapely>=2.0.0,<3.0",
     "transformers>=4.48.0,<5.0",
     "doclayout_yolo>=0.0.3",
-    "paddlepaddle>=2.6.0,<3.0",
     "paddleocr==2.9.0", # https://github.com/Moskize91/doc-page-extractor/issues/3
   ],
 )

doc_page_extractor-0.0.4/tests/test_history_bus.py

@@ -0,0 +1,55 @@
+import os
+import unittest
+
+from PIL import Image
+from doc_page_extractor import DocExtractor, Layout, LayoutClass
+
+
+class TestGroup(unittest.TestCase):
+  def test_history_bugs(self):
+    model_path = os.path.join(self._project_path(), "model")
+    image_path = os.path.join(self._project_path(), "tests", "images", "figure.png")
+    os.makedirs(model_path, exist_ok=True)
+
+    extractor = DocExtractor(model_path, "cpu")
+    layouts: list[tuple[LayoutClass, list[str]]]
+
+    with Image.open(image_path) as image:
+      result = extractor.extract(image, "ch")
+      layouts = [self._format_Layout(layout) for layout in result.layouts]
+
+    self.assertEqual(layouts, [
+      (LayoutClass.PLAIN_TEXT, [
+        "口的11.8%①。这既是江南农业落后的反映，又是它的原因。当战国以",
+        "后黄河流域因铁器牛耕的普及获得基本的开发，农区联结成一大片的",
+        "时候，南方农业开发始终没有突破星点状或斑块状分布的格局。由于",
+        "地旷人稀，耕作相当粗放，许多水田采取火耕水瓣的方式，旱田则多",
+        "行刀耕火种②。司马迁在《史记·货殖列传》中说：“总之，楚越之",
+        "地，地厂人希，饭稻囊鱼，或火耕而水瓣，果隋（蕨）赢（螺）蛤，",
+        "不待贾而足，地势饶食，无饥谨之患，以故皆偷生，无积聚而多",
+        "贫。”这种概括虽然未免太突出了南方经济的落后面，有一定片面性，",
+        "但大体还是反映了实际情形的。战国秦汉时期，南方与黄河流域农业",
+        "的差距显然拉大了。",
+      ]),
+      (LayoutClass.FIGURE, []),
+      (LayoutClass.FIGURE_CAPTION, [
+        "西晋陶水田犁耙模型（广东连县出土）"
+      ]),
+      (LayoutClass.FIGURE, []),
+      (LayoutClass.FIGURE_CAPTION, [
+        "南朝陶耙田模型 （广西苍梧倒水出土）"
+      ]),
+      (LayoutClass.PLAIN_TEXT, [
+        "①据赵文林、谢淑君：《中国人口史》（人民出版社1988年）有关资料统计。",
+        "②《盐铁论·通有》：“荆扬…………伐木而树谷，焚莱而播粟，火耕而水。”"
+      ]),
+      (LayoutClass.ABANDON, [
+        "136"
+      ]),
+    ])
+
+  def _format_Layout(self, layout: Layout) -> tuple[LayoutClass, list[str]]:
+    return layout.cls, [f.text.strip() for f in layout.fragments]
+
+  def _project_path(self) -> str:
+    return os.path.abspath(os.path.join(__file__, "..", ".."))