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

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

20
21
    if torchtext.__version__ >= "0.18.0":
        torchtext.disable_torchtext_deprecation_warning()
myhloli's avatar
myhloli committed
22
23
24
25
26
27
28
    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
29
    from doclayout_yolo import YOLOv10
赵小蒙'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
    if table_model_type == MODEL_NAME.STRUCT_EQTABLE:
47
        table_model = StructTableModel(model_path, max_time=max_time)
48
    elif table_model_type == MODEL_NAME.TABLE_MASTER:
49
50
51
52
53
        config = {
            "model_dir": model_path,
            "device": _device_
        }
        table_model = ppTableModel(config)
54
55
56
    else:
        logger.error("table model type not allow")
        exit(1)
57
    return table_model
58

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

60
61
62
def mfd_model_init(weight):
    mfd_model = YOLO(weight)
    return mfd_model
赵小蒙's avatar
update:  
赵小蒙 committed
63
64


65
def mfr_model_init(weight_dir, cfg_path, _device_='cpu'):
66
67
    args = argparse.Namespace(cfg_path=cfg_path, options=None)
    cfg = Config(args)
68
    cfg.config.model.pretrained = os.path.join(weight_dir, "pytorch_model.pth")
69
70
71
72
    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)
73
    model.to(_device_)
74
    model.eval()
75
    vis_processor = load_processor('formula_image_eval', cfg.config.datasets.formula_rec_eval.vis_processor.eval)
76
77
    mfr_transform = transforms.Compose([vis_processor, ])
    return [model, mfr_transform]
赵小蒙's avatar
update:  
赵小蒙 committed
78
79


80
81
82
83
84
def layout_model_init(weight, config_file, device):
    model = Layoutlmv3_Predictor(weight, config_file, device)
    return model


85
86
87
88
89
def doclayout_yolo_model_init(weight):
    model = YOLOv10(weight)
    return model


90
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):
91
    if lang is not None:
92
        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)
93
    else:
94
        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)
95
96
97
    return model


98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
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)
114
            return image
115
116


117
118
119
120
121
122
123
124
125
126
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):
127
128
129
130
131
132
        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]
133
134
135
136
137


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

    if model_name == AtomicModel.Layout:
138
139
140
141
142
143
144
145
146
147
        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"),
            )
148
149
150
151
152
153
154
155
156
157
158
159
160
    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"),
161
162
            kwargs.get("det_db_box_thresh"),
            kwargs.get("lang")
163
164
165
        )
    elif model_name == AtomicModel.Table:
        atom_model = table_model_init(
166
            kwargs.get("table_model_name"),
167
168
169
170
171
172
173
174
175
176
177
            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


178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
#  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


195
class CustomPEKModel:
196

197
198
199
200
201
202
203
204
205
206
207
208
209
210
    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')
211
        with open(config_path, "r", encoding='utf-8') as f:
212
213
            self.configs = yaml.load(f, Loader=yaml.FullLoader)
        # 初始化解析配置
214
215
216
217
218
219
220
221
222
223
224

        # 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)

225
        # table config
226
227
        self.table_config = kwargs.get("table_config")
        self.apply_table = self.table_config.get("enable", False)
228
        self.table_max_time = self.table_config.get("max_time", TABLE_MAX_TIME_VALUE)
229
230
231
        self.table_model_name = self.table_config.get("model", MODEL_NAME.TABLE_MASTER)

        # ocr config
232
        self.apply_ocr = ocr
233
        self.lang = kwargs.get("lang", None)
234

235
        logger.info(
236
237
238
            "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
239
            )
240
241
        )
        # 初始化解析方案
242
        self.device = kwargs.get("device", "cpu")
243
        logger.info("using device: {}".format(self.device))
244
        models_dir = kwargs.get("models_dir", os.path.join(root_dir, "resources", "models"))
245
        logger.info("using models_dir: {}".format(models_dir))
246

247
248
        atom_model_manager = AtomModelSingleton()

249
250
        # 初始化公式识别
        if self.apply_formula:
251

252
            # 初始化公式检测模型
253
254
            self.mfd_model = atom_model_manager.get_atom_model(
                atom_model_name=AtomicModel.MFD,
255
                mfd_weights=str(os.path.join(models_dir, self.configs["weights"][self.mfd_model_name]))
256
            )
257

258
            # 初始化公式解析模型
259
            mfr_weight_dir = str(os.path.join(models_dir, self.configs["weights"][self.mfr_model_name]))
260
            mfr_cfg_path = str(os.path.join(model_config_dir, "UniMERNet", "demo.yaml"))
261
262
263
264
265
266
            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
            )
267
268

        # 初始化layout模型
269
270
271
272
273
274
275
276
277
278
279
280
281
282
        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]))
            )
283
284
        # 初始化ocr
        if self.apply_ocr:
drunkpig's avatar
drunkpig committed
285

286
287
288
289
            # self.ocr_model = ModifiedPaddleOCR(show_log=show_log, det_db_box_thresh=0.3)
            self.ocr_model = atom_model_manager.get_atom_model(
                atom_model_name=AtomicModel.OCR,
                ocr_show_log=show_log,
290
291
                det_db_box_thresh=0.3,
                lang=self.lang
292
            )
293
        # init table model
294
        if self.apply_table:
295
            table_model_dir = self.configs["weights"][self.table_model_name]
296
297
            self.table_model = atom_model_manager.get_atom_model(
                atom_model_name=AtomicModel.Table,
298
                table_model_name=self.table_model_name,
299
300
301
302
                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
303

304
305
306
307
308
309
310
311
312
313
314
            home_directory = Path.home()
            det_source = os.path.join(models_dir, table_model_dir, DETECT_MODEL_DIR)
            rec_source = os.path.join(models_dir, table_model_dir, REC_MODEL_DIR)
            det_dest_dir = os.path.join(home_directory, PP_DET_DIRECTORY)
            rec_dest_dir = os.path.join(home_directory, PP_REC_DIRECTORY)

            if not os.path.exists(det_dest_dir):
                shutil.copytree(det_source, det_dest_dir)
            if not os.path.exists(rec_dest_dir):
                shutil.copytree(rec_source, rec_dest_dir)

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

317
318
    def __call__(self, image):

319
320
        page_start = time.time()

321
322
323
        latex_filling_list = []
        mf_image_list = []

324
325
        # layout检测
        layout_start = time.time()
326
327
328
329
330
331
        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 = []
332
            doclayout_yolo_res = self.layout_model.predict(image, imgsz=1024, conf=0.25, iou=0.45, verbose=True, device=self.device)[0]
333
334
335
336
337
338
339
340
            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)
341
        layout_cost = round(time.time() - layout_start, 2)
342
        logger.info(f"layout detection time: {layout_cost}")
343

344
345
        pil_img = Image.fromarray(image)

346
347
        if self.apply_formula:
            # 公式检测
348
            mfd_start = time.time()
349
            mfd_res = self.mfd_model.predict(image, imgsz=1888, conf=0.25, iou=0.45, verbose=True, device=self.device)[0]
350
            logger.info(f"mfd time: {round(time.time() - mfd_start, 2)}")
351
352
353
354
355
356
357
358
359
360
            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)
361
                bbox_img = pil_img.crop((xmin, ymin, xmax, ymax))
362
363
364
365
366
367
368
369
370
                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)
371
372
                with torch.no_grad():
                    output = self.mfr_model.generate({'image': mf_img})
373
374
375
376
377
                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}")
378

379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
        # 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)

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

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

427
                # OCR recognition
428
429
                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]
430

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

437
                        # Convert the coordinates back to the original coordinate system
438
439
440
441
442
443
444
445
446
447
448
449
                        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,
                        })

450
            ocr_cost = round(time.time() - ocr_start, 2)
451
            logger.info(f"ocr time: {ocr_cost}")
452

453
454
        # 表格识别 table recognition
        if self.apply_table:
455
456
457
458
            table_start = time.time()
            for res in table_res_list:
                new_image, _ = crop_img(res, pil_img)
                single_table_start_time = time.time()
459
                # logger.info("------------------table recognition processing begins-----------------")
460
461
                latex_code = None
                html_code = None
462
                if self.table_model_name == MODEL_NAME.STRUCT_EQTABLE:
463
                    with torch.no_grad():
464
465
466
                        table_result = self.table_model.predict(new_image, "html")
                        if len(table_result) > 0:
                            html_code = table_result[0]
467
468
                else:
                    html_code = self.table_model.img2html(new_image)
drunkpig's avatar
drunkpig committed
469

470
                run_time = time.time() - single_table_start_time
471
                # logger.info(f"------------table recognition processing ends within {run_time}s-----")
472
473
474
                if run_time > self.table_max_time:
                    logger.warning(f"------------table recognition processing exceeds max time {self.table_max_time}s----------")
                # 判断是否返回正常
475
476

                if latex_code:
477
                    expected_ending = latex_code.strip().endswith('end{tabular}') or latex_code.strip().endswith('end{table}')
478
479
480
                    if expected_ending:
                        res["latex"] = latex_code
                    else:
481
                        logger.warning(f"table recognition processing fails, not found expected LaTeX table end")
482
                elif html_code:
483
484
485
486
487
                    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")
488
                else:
489
490
491
492
                    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)}-----")
493

494
        return layout_res
liukaiwen's avatar
liukaiwen committed
495
496