magic-pdf 1.0.1__py3-none-any.whl → 1.2.0__py3-none-any.whl

magic_pdf/model/magic_model.py

@@ -450,11 +450,168 @@ class MagicModel:
             )
         return ret
 
+
+    def __tie_up_category_by_distance_v3(
+        self,
+        page_no: int,
+        subject_category_id: int,
+        object_category_id: int,
+        priority_pos: PosRelationEnum,
+    ):
+        subjects = self.__reduct_overlap(
+            list(
+                map(
+                    lambda x: {'bbox': x['bbox'], 'score': x['score']},
+                    filter(
+                        lambda x: x['category_id'] == subject_category_id,
+                        self.__model_list[page_no]['layout_dets'],
+                    ),
+                )
+            )
+        )
+        objects = self.__reduct_overlap(
+            list(
+                map(
+                    lambda x: {'bbox': x['bbox'], 'score': x['score']},
+                    filter(
+                        lambda x: x['category_id'] == object_category_id,
+                        self.__model_list[page_no]['layout_dets'],
+                    ),
+                )
+            )
+        )
+
+        ret = []
+        N, M = len(subjects), len(objects)
+        subjects.sort(key=lambda x: x['bbox'][0] ** 2 + x['bbox'][1] ** 2)
+        objects.sort(key=lambda x: x['bbox'][0] ** 2 + x['bbox'][1] ** 2)
+
+        OBJ_IDX_OFFSET = 10000
+        SUB_BIT_KIND, OBJ_BIT_KIND = 0, 1
+
+        all_boxes_with_idx = [(i, SUB_BIT_KIND, sub['bbox'][0], sub['bbox'][1]) for i, sub in enumerate(subjects)] + [(i + OBJ_IDX_OFFSET , OBJ_BIT_KIND, obj['bbox'][0], obj['bbox'][1]) for i, obj in enumerate(objects)]
+        seen_idx = set()
+        seen_sub_idx = set()
+
+        while N > len(seen_sub_idx):
+            candidates = []
+            for idx, kind, x0, y0 in all_boxes_with_idx:
+                if idx in seen_idx:
+                    continue
+                candidates.append((idx, kind, x0, y0))
+
+            if len(candidates) == 0:
+                break
+            left_x = min([v[2] for v in candidates])
+            top_y =  min([v[3] for v in candidates])
+
+            candidates.sort(key=lambda x: (x[2]-left_x) ** 2 + (x[3] - top_y) ** 2)
+
+
+            fst_idx, fst_kind, left_x, top_y = candidates[0]
+            candidates.sort(key=lambda x: (x[2] - left_x) ** 2 + (x[3] - top_y)**2)
+            nxt = None
+
+            for i in range(1, len(candidates)):
+                if candidates[i][1] ^ fst_kind == 1:
+                    nxt = candidates[i]
+                    break
+            if nxt is None:
+                break
+
+            if fst_kind == SUB_BIT_KIND:
+                sub_idx, obj_idx = fst_idx, nxt[0] - OBJ_IDX_OFFSET
+
+            else:
+                sub_idx, obj_idx = nxt[0], fst_idx - OBJ_IDX_OFFSET
+
+            pair_dis = bbox_distance(subjects[sub_idx]['bbox'], objects[obj_idx]['bbox'])
+            nearest_dis = float('inf')
+            for i in range(N):
+                if i in seen_idx:continue
+                nearest_dis = min(nearest_dis, bbox_distance(subjects[i]['bbox'], objects[obj_idx]['bbox']))
+
+            if pair_dis >= 3*nearest_dis:
+                seen_idx.add(sub_idx)
+                continue
+
+
+            seen_idx.add(sub_idx)
+            seen_idx.add(obj_idx + OBJ_IDX_OFFSET)
+            seen_sub_idx.add(sub_idx)
+
+            ret.append(
+                {
+                    'sub_bbox': {
+                        'bbox': subjects[sub_idx]['bbox'],
+                        'score': subjects[sub_idx]['score'],
+                    },
+                    'obj_bboxes': [
+                        {'score': objects[obj_idx]['score'], 'bbox': objects[obj_idx]['bbox']}
+                    ],
+                    'sub_idx': sub_idx,
+                }
+            )
+
+        for i in range(len(objects)):
+            j = i + OBJ_IDX_OFFSET
+            if j in seen_idx:
+                continue
+            seen_idx.add(j)
+            nearest_dis, nearest_sub_idx = float('inf'), -1
+            for k in range(len(subjects)):
+                dis = bbox_distance(objects[i]['bbox'], subjects[k]['bbox'])
+                if dis < nearest_dis:
+                    nearest_dis = dis
+                    nearest_sub_idx = k
+
+            for k in range(len(subjects)):
+                if k != nearest_sub_idx: continue
+                if k in seen_sub_idx:
+                    for kk in range(len(ret)):
+                        if ret[kk]['sub_idx'] == k:
+                            ret[kk]['obj_bboxes'].append({'score': objects[i]['score'], 'bbox': objects[i]['bbox']})
+                            break
+                else:
+                    ret.append(
+                        {
+                            'sub_bbox': {
+                                'bbox': subjects[k]['bbox'],
+                                'score': subjects[k]['score'],
+                            },
+                            'obj_bboxes': [
+                                {'score': objects[i]['score'], 'bbox': objects[i]['bbox']}
+                            ],
+                            'sub_idx': k,
+                        }
+                    )
+                seen_sub_idx.add(k)
+                seen_idx.add(k)
+
+
+        for i in range(len(subjects)):
+            if i in seen_sub_idx:
+                continue
+            ret.append(
+                {
+                    'sub_bbox': {
+                        'bbox': subjects[i]['bbox'],
+                        'score': subjects[i]['score'],
+                    },
+                    'obj_bboxes': [],
+                    'sub_idx': i,
+                }
+            )
+
+
+        return ret
+
+
     def get_imgs_v2(self, page_no: int):
-        with_captions = self.__tie_up_category_by_distance_v2(
+        with_captions = self.__tie_up_category_by_distance_v3(
             page_no, 3, 4, PosRelationEnum.BOTTOM
         )
-        with_footnotes = self.__tie_up_category_by_distance_v2(
+        with_footnotes = self.__tie_up_category_by_distance_v3(
             page_no, 3, CategoryId.ImageFootnote, PosRelationEnum.ALL
         )
         ret = []
@@ -470,10 +627,10 @@ class MagicModel:
         return ret
 
     def get_tables_v2(self, page_no: int) -> list:
-        with_captions = self.__tie_up_category_by_distance_v2(
+        with_captions = self.__tie_up_category_by_distance_v3(
             page_no, 5, 6, PosRelationEnum.UP
         )
-        with_footnotes = self.__tie_up_category_by_distance_v2(
+        with_footnotes = self.__tie_up_category_by_distance_v3(
             page_no, 5, 7, PosRelationEnum.ALL
         )
         ret = []

magic_pdf/model/pdf_extract_kit.py

@@ -69,6 +69,7 @@ class CustomPEKModel:
         self.apply_table = self.table_config.get('enable', False)
         self.table_max_time = self.table_config.get('max_time', TABLE_MAX_TIME_VALUE)
         self.table_model_name = self.table_config.get('model', MODEL_NAME.RAPID_TABLE)
+        self.table_sub_model_name = self.table_config.get('sub_model', None)
 
         # ocr config
         self.apply_ocr = ocr
@@ -88,13 +89,6 @@ class CustomPEKModel:
         # 初始化解析方案
         self.device = kwargs.get('device', 'cpu')
 
-        if str(self.device).startswith("npu"):
-            import torch_npu
-            os.environ['FLAGS_npu_jit_compile'] = '0'
-            os.environ['FLAGS_use_stride_kernel'] = '0'
-        elif str(self.device).startswith("mps"):
-            os.environ['PYTORCH_ENABLE_MPS_FALLBACK'] = '1'
-
         logger.info('using device: {}'.format(self.device))
         models_dir = kwargs.get(
             'models_dir', os.path.join(root_dir, 'resources', 'models')
@@ -144,7 +138,7 @@ class CustomPEKModel:
                         model_config_dir, 'layoutlmv3', 'layoutlmv3_base_inference.yaml'
                     )
                 ),
-                device=self.device,
+                device='cpu' if str(self.device).startswith("mps") else self.device,
             )
         elif self.layout_model_name == MODEL_NAME.DocLayout_YOLO:
             self.layout_model = atom_model_manager.get_atom_model(
@@ -174,6 +168,7 @@ class CustomPEKModel:
                 table_max_time=self.table_max_time,
                 device=self.device,
                 ocr_engine=self.ocr_model,
+                table_sub_model_name=self.table_sub_model_name
             )
 
         logger.info('DocAnalysis init done!')
@@ -192,24 +187,24 @@ class CustomPEKModel:
             layout_res = self.layout_model(image, ignore_catids=[])
         elif self.layout_model_name == MODEL_NAME.DocLayout_YOLO:
             # doclayout_yolo
-            if height > width:
-                input_res = {"poly":[0,0,width,0,width,height,0,height]}
-                new_image, useful_list = crop_img(input_res, pil_img, crop_paste_x=width//2, crop_paste_y=0)
-                paste_x, paste_y, xmin, ymin, xmax, ymax, new_width, new_height = useful_list
-                layout_res = self.layout_model.predict(new_image)
-                for res in layout_res:
-                    p1, p2, p3, p4, p5, p6, p7, p8 = res['poly']
-                    p1 = p1 - paste_x + xmin
-                    p2 = p2 - paste_y + ymin
-                    p3 = p3 - paste_x + xmin
-                    p4 = p4 - paste_y + ymin
-                    p5 = p5 - paste_x + xmin
-                    p6 = p6 - paste_y + ymin
-                    p7 = p7 - paste_x + xmin
-                    p8 = p8 - paste_y + ymin
-                    res['poly'] = [p1, p2, p3, p4, p5, p6, p7, p8]
-            else:
-                layout_res = self.layout_model.predict(image)
+            # if height > width:
+            #     input_res = {"poly":[0,0,width,0,width,height,0,height]}
+            #     new_image, useful_list = crop_img(input_res, pil_img, crop_paste_x=width//2, crop_paste_y=0)
+            #     paste_x, paste_y, xmin, ymin, xmax, ymax, new_width, new_height = useful_list
+            #     layout_res = self.layout_model.predict(new_image)
+            #     for res in layout_res:
+            #         p1, p2, p3, p4, p5, p6, p7, p8 = res['poly']
+            #         p1 = p1 - paste_x + xmin
+            #         p2 = p2 - paste_y + ymin
+            #         p3 = p3 - paste_x + xmin
+            #         p4 = p4 - paste_y + ymin
+            #         p5 = p5 - paste_x + xmin
+            #         p6 = p6 - paste_y + ymin
+            #         p7 = p7 - paste_x + xmin
+            #         p8 = p8 - paste_y + ymin
+            #         res['poly'] = [p1, p2, p3, p4, p5, p6, p7, p8]
+            # else:
+            layout_res = self.layout_model.predict(image)
 
         layout_cost = round(time.time() - layout_start, 2)
         logger.info(f'layout detection time: {layout_cost}')
@@ -228,7 +223,7 @@ class CustomPEKModel:
             logger.info(f'formula nums: {len(formula_list)}, mfr time: {mfr_cost}')
 
         # 清理显存
-        clean_vram(self.device, vram_threshold=8)
+        clean_vram(self.device, vram_threshold=6)
 
         # 从layout_res中获取ocr区域、表格区域、公式区域
         ocr_res_list, table_res_list, single_page_mfdetrec_res = (
@@ -276,7 +271,7 @@ class CustomPEKModel:
                 elif self.table_model_name == MODEL_NAME.TABLE_MASTER:
                     html_code = self.table_model.img2html(new_image)
                 elif self.table_model_name == MODEL_NAME.RAPID_TABLE:
-                    html_code, table_cell_bboxes, elapse = self.table_model.predict(
+                    html_code, table_cell_bboxes, logic_points, elapse = self.table_model.predict(
                         new_image
                     )
                 run_time = time.time() - single_table_start_time

magic_pdf/model/sub_modules/language_detection/yolov11/YOLOv11.py

@@ -1,4 +1,5 @@
 # Copyright (c) Opendatalab. All rights reserved.
+import time
 from collections import Counter
 from uuid import uuid4
 
@@ -102,9 +103,9 @@ class YOLOv11LangDetModel(object):
             temp_images = split_images(image)
             for temp_image in temp_images:
                 all_images.append(resize_images_to_224(temp_image))
-
-        images_lang_res = self.batch_predict(all_images, batch_size=8)
-        # logger.info(f"images_lang_res: {images_lang_res}")
+        # langdetect_start = time.time()
+        images_lang_res = self.batch_predict(all_images, batch_size=256)
+        # logger.info(f"image number of langdetect: {len(images_lang_res)}, langdetect time: {round(time.time() - langdetect_start, 2)}")
         if len(images_lang_res) > 0:
             count_dict = Counter(images_lang_res)
             language = max(count_dict, key=count_dict.get)

magic_pdf/model/sub_modules/layout/doclayout_yolo/DocLayoutYOLO.py

@@ -9,7 +9,11 @@ class DocLayoutYOLOModel(object):
     def predict(self, image):
         layout_res = []
         doclayout_yolo_res = self.model.predict(
-            image, imgsz=1024, conf=0.25, iou=0.45, verbose=False, device=self.device
+            image,
+            imgsz=1280,
+            conf=0.10,
+            iou=0.45,
+            verbose=False, device=self.device
         )[0]
         for xyxy, conf, cla in zip(
             doclayout_yolo_res.boxes.xyxy.cpu(),
@@ -32,8 +36,8 @@ class DocLayoutYOLOModel(object):
                 image_res.cpu()
                 for image_res in self.model.predict(
                     images[index : index + batch_size],
-                    imgsz=1024,
-                    conf=0.25,
+                    imgsz=1280,
+                    conf=0.10,
                     iou=0.45,
                     verbose=False,
                     device=self.device,

magic_pdf/model/sub_modules/mfr/unimernet/Unimernet.py

@@ -89,7 +89,7 @@ class UnimernetModel(object):
             mf_image_list.append(bbox_img)
 
         dataset = MathDataset(mf_image_list, transform=self.mfr_transform)
-        dataloader = DataLoader(dataset, batch_size=64, num_workers=0)
+        dataloader = DataLoader(dataset, batch_size=32, num_workers=0)
         mfr_res = []
         for mf_img in dataloader:
             mf_img = mf_img.to(self.device)

magic_pdf/model/sub_modules/model_init.py

@@ -4,24 +4,39 @@ from loguru import logger
 from magic_pdf.config.constants import MODEL_NAME
 from magic_pdf.model.model_list import AtomicModel
 from magic_pdf.model.sub_modules.language_detection.yolov11.YOLOv11 import YOLOv11LangDetModel
-from magic_pdf.model.sub_modules.layout.doclayout_yolo.DocLayoutYOLO import \
-    DocLayoutYOLOModel
-from magic_pdf.model.sub_modules.layout.layoutlmv3.model_init import \
-    Layoutlmv3_Predictor
+from magic_pdf.model.sub_modules.layout.doclayout_yolo.DocLayoutYOLO import DocLayoutYOLOModel
+from magic_pdf.model.sub_modules.layout.layoutlmv3.model_init import Layoutlmv3_Predictor
 from magic_pdf.model.sub_modules.mfd.yolov8.YOLOv8 import YOLOv8MFDModel
 from magic_pdf.model.sub_modules.mfr.unimernet.Unimernet import UnimernetModel
-from magic_pdf.model.sub_modules.ocr.paddleocr.ppocr_273_mod import \
-    ModifiedPaddleOCR
-from magic_pdf.model.sub_modules.table.rapidtable.rapid_table import \
-    RapidTableModel
-# from magic_pdf.model.sub_modules.ocr.paddleocr.ppocr_291_mod import ModifiedPaddleOCR
-from magic_pdf.model.sub_modules.table.structeqtable.struct_eqtable import \
-    StructTableModel
-from magic_pdf.model.sub_modules.table.tablemaster.tablemaster_paddle import \
-    TableMasterPaddleModel
-
-
-def table_model_init(table_model_type, model_path, max_time, _device_='cpu', ocr_engine=None):
+
+try:
+    from magic_pdf_ascend_plugin.libs.license_verifier import load_license, LicenseFormatError, LicenseSignatureError, LicenseExpiredError
+    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
+    license_key = load_license()
+    logger.info(f'Using Ascend Plugin Success, License id is {license_key["payload"]["id"]},'
+                f' License expired at {license_key["payload"]["date"]["end_date"]}')
+except Exception as e:
+    if isinstance(e, ImportError):
+        pass
+    elif isinstance(e, LicenseFormatError):
+        logger.error("Ascend Plugin: Invalid license format. Please check the license file.")
+    elif isinstance(e, LicenseSignatureError):
+        logger.error("Ascend Plugin: Invalid signature. The license may be tampered with.")
+    elif isinstance(e, LicenseExpiredError):
+        logger.error("Ascend Plugin: License has expired. Please renew your license.")
+    elif isinstance(e, FileNotFoundError):
+        logger.error("Ascend Plugin: Not found License file.")
+    else:
+        logger.error(f"Ascend Plugin: {e}")
+    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
+
+from magic_pdf.model.sub_modules.table.structeqtable.struct_eqtable import StructTableModel
+from magic_pdf.model.sub_modules.table.tablemaster.tablemaster_paddle import TableMasterPaddleModel
+
+def table_model_init(table_model_type, model_path, max_time, _device_='cpu', ocr_engine=None, table_sub_model_name=None):
     if table_model_type == MODEL_NAME.STRUCT_EQTABLE:
         table_model = StructTableModel(model_path, max_new_tokens=2048, max_time=max_time)
     elif table_model_type == MODEL_NAME.TABLE_MASTER:
@@ -31,7 +46,7 @@ def table_model_init(table_model_type, model_path, max_time, _device_='cpu', ocr
         }
         table_model = TableMasterPaddleModel(config)
     elif table_model_type == MODEL_NAME.RAPID_TABLE:
-        table_model = RapidTableModel(ocr_engine)
+        table_model = RapidTableModel(ocr_engine, table_sub_model_name)
     else:
         logger.error('table model type not allow')
         exit(1)
@@ -76,7 +91,6 @@ def ocr_model_init(show_log: bool = False,
                    use_dilation=True,
                    det_db_unclip_ratio=1.8,
                    ):
-
     if lang is not None and lang != '':
         model = ModifiedPaddleOCR(
             show_log=show_log,
@@ -163,7 +177,8 @@ def atom_model_init(model_name: str, **kwargs):
             kwargs.get('table_model_path'),
             kwargs.get('table_max_time'),
             kwargs.get('device'),
-            kwargs.get('ocr_engine')
+            kwargs.get('ocr_engine'),
+            kwargs.get('table_sub_model_name')
         )
     elif model_name == AtomicModel.LangDetect:
         if kwargs.get('langdetect_model_name') == MODEL_NAME.YOLO_V11_LangDetect:

magic_pdf/model/sub_modules/ocr/paddleocr/ocr_utils.py

@@ -7,6 +7,8 @@ import base64
 from magic_pdf.libs.boxbase import __is_overlaps_y_exceeds_threshold
 from magic_pdf.pre_proc.ocr_dict_merge import merge_spans_to_line
 
+import importlib.resources
+from paddleocr import PaddleOCR
 from ppocr.utils.utility import check_and_read
 
 
@@ -327,30 +329,35 @@ class ONNXModelSingleton:
         return self._models[key]
 
 def onnx_model_init(key):
-
-    import importlib.resources
-
-    resource_path = importlib.resources.path('rapidocr_onnxruntime.models','')
-
-    onnx_model = None
-    additional_ocr_params = {
-        "use_onnx": True,
-        "det_model_dir": f'{resource_path}/ch_PP-OCRv4_det_infer.onnx',
-        "rec_model_dir": f'{resource_path}/ch_PP-OCRv4_rec_infer.onnx',
-        "cls_model_dir": f'{resource_path}/ch_ppocr_mobile_v2.0_cls_infer.onnx',
-        "det_db_box_thresh": key[1],
-        "use_dilation": key[2],
-        "det_db_unclip_ratio": key[3],
-    }
-    # logger.info(f"additional_ocr_params: {additional_ocr_params}")
-    if key[0] is not None:
-        additional_ocr_params["lang"] = key[0]
-
-    from paddleocr import PaddleOCR
-    onnx_model = PaddleOCR(**additional_ocr_params)
-
-    if onnx_model is None:
-        logger.error('model init failed')
+    if len(key) < 4:
+        logger.error('Invalid key length, expected at least 4 elements')
         exit(1)
-    else:
-        return onnx_model
+
+    try:
+        with importlib.resources.path('rapidocr_onnxruntime.models', '') as resource_path:
+            additional_ocr_params = {
+                "use_onnx": True,
+                "det_model_dir": f'{resource_path}/ch_PP-OCRv4_det_infer.onnx',
+                "rec_model_dir": f'{resource_path}/ch_PP-OCRv4_rec_infer.onnx',
+                "cls_model_dir": f'{resource_path}/ch_ppocr_mobile_v2.0_cls_infer.onnx',
+                "det_db_box_thresh": key[1],
+                "use_dilation": key[2],
+                "det_db_unclip_ratio": key[3],
+            }
+
+            if key[0] is not None:
+                additional_ocr_params["lang"] = key[0]
+
+            # logger.info(f"additional_ocr_params: {additional_ocr_params}")
+
+            onnx_model = PaddleOCR(**additional_ocr_params)
+
+            if onnx_model is None:
+                logger.error('model init failed')
+                exit(1)
+            else:
+                return onnx_model
+
+    except Exception as e:
+        logger.exception(f'Error initializing model: {e}')
+        exit(1)

magic_pdf/model/sub_modules/table/rapidtable/rapid_table.py

@@ -2,12 +2,27 @@ import cv2
 import numpy as np
 import torch
 from loguru import logger
-from rapid_table import RapidTable
+from rapid_table import RapidTable, RapidTableInput
+from rapid_table.main import ModelType
+
+from magic_pdf.libs.config_reader import get_device
 
 
 class RapidTableModel(object):
-    def __init__(self, ocr_engine):
-        self.table_model = RapidTable()
+    def __init__(self, ocr_engine, table_sub_model_name):
+        sub_model_list = [model.value for model in ModelType]
+        if table_sub_model_name is None:
+            input_args = RapidTableInput()
+        elif table_sub_model_name in  sub_model_list:
+            if torch.cuda.is_available() and table_sub_model_name == "unitable":
+                input_args = RapidTableInput(model_type=table_sub_model_name, use_cuda=True, device=get_device())
+            else:
+                input_args = RapidTableInput(model_type=table_sub_model_name)
+        else:
+            raise ValueError(f"Invalid table_sub_model_name: {table_sub_model_name}. It must be one of {sub_model_list}")
+
+        self.table_model = RapidTable(input_args)
+
         # if ocr_engine is None:
         #     self.ocr_model_name = "RapidOCR"
         #     if torch.cuda.is_available():
@@ -45,7 +60,11 @@ class RapidTableModel(object):
             ocr_result = None
 
         if ocr_result:
-            html_code, table_cell_bboxes, elapse = self.table_model(np.asarray(image), ocr_result)
-            return html_code, table_cell_bboxes, elapse
+            table_results = self.table_model(np.asarray(image), ocr_result)
+            html_code = table_results.pred_html
+            table_cell_bboxes = table_results.cell_bboxes
+            logic_points = table_results.logic_points
+            elapse = table_results.elapse
+            return html_code, table_cell_bboxes, logic_points, elapse
         else:
-            return None, None, None
+            return None, None, None, None