Commit c3531d72 authored by myhloli's avatar myhloli
Browse files

refactor: update model paths and enhance RapidTableModel initialization

parent 83b965b7
import os
from pathlib import Path
import cv2 import cv2
import numpy as np import numpy as np
from loguru import logger from loguru import logger
from rapid_table import RapidTable, RapidTableInput 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): class RapidTableModel(object):
def __init__(self, ocr_engine): def __init__(self, ocr_engine):
root_dir = Path(__file__).absolute().parent.parent.parent slanet_plus_model_path = get_file_from_repos(ModelPath.slanet_plus)
slanet_plus_model_path = os.path.join(root_dir, 'resources', 'slanet_plus', 'slanet-plus.onnx')
input_args = RapidTableInput(model_type='slanet_plus', model_path=slanet_plus_model_path) input_args = RapidTableInput(model_type='slanet_plus', model_path=slanet_plus_model_path)
self.table_model = RapidTable(input_args) self.table_model = RapidTable(input_args)
self.ocr_engine = ocr_engine self.ocr_engine = ocr_engine
......
...@@ -46,6 +46,8 @@ class MakeMode: ...@@ -46,6 +46,8 @@ class MakeMode:
class ModelPath: 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_modelscope = "OpenDataLab/PDF-Extract-Kit-1.0"
pipeline_root_hf = "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" doclayout_yolo = "models/Layout/YOLO/doclayout_yolo_docstructbench_imgsz1280_2501.pt"
...@@ -53,8 +55,7 @@ class ModelPath: ...@@ -53,8 +55,7 @@ class ModelPath:
unimernet_small = "models/MFR/unimernet_hf_small_2503" unimernet_small = "models/MFR/unimernet_hf_small_2503"
pytorch_paddle = "models/OCR/paddleocr_torch" pytorch_paddle = "models/OCR/paddleocr_torch"
layout_reader = "models/ReadingOrder/layout_reader" layout_reader = "models/ReadingOrder/layout_reader"
vlm_root_hf = "opendatalab/MinerU2.0-2505-0.9B" slanet_plus = "models/TabRec/SlanetPlus/slanet-plus.onnx"
vlm_root_modelscope = "OpenDataLab/MinerU2.0-2505-0.9B"
class SplitFlag: class SplitFlag:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment