model_init.py 7.55 KB
Newer Older
1
import torch
2
3
from loguru import logger

4
from magic_pdf.config.constants import MODEL_NAME
5
from magic_pdf.model.model_list import AtomicModel
6
7
from magic_pdf.model.sub_modules.language_detection.yolov11.YOLOv11 import YOLOv11LangDetModel
from magic_pdf.model.sub_modules.layout.doclayout_yolo.DocLayoutYOLO import DocLayoutYOLOModel
8
9
from magic_pdf.model.sub_modules.mfd.yolov8.YOLOv8 import YOLOv8MFDModel
from magic_pdf.model.sub_modules.mfr.unimernet.Unimernet import UnimernetModel
10
11

try:
icecraft's avatar
icecraft committed
12
13
14
    from magic_pdf_ascend_plugin.libs.license_verifier import (
        LicenseExpiredError, LicenseFormatError, LicenseSignatureError,
        load_license)
15
16
    from magic_pdf_ascend_plugin.model_plugin.ocr.paddleocr.ppocr_273_npu import ModifiedPaddleOCR
    from magic_pdf_ascend_plugin.model_plugin.table.rapidtable.rapid_table_npu import RapidTableModel
17
    license_key = load_license()
18
19
    logger.info(f'Using Ascend Plugin Success, License id is {license_key["payload"]["id"]},'
                f' License expired at {license_key["payload"]["date"]["end_date"]}')
20
except Exception as e:
21
22
23
    if isinstance(e, ImportError):
        pass
    elif isinstance(e, LicenseFormatError):
icecraft's avatar
icecraft committed
24
        logger.error('Ascend Plugin: Invalid license format. Please check the license file.')
25
    elif isinstance(e, LicenseSignatureError):
icecraft's avatar
icecraft committed
26
        logger.error('Ascend Plugin: Invalid signature. The license may be tampered with.')
27
    elif isinstance(e, LicenseExpiredError):
icecraft's avatar
icecraft committed
28
        logger.error('Ascend Plugin: License has expired. Please renew your license.')
29
    elif isinstance(e, FileNotFoundError):
icecraft's avatar
icecraft committed
30
        logger.error('Ascend Plugin: Not found License file.')
31
    else:
icecraft's avatar
icecraft committed
32
        logger.error(f'Ascend Plugin: {e}')
33
34
35
    from magic_pdf.model.sub_modules.ocr.paddleocr.ppocr_273_mod import ModifiedPaddleOCR
    # from magic_pdf.model.sub_modules.ocr.paddleocr.ppocr_291_mod import ModifiedPaddleOCR
    from magic_pdf.model.sub_modules.table.rapidtable.rapid_table import RapidTableModel
36
37


38
def table_model_init(table_model_type, model_path, max_time, _device_='cpu', ocr_engine=None, table_sub_model_name=None):
39
    if table_model_type == MODEL_NAME.STRUCT_EQTABLE:
40
        from magic_pdf.model.sub_modules.table.structeqtable.struct_eqtable import StructTableModel
41
42
        table_model = StructTableModel(model_path, max_new_tokens=2048, max_time=max_time)
    elif table_model_type == MODEL_NAME.TABLE_MASTER:
43
        from magic_pdf.model.sub_modules.table.tablemaster.tablemaster_paddle import TableMasterPaddleModel
44
        config = {
45
46
            'model_dir': model_path,
            'device': _device_
47
48
49
        }
        table_model = TableMasterPaddleModel(config)
    elif table_model_type == MODEL_NAME.RAPID_TABLE:
50
        table_model = RapidTableModel(ocr_engine, table_sub_model_name)
51
    else:
52
        logger.error('table model type not allow')
53
54
55
56
57
58
        exit(1)

    return table_model


def mfd_model_init(weight, device='cpu'):
icecraft's avatar
icecraft committed
59
    if str(device).startswith('npu'):
60
        device = torch.device(device)
61
62
63
64
65
66
67
68
69
70
    mfd_model = YOLOv8MFDModel(weight, device)
    return mfd_model


def mfr_model_init(weight_dir, cfg_path, device='cpu'):
    mfr_model = UnimernetModel(weight_dir, cfg_path, device)
    return mfr_model


def layout_model_init(weight, config_file, device):
71
    from magic_pdf.model.sub_modules.layout.layoutlmv3.model_init import Layoutlmv3_Predictor
72
73
74
75
76
    model = Layoutlmv3_Predictor(weight, config_file, device)
    return model


def doclayout_yolo_model_init(weight, device='cpu'):
icecraft's avatar
icecraft committed
77
    if str(device).startswith('npu'):
78
        device = torch.device(device)
79
80
81
82
    model = DocLayoutYOLOModel(weight, device)
    return model


83
def langdetect_model_init(langdetect_model_weight, device='cpu'):
icecraft's avatar
icecraft committed
84
    if str(device).startswith('npu'):
85
        device = torch.device(device)
86
    model = YOLOv11LangDetModel(langdetect_model_weight, device)
87
88
89
    return model


90
91
92
93
94
95
def ocr_model_init(show_log: bool = False,
                   det_db_box_thresh=0.3,
                   lang=None,
                   use_dilation=True,
                   det_db_unclip_ratio=1.8,
                   ):
96
    if lang is not None and lang != '':
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
        model = ModifiedPaddleOCR(
            show_log=show_log,
            det_db_box_thresh=det_db_box_thresh,
            lang=lang,
            use_dilation=use_dilation,
            det_db_unclip_ratio=det_db_unclip_ratio,
        )
    else:
        model = ModifiedPaddleOCR(
            show_log=show_log,
            det_db_box_thresh=det_db_box_thresh,
            use_dilation=use_dilation,
            det_db_unclip_ratio=det_db_unclip_ratio,
        )
    return model


class AtomModelSingleton:
    _instance = None
    _models = {}

    def __new__(cls, *args, **kwargs):
        if cls._instance is None:
            cls._instance = super().__new__(cls)
        return cls._instance

    def get_atom_model(self, atom_model_name: str, **kwargs):
124

125
126
        lang = kwargs.get('lang', None)
        layout_model_name = kwargs.get('layout_model_name', None)
127
128
129
        table_model_name = kwargs.get('table_model_name', None)

        if atom_model_name in [AtomicModel.OCR]:
130
            key = (atom_model_name, lang)
131
132
133
134
135
136
137
        elif atom_model_name in [AtomicModel.Layout]:
            key = (atom_model_name, layout_model_name)
        elif atom_model_name in [AtomicModel.Table]:
            key = (atom_model_name, table_model_name)
        else:
            key = atom_model_name

138
139
        if key not in self._models:
            self._models[key] = atom_model_init(model_name=atom_model_name, **kwargs)
140
        return self._models[key]
141
142
143
144

def atom_model_init(model_name: str, **kwargs):
    atom_model = None
    if model_name == AtomicModel.Layout:
145
        if kwargs.get('layout_model_name') == MODEL_NAME.LAYOUTLMv3:
146
            atom_model = layout_model_init(
147
148
149
                kwargs.get('layout_weights'),
                kwargs.get('layout_config_file'),
                kwargs.get('device')
150
            )
151
        elif kwargs.get('layout_model_name') == MODEL_NAME.DocLayout_YOLO:
152
            atom_model = doclayout_yolo_model_init(
153
154
                kwargs.get('doclayout_yolo_weights'),
                kwargs.get('device')
155
            )
156
157
158
        else:
            logger.error('layout model name not allow')
            exit(1)
159
160
    elif model_name == AtomicModel.MFD:
        atom_model = mfd_model_init(
161
162
            kwargs.get('mfd_weights'),
            kwargs.get('device')
163
164
165
        )
    elif model_name == AtomicModel.MFR:
        atom_model = mfr_model_init(
166
167
168
            kwargs.get('mfr_weight_dir'),
            kwargs.get('mfr_cfg_path'),
            kwargs.get('device')
169
170
171
        )
    elif model_name == AtomicModel.OCR:
        atom_model = ocr_model_init(
172
173
            kwargs.get('ocr_show_log'),
            kwargs.get('det_db_box_thresh'),
174
            kwargs.get('lang'),
175
176
177
        )
    elif model_name == AtomicModel.Table:
        atom_model = table_model_init(
178
179
180
            kwargs.get('table_model_name'),
            kwargs.get('table_model_path'),
            kwargs.get('table_max_time'),
181
            kwargs.get('device'),
182
183
            kwargs.get('ocr_engine'),
            kwargs.get('table_sub_model_name')
184
        )
185
186
187
    elif model_name == AtomicModel.LangDetect:
        if kwargs.get('langdetect_model_name') == MODEL_NAME.YOLO_V11_LangDetect:
            atom_model = langdetect_model_init(
188
                kwargs.get('langdetect_model_weight'),
189
190
191
192
193
                kwargs.get('device')
            )
        else:
            logger.error('langdetect model name not allow')
            exit(1)
194
    else:
195
        logger.error('model name not allow')
196
197
198
        exit(1)

    if atom_model is None:
199
        logger.error('model init failed')
200
201
202
        exit(1)
    else:
        return atom_model