浏览代码

refactor: update model paths and enhance RapidTableModel initialization

myhloli 5 月之前
父节点
当前提交
c3531d72ae
共有 3 个文件被更改,包括 7 次插入6 次删除
  1. 4 4
      mineru/model/table/rapid_table.py
  2. 二进制
      mineru/resources/slanet_plus/slanet-plus.onnx
  3. 3 2
      mineru/utils/enum_class.py

+ 4 - 4
mineru/model/table/rapid_table.py

@@ -1,15 +1,15 @@
-import os
-from pathlib import Path
 import cv2
 import numpy as np
 from loguru import logger
 from rapid_table import RapidTable, RapidTableInput
 
+from mineru.utils.enum_class import ModelPath
+from mineru.utils.models_download_utils import get_file_from_repos
+
 
 class RapidTableModel(object):
     def __init__(self, ocr_engine):
-        root_dir = Path(__file__).absolute().parent.parent.parent
-        slanet_plus_model_path = os.path.join(root_dir, 'resources', 'slanet_plus', 'slanet-plus.onnx')
+        slanet_plus_model_path = get_file_from_repos(ModelPath.slanet_plus)
         input_args = RapidTableInput(model_type='slanet_plus', model_path=slanet_plus_model_path)
         self.table_model = RapidTable(input_args)
         self.ocr_engine = ocr_engine

二进制
mineru/resources/slanet_plus/slanet-plus.onnx


+ 3 - 2
mineru/utils/enum_class.py

@@ -46,6 +46,8 @@ class MakeMode:
 
 
 class ModelPath:
+    vlm_root_hf = "opendatalab/MinerU2.0-2505-0.9B"
+    vlm_root_modelscope = "OpenDataLab/MinerU2.0-2505-0.9B"
     pipeline_root_modelscope = "OpenDataLab/PDF-Extract-Kit-1.0"
     pipeline_root_hf = "opendatalab/PDF-Extract-Kit-1.0"
     doclayout_yolo = "models/Layout/YOLO/doclayout_yolo_docstructbench_imgsz1280_2501.pt"
@@ -53,8 +55,7 @@ class ModelPath:
     unimernet_small = "models/MFR/unimernet_hf_small_2503"
     pytorch_paddle = "models/OCR/paddleocr_torch"
     layout_reader = "models/ReadingOrder/layout_reader"
-    vlm_root_hf = "opendatalab/MinerU2.0-2505-0.9B"
-    vlm_root_modelscope = "OpenDataLab/MinerU2.0-2505-0.9B"
+    slanet_plus = "models/TabRec/SlanetPlus/slanet-plus.onnx"
 
 
 class SplitFlag: