pdf_extract_kit.py 8.04 KB
Newer Older
赵小蒙's avatar
update:  
赵小蒙 committed
1
import os
2
3
4
import time

import cv2
赵小蒙's avatar
update:  
赵小蒙 committed
5
6
import numpy as np
import yaml
7
from PIL import Image
赵小蒙's avatar
update:  
赵小蒙 committed
8
from ultralytics import YOLO
9
from loguru import logger
10

11
from magic_pdf.model.pek_sub_modules.layoutlmv3.model_init import Layoutlmv3_Predictor
赵小蒙's avatar
update:  
赵小蒙 committed
12
from unimernet.common.config import Config
13
import unimernet.tasks as tasks
赵小蒙's avatar
update:  
赵小蒙 committed
14
from unimernet.processors import load_processor
15
16
import argparse
from torchvision import transforms
17
from torch.utils.data import Dataset, DataLoader
赵小蒙's avatar
update:  
赵小蒙 committed
18

19
from magic_pdf.model.pek_sub_modules.post_process import get_croped_image, latex_rm_whitespace
20
from magic_pdf.model.pek_sub_modules.self_modify import ModifiedPaddleOCR
赵小蒙's avatar
update:  
赵小蒙 committed
21
22


23
24
def layout_model_init(weight, config_file, device):
    model = Layoutlmv3_Predictor(weight, config_file, device)
25
    return model
赵小蒙's avatar
update:  
赵小蒙 committed
26
27


28
29
30
31
32
33
34
35
36
37
38
def mfr_model_init(weight_dir, cfg_path, device='cpu'):
    args = argparse.Namespace(cfg_path=cfg_path, options=None)
    cfg = Config(args)
    cfg.config.model.pretrained = os.path.join(weight_dir, "pytorch_model.bin")
    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)
    model = model.to(device)
    vis_processor = load_processor('formula_image_eval', cfg.config.datasets.formula_rec_eval.vis_processor.eval)
    return model, vis_processor
赵小蒙's avatar
update:  
赵小蒙 committed
39
40


41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
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)
        return image


60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
class CustomPEKModel:
    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')
        with open(config_path, "r") as f:
            self.configs = yaml.load(f, Loader=yaml.FullLoader)
        # 初始化解析配置
        self.apply_layout = kwargs.get("apply_layout", self.configs["config"]["layout"])
        self.apply_formula = kwargs.get("apply_formula", self.configs["config"]["formula"])
        self.apply_ocr = ocr
        logger.info(
            "DocAnalysis init, this may take some times. apply_layout: {}, apply_formula: {}, apply_ocr: {}".format(
                self.apply_layout, self.apply_formula, self.apply_ocr
赵小蒙's avatar
update:  
赵小蒙 committed
84
            )
85
86
87
        )
        assert self.apply_layout, "DocAnalysis must contain layout model."
        # 初始化解析方案
88
        self.device = kwargs.get("device", self.configs["config"]["device"])
89
        logger.info("using device: {}".format(self.device))
90
        models_dir = kwargs.get("models_dir", os.path.join(root_dir, "resources", "models"))
91
92
        # 初始化layout模型
        self.layout_model = layout_model_init(
93
            os.path.join(models_dir, self.configs['weights']['layout']),
94
95
            os.path.join(model_config_dir, "layoutlmv3", "layoutlmv3_base_inference.yaml"),
            device=self.device
96
97
98
99
        )
        # 初始化公式识别
        if self.apply_formula:
            # 初始化公式检测模型
100
            self.mfd_model = YOLO(model=str(os.path.join(models_dir, self.configs["weights"]["mfd"])))
101
102
103
            # 初始化公式解析模型
            mfr_config_path = os.path.join(model_config_dir, 'UniMERNet', 'demo.yaml')
            self.mfr_model, mfr_vis_processors = mfr_model_init(
104
105
106
107
                os.path.join(models_dir, self.configs["weights"]["mfr"]),
                mfr_config_path,
                device=self.device
            )
108
109
110
111
            self.mfr_transform = transforms.Compose([mfr_vis_processors, ])
        # 初始化ocr
        if self.apply_ocr:
            self.ocr_model = ModifiedPaddleOCR(show_log=show_log)
赵小蒙's avatar
update:  
赵小蒙 committed
112

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

115
116
    def __call__(self, image):

117
118
119
        latex_filling_list = []
        mf_image_list = []

120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
        # layout检测
        layout_start = time.time()
        layout_res = self.layout_model(image, ignore_catids=[])
        layout_cost = round(time.time() - layout_start, 2)
        logger.info(f"layout detection cost: {layout_cost}")

        # 公式检测
        mfd_res = self.mfd_model.predict(image, imgsz=1888, conf=0.25, iou=0.45, verbose=True)[0]
        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)
            bbox_img = get_croped_image(Image.fromarray(image), [xmin, ymin, xmax, ymax])
            mf_image_list.append(bbox_img)

        # 公式识别
        mfr_start = time.time()
143
        dataset = MathDataset(mf_image_list, transform=self.mfr_transform)
144
        dataloader = DataLoader(dataset, batch_size=64, num_workers=0)
145
        mfr_res = []
146
147
148
        for mf_img in dataloader:
            mf_img = mf_img.to(self.device)
            output = self.mfr_model.generate({'image': mf_img})
149
150
151
            mfr_res.extend(output['pred_str'])
        for res, latex in zip(latex_filling_list, mfr_res):
            res['latex'] = latex_rm_whitespace(latex)
152
153
        mfr_cost = round(time.time() - mfr_start, 2)
        logger.info(f"formula nums: {len(mf_image_list)}, mfr time: {mfr_cost}")
154

myhloli's avatar
myhloli committed
155
        # ocr识别
156
        if self.apply_ocr:
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
            ocr_start = time.time()
            pil_img = Image.fromarray(image)
            single_page_mfdetrec_res = []
            for res in layout_res:
                if int(res['category_id']) in [13, 14]:
                    xmin, ymin = int(res['poly'][0]), int(res['poly'][1])
                    xmax, ymax = int(res['poly'][4]), int(res['poly'][5])
                    single_page_mfdetrec_res.append({
                        "bbox": [xmin, ymin, xmax, ymax],
                    })
            for res in layout_res:
                if int(res['category_id']) in [0, 1, 2, 4, 6, 7]:  # 需要进行ocr的类别
                    xmin, ymin = int(res['poly'][0]), int(res['poly'][1])
                    xmax, ymax = int(res['poly'][4]), int(res['poly'][5])
                    crop_box = (xmin, ymin, xmax, ymax)
                    cropped_img = Image.new('RGB', pil_img.size, 'white')
                    cropped_img.paste(pil_img.crop(crop_box), crop_box)
                    cropped_img = cv2.cvtColor(np.asarray(cropped_img), cv2.COLOR_RGB2BGR)
                    ocr_res = self.ocr_model.ocr(cropped_img, mfd_res=single_page_mfdetrec_res)[0]
                    if ocr_res:
                        for box_ocr_res in ocr_res:
                            p1, p2, p3, p4 = box_ocr_res[0]
                            text, score = box_ocr_res[1]
                            layout_res.append({
                                'category_id': 15,
                                'poly': p1 + p2 + p3 + p4,
                                'score': round(score, 2),
                                'text': text,
                            })
            ocr_cost = round(time.time() - ocr_start, 2)
            logger.info(f"ocr cost: {ocr_cost}")

        return layout_res