pdf_extract_kit.py 21.6 KB
Newer Older
1
from loguru import logger
myhloli's avatar
myhloli committed
2
import os
3
import time
4
from magic_pdf.libs.Constants import *
5
from magic_pdf.libs.clean_memory import clean_memory
6
from magic_pdf.model.model_list import AtomicModel
7
8

os.environ['NO_ALBUMENTATIONS_UPDATE'] = '1'  # 禁止albumentations检查更新
9
os.environ['YOLO_VERBOSE'] = 'False'  # disable yolo logger
myhloli's avatar
myhloli committed
10
11
12
13
14
15
try:
    import cv2
    import yaml
    import argparse
    import numpy as np
    import torch
16
    import torchtext
17

18
19
    if torchtext.__version__ >= "0.18.0":
        torchtext.disable_torchtext_deprecation_warning()
myhloli's avatar
myhloli committed
20
21
22
23
24
25
26
    from PIL import Image
    from torchvision import transforms
    from torch.utils.data import Dataset, DataLoader
    from ultralytics import YOLO
    from unimernet.common.config import Config
    import unimernet.tasks as tasks
    from unimernet.processors import load_processor
27
    from doclayout_yolo import YOLOv10
28
    from rapid_table import RapidTable
29
    from rapidocr_paddle import RapidOCR
赵小蒙's avatar
update:  
赵小蒙 committed
30

31
32
except ImportError as e:
    logger.exception(e)
33
34
    logger.error(
        'Required dependency not installed, please install by \n'
35
        '"pip install magic-pdf[full] --extra-index-url https://myhloli.github.io/wheels/"')
myhloli's avatar
myhloli committed
36
    exit(1)
赵小蒙's avatar
update:  
赵小蒙 committed
37

38
from magic_pdf.model.pek_sub_modules.layoutlmv3.model_init import Layoutlmv3_Predictor
39
from magic_pdf.model.pek_sub_modules.post_process import latex_rm_whitespace
40
from magic_pdf.model.pek_sub_modules.self_modify import ModifiedPaddleOCR
41
from magic_pdf.model.pek_sub_modules.structeqtable.StructTableModel import StructTableModel
42
43
44
45
from magic_pdf.model.ppTableModel import ppTableModel


def table_model_init(table_model_type, model_path, max_time, _device_='cpu'):
46
    ocr_engine = None
47
    if table_model_type == MODEL_NAME.STRUCT_EQTABLE:
48
        table_model = StructTableModel(model_path, max_time=max_time)
49
    elif table_model_type == MODEL_NAME.TABLE_MASTER:
50
51
52
53
54
        config = {
            "model_dir": model_path,
            "device": _device_
        }
        table_model = ppTableModel(config)
55
56
    elif table_model_type == MODEL_NAME.RAPID_TABLE:
        table_model = RapidTable()
57
        ocr_engine = RapidOCR(det_use_cuda=True, cls_use_cuda=True, rec_use_cuda=True)
58
59
60
    else:
        logger.error("table model type not allow")
        exit(1)
61

62
63
64
65
    if ocr_engine:
        return [table_model, ocr_engine]
    else:
        return table_model
66

赵小蒙's avatar
update:  
赵小蒙 committed
67

68
69
70
def mfd_model_init(weight):
    mfd_model = YOLO(weight)
    return mfd_model
赵小蒙's avatar
update:  
赵小蒙 committed
71
72


73
def mfr_model_init(weight_dir, cfg_path, _device_='cpu'):
74
75
    args = argparse.Namespace(cfg_path=cfg_path, options=None)
    cfg = Config(args)
76
    cfg.config.model.pretrained = os.path.join(weight_dir, "pytorch_model.pth")
77
78
79
80
    cfg.config.model.model_config.model_name = weight_dir
    cfg.config.model.tokenizer_config.path = weight_dir
    task = tasks.setup_task(cfg)
    model = task.build_model(cfg)
81
    model.to(_device_)
82
    model.eval()
83
    vis_processor = load_processor('formula_image_eval', cfg.config.datasets.formula_rec_eval.vis_processor.eval)
84
85
    mfr_transform = transforms.Compose([vis_processor, ])
    return [model, mfr_transform]
赵小蒙's avatar
update:  
赵小蒙 committed
86
87


88
89
90
91
92
def layout_model_init(weight, config_file, device):
    model = Layoutlmv3_Predictor(weight, config_file, device)
    return model


93
94
95
96
97
def doclayout_yolo_model_init(weight):
    model = YOLOv10(weight)
    return model


98
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):
99
    if lang is not None:
100
        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)
101
    else:
102
        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)
103
104
105
    return model


106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
class MathDataset(Dataset):
    def __init__(self, image_paths, transform=None):
        self.image_paths = image_paths
        self.transform = transform

    def __len__(self):
        return len(self.image_paths)

    def __getitem__(self, idx):
        # if not pil image, then convert to pil image
        if isinstance(self.image_paths[idx], str):
            raw_image = Image.open(self.image_paths[idx])
        else:
            raw_image = self.image_paths[idx]
        if self.transform:
            image = self.transform(raw_image)
122
            return image
123
124


125
126
127
128
129
130
131
132
133
134
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):
135
136
137
138
139
140
        lang = kwargs.get("lang", None)
        layout_model_name = kwargs.get("layout_model_name", None)
        key = (atom_model_name, layout_model_name, lang)
        if key not in self._models:
            self._models[key] = atom_model_init(model_name=atom_model_name, **kwargs)
        return self._models[key]
141
142
143
144
145


def atom_model_init(model_name: str, **kwargs):

    if model_name == AtomicModel.Layout:
146
147
148
149
150
151
152
153
154
155
        if kwargs.get("layout_model_name") == MODEL_NAME.LAYOUTLMv3:
            atom_model = layout_model_init(
                kwargs.get("layout_weights"),
                kwargs.get("layout_config_file"),
                kwargs.get("device")
            )
        elif kwargs.get("layout_model_name") == MODEL_NAME.DocLayout_YOLO:
            atom_model = doclayout_yolo_model_init(
                kwargs.get("doclayout_yolo_weights"),
            )
156
157
158
159
160
161
162
163
164
165
166
167
168
    elif model_name == AtomicModel.MFD:
        atom_model = mfd_model_init(
            kwargs.get("mfd_weights")
        )
    elif model_name == AtomicModel.MFR:
        atom_model = mfr_model_init(
            kwargs.get("mfr_weight_dir"),
            kwargs.get("mfr_cfg_path"),
            kwargs.get("device")
        )
    elif model_name == AtomicModel.OCR:
        atom_model = ocr_model_init(
            kwargs.get("ocr_show_log"),
169
170
            kwargs.get("det_db_box_thresh"),
            kwargs.get("lang")
171
172
173
        )
    elif model_name == AtomicModel.Table:
        atom_model = table_model_init(
174
            kwargs.get("table_model_name"),
175
176
177
178
179
180
181
182
183
184
185
            kwargs.get("table_model_path"),
            kwargs.get("table_max_time"),
            kwargs.get("device")
        )
    else:
        logger.error("model name not allow")
        exit(1)

    return atom_model


186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
#  Unified crop img logic
def crop_img(input_res, input_pil_img, crop_paste_x=0, crop_paste_y=0):
    crop_xmin, crop_ymin = int(input_res['poly'][0]), int(input_res['poly'][1])
    crop_xmax, crop_ymax = int(input_res['poly'][4]), int(input_res['poly'][5])
    # Create a white background with an additional width and height of 50
    crop_new_width = crop_xmax - crop_xmin + crop_paste_x * 2
    crop_new_height = crop_ymax - crop_ymin + crop_paste_y * 2
    return_image = Image.new('RGB', (crop_new_width, crop_new_height), 'white')

    # Crop image
    crop_box = (crop_xmin, crop_ymin, crop_xmax, crop_ymax)
    cropped_img = input_pil_img.crop(crop_box)
    return_image.paste(cropped_img, (crop_paste_x, crop_paste_y))
    return_list = [crop_paste_x, crop_paste_y, crop_xmin, crop_ymin, crop_xmax, crop_ymax, crop_new_width, crop_new_height]
    return return_image, return_list


203
class CustomPEKModel:
204

205
206
207
208
209
210
211
212
213
214
215
216
217
218
    def __init__(self, ocr: bool = False, show_log: bool = False, **kwargs):
        """
        ======== model init ========
        """
        # 获取当前文件(即 pdf_extract_kit.py)的绝对路径
        current_file_path = os.path.abspath(__file__)
        # 获取当前文件所在的目录(model)
        current_dir = os.path.dirname(current_file_path)
        # 上一级目录(magic_pdf)
        root_dir = os.path.dirname(current_dir)
        # model_config目录
        model_config_dir = os.path.join(root_dir, 'resources', 'model_config')
        # 构建 model_configs.yaml 文件的完整路径
        config_path = os.path.join(model_config_dir, 'model_configs.yaml')
219
        with open(config_path, "r", encoding='utf-8') as f:
220
221
            self.configs = yaml.load(f, Loader=yaml.FullLoader)
        # 初始化解析配置
222
223
224
225
226
227
228
229
230
231
232

        # layout config
        self.layout_config = kwargs.get("layout_config")
        self.layout_model_name = self.layout_config.get("model", MODEL_NAME.DocLayout_YOLO)

        # formula config
        self.formula_config = kwargs.get("formula_config")
        self.mfd_model_name = self.formula_config.get("mfd_model", MODEL_NAME.YOLO_V8_MFD)
        self.mfr_model_name = self.formula_config.get("mfr_model", MODEL_NAME.UniMerNet_v2_Small)
        self.apply_formula = self.formula_config.get("enable", True)

233
        # table config
234
235
        self.table_config = kwargs.get("table_config")
        self.apply_table = self.table_config.get("enable", False)
236
        self.table_max_time = self.table_config.get("max_time", TABLE_MAX_TIME_VALUE)
237
        self.table_model_name = self.table_config.get("model", MODEL_NAME.RAPID_TABLE)
238
239

        # ocr config
240
        self.apply_ocr = ocr
241
        self.lang = kwargs.get("lang", None)
242

243
        logger.info(
244
245
246
            "DocAnalysis init, this may take some times, layout_model: {}, apply_formula: {}, apply_ocr: {}, "
            "apply_table: {}, table_model: {}, lang: {}".format(
                self.layout_model_name, self.apply_formula, self.apply_ocr, self.apply_table, self.table_model_name, self.lang
赵小蒙's avatar
update:  
赵小蒙 committed
247
            )
248
249
        )
        # 初始化解析方案
250
        self.device = kwargs.get("device", "cpu")
251
        logger.info("using device: {}".format(self.device))
252
        models_dir = kwargs.get("models_dir", os.path.join(root_dir, "resources", "models"))
253
        logger.info("using models_dir: {}".format(models_dir))
254

255
256
        atom_model_manager = AtomModelSingleton()

257
258
        # 初始化公式识别
        if self.apply_formula:
259

260
            # 初始化公式检测模型
261
262
            self.mfd_model = atom_model_manager.get_atom_model(
                atom_model_name=AtomicModel.MFD,
263
                mfd_weights=str(os.path.join(models_dir, self.configs["weights"][self.mfd_model_name]))
264
            )
265

266
            # 初始化公式解析模型
267
            mfr_weight_dir = str(os.path.join(models_dir, self.configs["weights"][self.mfr_model_name]))
268
            mfr_cfg_path = str(os.path.join(model_config_dir, "UniMERNet", "demo.yaml"))
269
270
271
272
273
274
            self.mfr_model, self.mfr_transform = atom_model_manager.get_atom_model(
                atom_model_name=AtomicModel.MFR,
                mfr_weight_dir=mfr_weight_dir,
                mfr_cfg_path=mfr_cfg_path,
                device=self.device
            )
275
276

        # 初始化layout模型
277
278
279
280
281
282
283
284
285
286
287
288
289
290
        if self.layout_model_name == MODEL_NAME.LAYOUTLMv3:
            self.layout_model = atom_model_manager.get_atom_model(
                atom_model_name=AtomicModel.Layout,
                layout_model_name=MODEL_NAME.LAYOUTLMv3,
                layout_weights=str(os.path.join(models_dir, self.configs['weights'][self.layout_model_name])),
                layout_config_file=str(os.path.join(model_config_dir, "layoutlmv3", "layoutlmv3_base_inference.yaml")),
                device=self.device
            )
        elif self.layout_model_name == MODEL_NAME.DocLayout_YOLO:
            self.layout_model = atom_model_manager.get_atom_model(
                atom_model_name=AtomicModel.Layout,
                layout_model_name=MODEL_NAME.DocLayout_YOLO,
                doclayout_yolo_weights=str(os.path.join(models_dir, self.configs['weights'][self.layout_model_name]))
            )
291
        # 初始化ocr
292
293
294
295
296
297
298
        if self.apply_ocr:
            self.ocr_model = atom_model_manager.get_atom_model(
                atom_model_name=AtomicModel.OCR,
                ocr_show_log=show_log,
                det_db_box_thresh=0.3,
                lang=self.lang
            )
299
        # init table model
300
        if self.apply_table:
301
            table_model_dir = self.configs["weights"][self.table_model_name]
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
            if self.table_model_name in [MODEL_NAME.STRUCT_EQTABLE, MODEL_NAME.TABLE_MASTER]:
                self.table_model = atom_model_manager.get_atom_model(
                    atom_model_name=AtomicModel.Table,
                    table_model_name=self.table_model_name,
                    table_model_path=str(os.path.join(models_dir, table_model_dir)),
                    table_max_time=self.table_max_time,
                    device=self.device
                )
            elif self.table_model_name in [MODEL_NAME.RAPID_TABLE]:
                self.table_model, self.ocr_engine =atom_model_manager.get_atom_model(
                    atom_model_name=AtomicModel.Table,
                    table_model_name=self.table_model_name,
                    table_model_path=str(os.path.join(models_dir, table_model_dir)),
                    table_max_time=self.table_max_time,
                    device=self.device
                )
drunkpig's avatar
drunkpig committed
318

319
        logger.info('DocAnalysis init done!')
赵小蒙's avatar
update:  
赵小蒙 committed
320

321
322
    def __call__(self, image):

323
324
        page_start = time.time()

325
326
327
        latex_filling_list = []
        mf_image_list = []

328
329
        # layout检测
        layout_start = time.time()
330
331
332
333
334
335
        if self.layout_model_name == MODEL_NAME.LAYOUTLMv3:
            # layoutlmv3
            layout_res = self.layout_model(image, ignore_catids=[])
        elif self.layout_model_name == MODEL_NAME.DocLayout_YOLO:
            # doclayout_yolo
            layout_res = []
336
            doclayout_yolo_res = self.layout_model.predict(image, imgsz=1024, conf=0.25, iou=0.45, verbose=True, device=self.device)[0]
337
338
339
340
341
342
343
344
            for xyxy, conf, cla in zip(doclayout_yolo_res.boxes.xyxy.cpu(), doclayout_yolo_res.boxes.conf.cpu(), doclayout_yolo_res.boxes.cls.cpu()):
                xmin, ymin, xmax, ymax = [int(p.item()) for p in xyxy]
                new_item = {
                    'category_id': int(cla.item()),
                    'poly': [xmin, ymin, xmax, ymin, xmax, ymax, xmin, ymax],
                    'score': round(float(conf.item()), 3),
                }
                layout_res.append(new_item)
345
        layout_cost = round(time.time() - layout_start, 2)
346
        logger.info(f"layout detection time: {layout_cost}")
347

348
349
        pil_img = Image.fromarray(image)

350
351
        if self.apply_formula:
            # 公式检测
352
            mfd_start = time.time()
353
            mfd_res = self.mfd_model.predict(image, imgsz=1888, conf=0.25, iou=0.45, verbose=True, device=self.device)[0]
354
            logger.info(f"mfd time: {round(time.time() - mfd_start, 2)}")
355
356
357
358
359
360
361
362
363
364
            for xyxy, conf, cla in zip(mfd_res.boxes.xyxy.cpu(), mfd_res.boxes.conf.cpu(), mfd_res.boxes.cls.cpu()):
                xmin, ymin, xmax, ymax = [int(p.item()) for p in xyxy]
                new_item = {
                    'category_id': 13 + int(cla.item()),
                    'poly': [xmin, ymin, xmax, ymin, xmax, ymax, xmin, ymax],
                    'score': round(float(conf.item()), 2),
                    'latex': '',
                }
                layout_res.append(new_item)
                latex_filling_list.append(new_item)
365
                bbox_img = pil_img.crop((xmin, ymin, xmax, ymax))
366
367
368
369
370
371
372
373
374
                mf_image_list.append(bbox_img)

            # 公式识别
            mfr_start = time.time()
            dataset = MathDataset(mf_image_list, transform=self.mfr_transform)
            dataloader = DataLoader(dataset, batch_size=64, num_workers=0)
            mfr_res = []
            for mf_img in dataloader:
                mf_img = mf_img.to(self.device)
375
376
                with torch.no_grad():
                    output = self.mfr_model.generate({'image': mf_img})
377
378
379
380
381
                mfr_res.extend(output['pred_str'])
            for res, latex in zip(latex_filling_list, mfr_res):
                res['latex'] = latex_rm_whitespace(latex)
            mfr_cost = round(time.time() - mfr_start, 2)
            logger.info(f"formula nums: {len(mf_image_list)}, mfr time: {mfr_cost}")
382

383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
        # Select regions for OCR / formula regions / table regions
        ocr_res_list = []
        table_res_list = []
        single_page_mfdetrec_res = []
        for res in layout_res:
            if int(res['category_id']) in [13, 14]:
                single_page_mfdetrec_res.append({
                    "bbox": [int(res['poly'][0]), int(res['poly'][1]),
                             int(res['poly'][4]), int(res['poly'][5])],
                })
            elif int(res['category_id']) in [0, 1, 2, 4, 6, 7]:
                ocr_res_list.append(res)
            elif int(res['category_id']) in [5]:
                table_res_list.append(res)

398
        if torch.cuda.is_available() and self.device != 'cpu':
399
400
            total_memory = torch.cuda.get_device_properties(self.device).total_memory / (1024 ** 3)  # 将字节转换为 GB
            if total_memory <= 8:
401
402
403
404
                gc_start = time.time()
                clean_memory()
                gc_time = round(time.time() - gc_start, 2)
                logger.info(f"gc time: {gc_time}")
405

myhloli's avatar
myhloli committed
406
        # ocr识别
407
        if self.apply_ocr:
408
            ocr_start = time.time()
409
            # Process each area that requires OCR processing
410
            for res in ocr_res_list:
411
412
413
                new_image, useful_list = crop_img(res, pil_img, crop_paste_x=50, crop_paste_y=50)
                paste_x, paste_y, xmin, ymin, xmax, ymax, new_width, new_height = useful_list
                # Adjust the coordinates of the formula area
414
415
416
                adjusted_mfdetrec_res = []
                for mf_res in single_page_mfdetrec_res:
                    mf_xmin, mf_ymin, mf_xmax, mf_ymax = mf_res["bbox"]
417
                    # Adjust the coordinates of the formula area to the coordinates relative to the cropping area
418
419
420
421
                    x0 = mf_xmin - xmin + paste_x
                    y0 = mf_ymin - ymin + paste_y
                    x1 = mf_xmax - xmin + paste_x
                    y1 = mf_ymax - ymin + paste_y
422
                    # Filter formula blocks outside the graph
423
                    if any([x1 < 0, y1 < 0]) or any([x0 > new_width, y0 > new_height]):
424
425
426
427
428
429
                        continue
                    else:
                        adjusted_mfdetrec_res.append({
                            "bbox": [x0, y0, x1, y1],
                        })

430
                # OCR recognition
431
432
                new_image = cv2.cvtColor(np.asarray(new_image), cv2.COLOR_RGB2BGR)
                ocr_res = self.ocr_model.ocr(new_image, mfd_res=adjusted_mfdetrec_res)[0]
433

434
                # Integration results
435
436
437
438
439
                if ocr_res:
                    for box_ocr_res in ocr_res:
                        p1, p2, p3, p4 = box_ocr_res[0]
                        text, score = box_ocr_res[1]

440
                        # Convert the coordinates back to the original coordinate system
441
442
443
444
445
446
447
448
449
450
451
452
                        p1 = [p1[0] - paste_x + xmin, p1[1] - paste_y + ymin]
                        p2 = [p2[0] - paste_x + xmin, p2[1] - paste_y + ymin]
                        p3 = [p3[0] - paste_x + xmin, p3[1] - paste_y + ymin]
                        p4 = [p4[0] - paste_x + xmin, p4[1] - paste_y + ymin]

                        layout_res.append({
                            'category_id': 15,
                            'poly': p1 + p2 + p3 + p4,
                            'score': round(score, 2),
                            'text': text,
                        })

453
            ocr_cost = round(time.time() - ocr_start, 2)
454
            logger.info(f"ocr time: {ocr_cost}")
455

456
457
        # 表格识别 table recognition
        if self.apply_table:
458
459
460
461
            table_start = time.time()
            for res in table_res_list:
                new_image, _ = crop_img(res, pil_img)
                single_table_start_time = time.time()
462
                # logger.info("------------------table recognition processing begins-----------------")
463
464
                latex_code = None
                html_code = None
465
                if self.table_model_name == MODEL_NAME.STRUCT_EQTABLE:
466
                    with torch.no_grad():
467
468
469
                        table_result = self.table_model.predict(new_image, "html")
                        if len(table_result) > 0:
                            html_code = table_result[0]
470
                elif self.table_model_name == MODEL_NAME.TABLE_MASTER:
471
                    html_code = self.table_model.img2html(new_image)
472
                elif self.table_model_name == MODEL_NAME.RAPID_TABLE:
473
474
                    ocr_result, _ = self.ocr_engine(np.asarray(new_image))
                    html_code, table_cell_bboxes, elapse = self.table_model(np.asarray(new_image), ocr_result)
drunkpig's avatar
drunkpig committed
475

476
                run_time = time.time() - single_table_start_time
477
                # logger.info(f"------------table recognition processing ends within {run_time}s-----")
478
479
480
                if run_time > self.table_max_time:
                    logger.warning(f"------------table recognition processing exceeds max time {self.table_max_time}s----------")
                # 判断是否返回正常
481
482

                if latex_code:
483
                    expected_ending = latex_code.strip().endswith('end{tabular}') or latex_code.strip().endswith('end{table}')
484
485
486
                    if expected_ending:
                        res["latex"] = latex_code
                    else:
487
                        logger.warning(f"table recognition processing fails, not found expected LaTeX table end")
488
                elif html_code:
489
490
491
492
493
                    expected_ending = html_code.strip().endswith('</html>') or html_code.strip().endswith('</table>')
                    if expected_ending:
                        res["html"] = html_code
                    else:
                        logger.warning(f"table recognition processing fails, not found expected HTML table end")
494
                else:
495
496
497
498
                    logger.warning(f"table recognition processing fails, not get latex or html return")
            logger.info(f"table time: {round(time.time() - table_start, 2)}")

        logger.info(f"-----page total time: {round(time.time() - page_start, 2)}-----")
499

500
        return layout_res
liukaiwen's avatar
liukaiwen committed
501
502