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

doc_page_extractor/onnxocr/utils.py

@@ -1,71 +0,0 @@
-import numpy as np
-import cv2
-
-def get_rotate_crop_image(img, points):
-  """
-  img_height, img_width = img.shape[0:2]
-  left = int(np.min(points[:, 0]))
-  right = int(np.max(points[:, 0]))
-  top = int(np.min(points[:, 1]))
-  bottom = int(np.max(points[:, 1]))
-  img_crop = img[top:bottom, left:right, :].copy()
-  points[:, 0] = points[:, 0] - left
-  points[:, 1] = points[:, 1] - top
-  """
-  assert len(points) == 4, "shape of points must be 4*2"
-  img_crop_width = int(
-    max(
-      np.linalg.norm(points[0] - points[1]), np.linalg.norm(points[2] - points[3])
-    )
-  )
-  img_crop_height = int(
-    max(
-      np.linalg.norm(points[0] - points[3]), np.linalg.norm(points[1] - points[2])
-    )
-  )
-  pts_std = np.float32(
-    [
-      [0, 0],
-      [img_crop_width, 0],
-      [img_crop_width, img_crop_height],
-      [0, img_crop_height],
-    ]
-  )
-  M = cv2.getPerspectiveTransform(points, pts_std)
-  dst_img = cv2.warpPerspective(
-    img,
-    M,
-    (img_crop_width, img_crop_height),
-    borderMode=cv2.BORDER_REPLICATE,
-    flags=cv2.INTER_CUBIC,
-  )
-  dst_img_height, dst_img_width = dst_img.shape[0:2]
-  if dst_img_height * 1.0 / dst_img_width >= 1.5:
-    dst_img = np.rot90(dst_img)
-  return dst_img
-
-
-def get_minarea_rect_crop(img, points):
-  bounding_box = cv2.minAreaRect(np.array(points).astype(np.int32))
-  points = sorted(list(cv2.boxPoints(bounding_box)), key=lambda x: x[0])
-
-  index_a, index_b, index_c, index_d = 0, 1, 2, 3
-  if points[1][1] > points[0][1]:
-    index_a = 0
-    index_d = 1
-  else:
-    index_a = 1
-    index_d = 0
-  if points[3][1] > points[2][1]:
-    index_b = 2
-    index_c = 3
-  else:
-    index_b = 3
-    index_c = 2
-
-  box = [points[index_a], points[index_b], points[index_c], points[index_d]]
-  crop_img = get_rotate_crop_image(img, np.array(box))
-  return crop_img
-
-def str2bool(v):
-  return v.lower() in ("true", "t", "1")

doc_page_extractor/overlap.py

@@ -1,167 +0,0 @@
-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 merge_fragments_as_line(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)
-
-    if proto_texts_len == 0:
-      continue
-
-    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 height == 0:
-      continue
-
-    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)
-    if polygon2_width == 0.0 or polygon2_height == 0.0:
-      return 0.0
-    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/raw_optimizer.py

@@ -1,104 +0,0 @@
-import numpy as np
-
-from dataclasses import dataclass
-from PIL.Image import Image
-from math import pi
-from .types import OCRFragment, Layout
-from .rotation import calculate_rotation, RotationAdjuster
-from .rectangle import Rectangle
-
-
-_TINY_ROTATION = 0.005 # below this angle, we consider the text is horizontal
-
-
-@dataclass
-class _RotationContext:
-  to_origin: RotationAdjuster
-  to_new: RotationAdjuster
-  fragment_origin_rectangles: list[Rectangle]
-
-class RawOptimizer:
-  def __init__(
-      self,
-      raw: Image,
-      adjust_points: bool,
-    ):
-    self._raw: Image = raw
-    self._image: Image = raw
-    self._adjust_points: bool = adjust_points
-    self._fragments: list[OCRFragment]
-    self._rotation: float = 0.0
-    self._rotation_context: _RotationContext | None  = None
-
-  @property
-  def image(self) -> Image:
-    return self._image
-
-  @property
-  def adjusted_image(self) -> Image | None:
-    if self._adjust_points and self._image != self._raw:
-      return self._image
-
-  @property
-  def rotation(self) -> float:
-    return self._rotation
-
-  @property
-  def image_np(self) -> np.ndarray:
-    return np.array(self._raw)
-
-  def receive_raw_fragments(self, fragments: list[OCRFragment]):
-    self._fragments = fragments
-    self._rotation = calculate_rotation(fragments)
-
-    if abs(self._rotation) < _TINY_ROTATION:
-      return
-
-    origin_size = self._raw.size
-    self._image = self._raw.rotate(
-      angle=self._rotation * 180 / pi,
-      fillcolor=(255, 255, 255),
-      expand=True,
-    )
-    self._rotation_context = _RotationContext(
-      fragment_origin_rectangles=[f.rect for f in fragments],
-      to_origin=RotationAdjuster(
-        origin_size=origin_size,
-        new_size=self._image.size,
-        rotation=self._rotation,
-        to_origin_coordinate=True,
-      ),
-      to_new=RotationAdjuster(
-        origin_size=origin_size,
-        new_size=self._image.size,
-        rotation=self._rotation,
-        to_origin_coordinate=False,
-      ),
-    )
-    adjuster = self._rotation_context.to_new
-
-    for fragment in fragments:
-      rect = fragment.rect
-      fragment.rect = Rectangle(
-        lt=adjuster.adjust(rect.lt),
-        rt=adjuster.adjust(rect.rt),
-        lb=adjuster.adjust(rect.lb),
-        rb=adjuster.adjust(rect.rb),
-      )
-
-  def receive_raw_layouts(self, layouts: list[Layout]):
-    if self._adjust_points or self._rotation_context is None:
-      return
-
-    for fragment, origin_rect in zip(self._fragments, self._rotation_context.fragment_origin_rectangles):
-      fragment.rect = origin_rect
-
-    adjuster = self._rotation_context.to_origin
-
-    for layout in layouts:
-      layout.rect = Rectangle(
-        lt=adjuster.adjust(layout.rect.lt),
-        rt=adjuster.adjust(layout.rect.rt),
-        lb=adjuster.adjust(layout.rect.lb),
-        rb=adjuster.adjust(layout.rect.rb),
-      )

doc_page_extractor/rectangle.py

@@ -1,72 +0,0 @@
-from typing import Generator
-from dataclasses import dataclass
-from math import sqrt
-from shapely.geometry import Polygon
-
-
-Point = tuple[float, float]
-
-@dataclass
-class Rectangle:
-  lt: Point
-  rt: Point
-  lb: Point
-  rb: Point
-
-  def __iter__(self) -> Generator[Point, None, None]:
-    yield self.lt
-    yield self.lb
-    yield self.rb
-    yield self.rt
-
-  @property
-  def is_valid(self) -> bool:
-    return Polygon(self).is_valid
-
-  @property
-  def segments(self) -> Generator[tuple[Point, Point], None, None]:
-    yield (self.lt, self.lb)
-    yield (self.lb, self.rb)
-    yield (self.rb, self.rt)
-    yield (self.rt, self.lt)
-
-  @property
-  def area(self) -> float:
-    return Polygon(self).area
-
-  @property
-  def size(self) -> tuple[float, float]:
-    width: float = 0.0
-    height: float = 0.0
-    for i, (p1, p2) in enumerate(self.segments):
-      dx = p1[0] - p2[0]
-      dy = p1[1] - p2[1]
-      distance = sqrt(dx * dx + dy * dy)
-      if i % 2 == 0:
-        height += distance
-      else:
-        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)
-  if not poly1.is_valid or not poly2.is_valid:
-    return 0.0
-  intersection = poly1.intersection(poly2)
-  if intersection.is_empty:
-    return 0.0
-  return intersection.area

doc_page_extractor/rotation.py

@@ -1,158 +0,0 @@
-from math import pi, atan2, sqrt, sin, cos
-from .types import OCRFragment
-from .rectangle import Point, Rectangle
-
-
-class RotationAdjuster:
-  def __init__(
-      self,
-      origin_size: tuple[int, int],
-      new_size: tuple[int, int],
-      rotation: float,
-      to_origin_coordinate: bool,
-    ):
-    from_size: tuple[int, int]
-    to_size: tuple[int, int]
-    if to_origin_coordinate:
-      from_size = new_size
-      to_size = origin_size
-    else:
-      from_size = origin_size
-      to_size = new_size
-      rotation = -rotation
-
-    self._rotation: float = rotation
-    self._center_offset: tuple[float, float] = (
-      - from_size[0] / 2.0,
-      - from_size[1] / 2.0,
-    )
-    self._new_offset: tuple[float, float] = (
-      to_size[0] / 2.0,
-      to_size[1] / 2.0,
-    )
-
-  def adjust(self, point: Point) -> Point:
-    x, y = point
-    x += self._center_offset[0]
-    y += self._center_offset[1]
-
-    if x != 0.0 or y != 0.0:
-      radius = sqrt(x*x + y*y)
-      angle = atan2(y, x) + self._rotation
-      x = radius * cos(angle)
-      y = radius * sin(angle)
-
-    x += self._new_offset[0]
-    y += self._new_offset[1]
-
-    return x, y
-
-# to [0, pi)
-def normal_vertical_rotation(rotation: float) -> float:
-  while rotation >= 2 * pi:
-    rotation -= 2 * pi
-  while rotation <= - 2 * pi:
-    rotation += 2 * pi
-  if rotation < 0.0:
-    rotation += pi
-  return rotation
-
-def calculate_rotation(fragments: list[OCRFragment]):
-  horizontal_rotations: list[float] = []
-  vertical_rotations: list[float] = []
-
-  for fragment in fragments:
-    result = _rotation_with(fragment.rect)
-    if result is not None:
-      horizontal_rotations.extend(result[0])
-      vertical_rotations.extend(result[1])
-
-  if len(horizontal_rotations) == 0 or len(vertical_rotations) == 0:
-    return 0.0
-
-  horizontal_rotations = _normal_horizontal_rotations(horizontal_rotations)
-  horizontal_rotation = _find_median(horizontal_rotations)
-  vertical_rotation = _find_median(vertical_rotations)
-
-  return (vertical_rotation - 0.5 * pi + horizontal_rotation) / 2.0
-
-# @return horizontal [-pi/2, pi/2), vertical [0, pi)
-def calculate_rotation_with_rect(rect: Rectangle) -> tuple[float, float]:
-  result = _rotation_with(rect)
-  if result is None:
-    return 0.0, 0.5 * pi
-
-  horizontal_rotations, vertical_rotations = result
-  horizontal_rotations = _normal_horizontal_rotations(horizontal_rotations)
-  horizontal_rotation = _find_mean(horizontal_rotations)
-  vertical_rotation = _find_mean(vertical_rotations)
-
-  return horizontal_rotation, vertical_rotation
-
-def _rotation_with(rect: Rectangle):
-  rotations0: list[float] = []
-  rotations1: list[float] = []
-
-  for i, (p1, p2) in enumerate(rect.segments):
-    dx = p2[0] - p1[0]
-    dy = p2[1] - p1[1]
-    if dx == 0.0 and dy == 0.0:
-      return None
-    rotation: float = atan2(dy, dx)
-    if rotation < 0.0:
-      rotation += pi
-    if i % 2 == 0:
-      rotations0.append(rotation)
-    else:
-      rotations1.append(rotation)
-
-  if _is_vertical(rotations0):
-    return rotations1, rotations0
-  else:
-    return rotations0, rotations1
-
-# [0, pi) --> [-pi/2, pi/2)
-def _normal_horizontal_rotations(rotations: list[float]) -> list[float]:
-  for i, rotation in enumerate(rotations):
-    if rotation > 0.5 * pi:
-      rotations[i] = rotation - pi
-  return rotations
-
-def _find_median(rotations: list[float]):
-  rotations.sort()
-  n = len(rotations)
-
-  if n % 2 == 1:
-    return rotations[n // 2]
-  else:
-    mid1 = rotations[n // 2 - 1]
-    mid2 = rotations[n // 2]
-    return (mid1 + mid2) / 2
-
-def _find_mean(rotations: list[float]) -> float:
-  if len(rotations) == 0:
-    return 0.0
-  return sum(rotations) / len(rotations)
-
-# rotation is in [0, pi)
-def _is_vertical(rotations: list[float]):
-  horizontal_count: int = 0
-  vertical_count: int = 0
-  horizontal_delta: float = 0.0
-  vertical_delta: float = 0.0
-
-  for rotation in rotations:
-    if rotation < 0.25 * pi: # [0, pi/4)
-      horizontal_count += 1
-      horizontal_delta += rotation
-    elif rotation < 0.75 * pi: # [pi/4, 3pi/4)
-      vertical_count += 1
-      vertical_delta += abs(rotation - 0.5 * pi)
-    else: # [3pi/4, pi)
-      horizontal_count += 1
-      horizontal_delta += pi - rotation
-
-  if vertical_count == horizontal_delta:
-    return vertical_delta < horizontal_delta
-  else:
-    return vertical_count > horizontal_count

doc_page_extractor/struct_eqtable/__init__.py

@@ -1,49 +0,0 @@
-from .pix2s import Pix2Struct, Pix2StructTensorRT
-from .internvl import InternVL, InternVL_LMDeploy
-
-from transformers import AutoConfig
-
-
-__ALL_MODELS__ = {
-    'Pix2Struct': Pix2Struct,
-    'Pix2StructTensorRT': Pix2StructTensorRT,
-    'InternVL': InternVL,
-    'InternVL_LMDeploy': InternVL_LMDeploy,
-}
-
-
-def get_model_name(model_path):
-    model_config = AutoConfig.from_pretrained(
-        model_path,
-        trust_remote_code=True,
-    )
-
-    if 'Pix2Struct' in model_config.architectures[0]:
-        model_name = 'Pix2Struct'
-    elif 'InternVL' in model_config.architectures[0]:
-        model_name = 'InternVL'
-    else:
-        raise ValueError(f"Unsupported model type: {model_config.architectures[0]}")
-
-    return model_name
-
-
-def build_model(
-        model_ckpt='U4R/StructTable-InternVL2-1B',
-        cache_dir=None,
-        local_files_only=None,
-        **kwargs,
-    ):
-    model_name = get_model_name(model_ckpt)
-    if model_name == 'InternVL' and kwargs.get('lmdeploy', False):
-        model_name = 'InternVL_LMDeploy'
-    elif model_name == 'Pix2Struct' and kwargs.get('tensorrt_path', None):
-        model_name = 'Pix2StructTensorRT'
-
-    model = __ALL_MODELS__[model_name](
-        model_ckpt,
-        cache_dir=cache_dir,
-        local_files_only=local_files_only,
-        **kwargs
-    )
-    return model

doc_page_extractor/struct_eqtable/internvl/__init__.py

@@ -1,2 +0,0 @@
-from .internvl import InternVL
-from .internvl_lmdeploy import InternVL_LMDeploy