doc-page-extractor 0.0.9__py3-none-any.whl → 0.1.0__py3-none-any.whl

doc_page_extractor/__init__.py

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

doc_page_extractor/extractor.py

@@ -1,20 +1,30 @@
 import os
 
-from typing import Literal, Iterable
+from typing import Literal, Generator
 from pathlib import Path
 from PIL.Image import Image
-from transformers import LayoutLMv3ForTokenClassification
 from doclayout_yolo import YOLOv10
 
-from .layoutreader import prepare_inputs, boxes2inputs, parse_logits
 from .ocr import OCR
 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 .overlap import regroup_lines, remove_overlap_layouts
-from .utils import ensure_dir
+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:
@@ -23,15 +33,29 @@ class DocExtractor:
       model_dir_path: str,
       device: Literal["cpu", "cuda"] = "cpu",
       ocr_for_each_layouts: bool = True,
-      order_by_layoutreader: bool = False,
+      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._order_by_layoutreader: bool = order_by_layoutreader
-    self._ocr: OCR = OCR(device, model_dir_path)
+    self._extract_formula: bool = extract_formula
+    self._extract_table_format: TableLayoutParsedFormat | None = extract_table_format
     self._yolo: YOLOv10 | None = None
-    self._layout: LayoutLMv3ForTokenClassification | 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,
@@ -42,26 +66,21 @@ class DocExtractor:
     raw_optimizer = RawOptimizer(image, adjust_points)
     fragments = list(self._ocr.search_fragments(raw_optimizer.image_np))
     raw_optimizer.receive_raw_fragments(fragments)
-
-    layouts = self._get_layouts(raw_optimizer.image)
+    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)
 
-    if self._order_by_layoutreader:
-      width, height = raw_optimizer.image.size
-      self._order_fragments_by_ai(width, height, layouts)
-    else:
-      self._order_fragments_by_y(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 = regroup_lines(layout.fragments)
-      layout.fragments.sort(key=lambda fragment: fragment.order)
+      layout.fragments = merge_fragments_as_line(layout.fragments)
 
-    layouts = self._sort_layouts(layouts)
     raw_optimizer.receive_raw_layouts(layouts)
 
     return ExtractedResult(
@@ -71,7 +90,7 @@ class DocExtractor:
       adjusted_image=raw_optimizer.adjusted_image,
     )
 
-  def _get_layouts(self, source: Image) -> list[Layout]:
+  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(
@@ -81,7 +100,6 @@ class DocExtractor:
       device=self._device    # Device to use (e.g., "cuda" or "cpu")
     )
     boxes = det_res[0].__dict__["boxes"]
-    layouts: list[Layout] = []
 
     for cls_id, rect in zip(boxes.cls, boxes.xyxy):
       cls_id = cls_id.item()
@@ -98,9 +116,12 @@ class DocExtractor:
         lb=(x1, y2),
         rb=(x2, y2),
       )
-      layouts.append(Layout(cls, rect, []))
-
-    return layouts
+      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)
@@ -116,6 +137,17 @@ class DocExtractor:
     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] = []
@@ -158,67 +190,16 @@ class DocExtractor:
 
   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(self._model_dir_path, yolo_model_name))
+      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 _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]):
-    if width == 0 or height == 0:
-      return
-
-    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(
-        os.path.join(self._model_dir_path, "layoutreader"),
-      )
-      self._layout = LayoutLMv3ForTokenClassification.from_pretrained(
-        pretrained_model_name_or_path="hantian/layoutreader",
-        cache_dir=cache_dir,
-        local_files_only=os.path.exists(os.path.join(cache_dir, "models--hantian--layoutreader")),
-      )
-    return self._layout
-
   def _should_keep_layout(self, layout: Layout) -> bool:
     if len(layout.fragments) > 0:
       return True
@@ -229,78 +210,3 @@ class DocExtractor:
       cls == LayoutClass.ISOLATE_FORMULA
     )
 
-  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]] = []
-
-    for i, layout in enumerate(layouts):
-      if len(layout.fragments) > 0:
-        sorted_layouts.append((i, layout))
-      else:
-        empty_layouts.append((i, layout))
-
-    # 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 = -1
-      max_less_layout: Layout | None = None
-      max_less_index_in_enumerated: int = -1
-      for j, (k, sorted_layout) in enumerate(sorted_layouts):
-        if k < i and k > max_less_index:
-          max_less_index = k
-          max_less_layout = sorted_layout
-          max_less_index_in_enumerated = j
-
-      if max_less_layout is None:
-        sorted_layouts.insert(0, (i, layout))
-      else:
-        sorted_layouts.insert(max_less_index_in_enumerated + 1, (i, layout))
-
-    return [layout for _, layout in sorted_layouts]
-
-  def _collect_rate_boxes(self, fragments: Iterable[OCRFragment]):
-    boxes = self._get_boxes(fragments)
-    left = float("inf")
-    top = float("inf")
-    right = float("-inf")
-    bottom = float("-inf")
-
-    for _left, _top, _right, _bottom in boxes:
-      left = min(left, _left)
-      top = min(top, _top)
-      right = max(right, _right)
-      bottom = max(bottom, _bottom)
-
-    width = right - left
-    height = bottom - top
-
-    if width == 0 or height == 0:
-      return
-
-    for _left, _top, _right, _bottom in boxes:
-      yield (
-        (_left - left) / width,
-        (_top - top) / height,
-        (_right - left) / width,
-        (_bottom - top) / height,
-      )
-
-  def _get_boxes(self, fragments: Iterable[OCRFragment]):
-    boxes: list[tuple[float, float, float, float]] = []
-    for fragment in fragments:
-      left: float = float("inf")
-      top: float = float("inf")
-      right: float = float("-inf")
-      bottom: float = float("-inf")
-      for x, y in fragment.rect:
-        left = min(left, x)
-        top = min(top, y)
-        right = max(right, x)
-        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/latex.py

@@ -0,0 +1,57 @@
+import os
+import torch
+import requests
+
+from munch import Munch
+from pix2tex.cli import LatexOCR
+from PIL.Image import Image
+from .utils import expand_image
+
+
+class LaTeX:
+  def __init__(self, model_path: str):
+    self._model_path: str = model_path
+    self._model: LatexOCR | None = None
+
+  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._model is None:
+      if not os.path.exists(self._model_path):
+        self._download_model()
+
+      self._model = LatexOCR(Munch({
+        "config": os.path.join("settings", "config.yaml"),
+        "checkpoint": os.path.join(self._model_path, "weights.pth"),
+        "no_cuda": True,
+        "no_resize": False,
+      }))
+    return self._model
+
+  # from https://github.com/lukas-blecher/LaTeX-OCR/blob/5c1ac929bd19a7ecf86d5fb8d94771c8969fcb80/pix2tex/model/checkpoints/get_latest_checkpoint.py#L37-L45
+  def _download_model(self):
+    os.makedirs(self._model_path, exist_ok=True)
+    tag = "v0.0.1"
+    files: list[tuple[str, str]] = (
+      ("weights.pth", f"https://github.com/lukas-blecher/LaTeX-OCR/releases/download/{tag}/weights.pth"),
+      ("image_resizer.pth", f"https://github.com/lukas-blecher/LaTeX-OCR/releases/download/{tag}/image_resizer.pth")
+    )
+    for file_name, url in files:
+      file_path = os.path.join(self._model_path, file_name)
+      try:
+        with open(file_path, "wb") as file:
+          response = requests.get(url, stream=True, timeout=15)
+          response.raise_for_status()
+          for chunk in response.iter_content(chunk_size=8192):
+            if chunk:  # 过滤掉保持连接的新块
+              file.write(chunk)
+          file.flush()
+
+      except BaseException as e:
+        if os.path.exists(file_path):
+          os.remove(file_path)
+        raise e

doc_page_extractor/layout_order.py

@@ -0,0 +1,240 @@
+import os
+import torch
+
+from typing import Generator
+from dataclasses import dataclass
+from transformers import LayoutLMv3ForTokenClassification
+
+from .types import Layout, LayoutClass
+from .layoutreader import prepare_inputs, boxes2inputs, parse_logits
+from .utils import ensure_dir
+
+
+@dataclass
+class _BBox:
+  layout_index: int
+  fragment_index: int
+  virtual: bool
+  order: int
+  value: tuple[float, float, float, float]
+
+class LayoutOrder:
+  def __init__(self, model_path: str):
+    self._model_path: str = model_path
+    self._model: LayoutLMv3ForTokenClassification | None = None
+
+  def _get_model(self) -> LayoutLMv3ForTokenClassification:
+    if self._model is None:
+      model_path = ensure_dir(self._model_path)
+      self._model = LayoutLMv3ForTokenClassification.from_pretrained(
+        pretrained_model_name_or_path="hantian/layoutreader",
+        cache_dir=model_path,
+        local_files_only=os.path.exists(os.path.join(model_path, "models--hantian--layoutreader")),
+      )
+    return self._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.value, 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) -> _BBox:
+    x0, y0, x1, y1 = bbox
+    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

doc_page_extractor/ocr.py

@@ -58,7 +58,6 @@ class OCR:
     self._text_system: TextSystem | None = None
 
   def search_fragments(self, image: np.ndarray) -> Generator[OCRFragment, None, None]:
-    index: int = 0
     for box, res in self._ocr(image):
       text, rank = res
       if is_space_text(text):
@@ -74,12 +73,11 @@ class OCR:
         continue
 
       yield OCRFragment(
-        order=index,
+        order=0,
         text=text,
         rank=rank,
         rect=rect,
       )
-      index += 1
 
   def _ocr(self, image: np.ndarray) -> Generator[tuple[list[list[float]], tuple[str, float]], None, None]:
     text_system = self._get_text_system()

doc_page_extractor/overlap.py

@@ -60,7 +60,7 @@ class _OverlapMatrixContext:
          rate >= _INCLUDES_MIN_RATE:
         yield i
 
-def regroup_lines(origin_fragments: list[OCRFragment]) -> list[OCRFragment]:
+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:

doc_page_extractor/struct_eqtable/__init__.py

@@ -0,0 +1,49 @@
+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