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 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
......
......@@ -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:
......
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