magic-pdf 0.10.5__py3-none-any.whl → 1.0.0__py3-none-any.whl

magic_pdf/libs/draw_bbox.py

@@ -1,7 +1,8 @@
 import fitz
 from magic_pdf.config.constants import CROSS_PAGE
-from magic_pdf.config.ocr_content_type import BlockType, CategoryId, ContentType
-from magic_pdf.data.dataset import PymuDocDataset
+from magic_pdf.config.ocr_content_type import (BlockType, CategoryId,
+                                               ContentType)
+from magic_pdf.data.dataset import Dataset
 from magic_pdf.model.magic_model import MagicModel
 
 
@@ -194,7 +195,7 @@ def draw_layout_bbox(pdf_info, pdf_bytes, out_path, filename):
         )
 
     # Save the PDF
-    pdf_docs.save(f'{out_path}/{filename}_layout.pdf')
+    pdf_docs.save(f'{out_path}/{filename}')
 
 
 def draw_span_bbox(pdf_info, pdf_bytes, out_path, filename):
@@ -282,18 +283,17 @@ def draw_span_bbox(pdf_info, pdf_bytes, out_path, filename):
         draw_bbox_without_number(i, dropped_list, page, [158, 158, 158], False)
 
     # Save the PDF
-    pdf_docs.save(f'{out_path}/{filename}_spans.pdf')
+    pdf_docs.save(f'{out_path}/{filename}')
 
 
-def draw_model_bbox(model_list: list, pdf_bytes, out_path, filename):
+def draw_model_bbox(model_list, dataset: Dataset, out_path, filename):
     dropped_bbox_list = []
     tables_body_list, tables_caption_list, tables_footnote_list = [], [], []
     imgs_body_list, imgs_caption_list, imgs_footnote_list = [], [], []
     titles_list = []
     texts_list = []
     interequations_list = []
-    pdf_docs = fitz.open('pdf', pdf_bytes)
-    magic_model = MagicModel(model_list, PymuDocDataset(pdf_bytes))
+    magic_model = MagicModel(model_list, dataset)
     for i in range(len(model_list)):
         page_dropped_list = []
         tables_body, tables_caption, tables_footnote = [], [], []
@@ -337,7 +337,8 @@ def draw_model_bbox(model_list: list, pdf_bytes, out_path, filename):
         dropped_bbox_list.append(page_dropped_list)
         imgs_footnote_list.append(imgs_footnote)
 
-    for i, page in enumerate(pdf_docs):
+    for i in range(len(dataset)):
+        page = dataset.get_page(i)
         draw_bbox_with_number(
             i, dropped_bbox_list, page, [158, 158, 158], True
         )  # color !
@@ -352,7 +353,7 @@ def draw_model_bbox(model_list: list, pdf_bytes, out_path, filename):
         draw_bbox_with_number(i, interequations_list, page, [0, 255, 0], True)
 
     # Save the PDF
-    pdf_docs.save(f'{out_path}/{filename}_model.pdf')
+    dataset.dump_to_file(f'{out_path}/{filename}')
 
 
 def draw_line_sort_bbox(pdf_info, pdf_bytes, out_path, filename):
@@ -390,20 +391,16 @@ def draw_line_sort_bbox(pdf_info, pdf_bytes, out_path, filename):
     for i, page in enumerate(pdf_docs):
         draw_bbox_with_number(i, layout_bbox_list, page, [255, 0, 0], False)
 
-    pdf_docs.save(f'{out_path}/{filename}_line_sort.pdf')
-
+    pdf_docs.save(f'{out_path}/{filename}')
 
-def draw_layout_sort_bbox(pdf_info, pdf_bytes, out_path, filename):
-    layout_bbox_list = []
 
-    for page in pdf_info:
-        page_block_list = []
-        for block in page['para_blocks']:
-            bbox = block['bbox']
-            page_block_list.append(bbox)
-        layout_bbox_list.append(page_block_list)
+def draw_char_bbox(pdf_bytes, out_path, filename):
     pdf_docs = fitz.open('pdf', pdf_bytes)
     for i, page in enumerate(pdf_docs):
-        draw_bbox_with_number(i, layout_bbox_list, page, [255, 0, 0], False)
-
-    pdf_docs.save(f'{out_path}/{filename}_layout_sort.pdf')
+        for block in page.get_text('rawdict', flags=fitz.TEXT_PRESERVE_LIGATURES | fitz.TEXT_PRESERVE_WHITESPACE | fitz.TEXT_MEDIABOX_CLIP)['blocks']:
+            for line in block['lines']:
+                for span in line['spans']:
+                    for char in span['chars']:
+                        char_bbox = char['bbox']
+                        page.draw_rect(char_bbox, color=[1, 0, 0], fill=None, fill_opacity=1, width=0.3, overlay=True,)
+    pdf_docs.save(f'{out_path}/{filename}')

magic_pdf/libs/language.py

@@ -16,11 +16,14 @@ def detect_lang(text: str) -> str:
 
     if len(text) == 0:
         return ""
+
+    text = text.replace("\n", "")
     try:
         lang_upper = detect_language(text)
     except:
         html_no_ctrl_chars = ''.join([l for l in text if unicodedata.category(l)[0] not in ['C', ]])
         lang_upper = detect_language(html_no_ctrl_chars)
+
     try:
         lang = lang_upper.lower()
     except:

magic_pdf/libs/pdf_check.py

@@ -1,9 +1,9 @@
 import fitz
 import numpy as np
 from loguru import logger
-# import re
-# from io import BytesIO
-# from pdfminer.high_level import extract_text
+import re
+from io import BytesIO
+from pdfminer.high_level import extract_text
 
 
 def calculate_sample_count(total_page: int):
@@ -33,33 +33,33 @@ def extract_pages(src_pdf_bytes: bytes) -> fitz.Document:
     return sample_docs
 
 
-# def detect_invalid_chars(src_pdf_bytes: bytes) -> bool:
-#     """"
-#     检测PDF中是否包含非法字符
-#     """
-#     '''pdfminer比较慢,需要先随机抽取10页左右的sample'''
-#     sample_docs = extract_pages(src_pdf_bytes)
-#     sample_pdf_bytes = sample_docs.tobytes()
-#     sample_pdf_file_like_object = BytesIO(sample_pdf_bytes)
-#     text = extract_text(sample_pdf_file_like_object)
-#     text = text.replace("\n", "")
-#     # logger.info(text)
-#     '''乱码文本用pdfminer提取出来的文本特征是(cid:xxx)'''
-#     cid_pattern = re.compile(r'\(cid:\d+\)')
-#     matches = cid_pattern.findall(text)
-#     cid_count = len(matches)
-#     cid_len = sum(len(match) for match in matches)
-#     text_len = len(text)
-#     if text_len == 0:
-#         cid_chars_radio = 0
-#     else:
-#         cid_chars_radio = cid_count/(cid_count + text_len - cid_len)
-#     logger.info(f"cid_count: {cid_count}, text_len: {text_len}, cid_chars_radio: {cid_chars_radio}")
-#     '''当一篇文章存在5%以上的文本是乱码时,认为该文档为乱码文档'''
-#     if cid_chars_radio > 0.05:
-#         return False  # 乱码文档
-#     else:
-#         return True   # 正常文档
+def detect_invalid_chars(src_pdf_bytes: bytes) -> bool:
+    """"
+    检测PDF中是否包含非法字符
+    """
+    '''pdfminer比较慢,需要先随机抽取10页左右的sample'''
+    sample_docs = extract_pages(src_pdf_bytes)
+    sample_pdf_bytes = sample_docs.tobytes()
+    sample_pdf_file_like_object = BytesIO(sample_pdf_bytes)
+    text = extract_text(sample_pdf_file_like_object)
+    text = text.replace("\n", "")
+    # logger.info(text)
+    '''乱码文本用pdfminer提取出来的文本特征是(cid:xxx)'''
+    cid_pattern = re.compile(r'\(cid:\d+\)')
+    matches = cid_pattern.findall(text)
+    cid_count = len(matches)
+    cid_len = sum(len(match) for match in matches)
+    text_len = len(text)
+    if text_len == 0:
+        cid_chars_radio = 0
+    else:
+        cid_chars_radio = cid_count/(cid_count + text_len - cid_len)
+    logger.info(f"cid_count: {cid_count}, text_len: {text_len}, cid_chars_radio: {cid_chars_radio}")
+    '''当一篇文章存在5%以上的文本是乱码时,认为该文档为乱码文档'''
+    if cid_chars_radio > 0.05:
+        return False  # 乱码文档
+    else:
+        return True   # 正常文档
 
 
 def count_replacement_characters(text: str) -> int:

magic_pdf/libs/version.py

@@ -1 +1 @@
-__version__ = "0.10.5"
+__version__ = "1.0.0"

magic_pdf/model/__init__.py

@@ -1,2 +1,2 @@
 __use_inside_model__ = True
-__model_mode__ = "full"
+__model_mode__ = 'full'

magic_pdf/model/batch_analyze.py

@@ -0,0 +1,275 @@
+import time
+
+import cv2
+import numpy as np
+import torch
+from loguru import logger
+from PIL import Image
+
+from magic_pdf.config.constants import MODEL_NAME
+from magic_pdf.config.exceptions import CUDA_NOT_AVAILABLE
+from magic_pdf.data.dataset import Dataset
+from magic_pdf.libs.clean_memory import clean_memory
+from magic_pdf.libs.config_reader import get_device
+from magic_pdf.model.doc_analyze_by_custom_model import ModelSingleton
+from magic_pdf.model.pdf_extract_kit import CustomPEKModel
+from magic_pdf.model.sub_modules.model_utils import (
+    clean_vram, crop_img, get_res_list_from_layout_res)
+from magic_pdf.model.sub_modules.ocr.paddleocr.ocr_utils import (
+    get_adjusted_mfdetrec_res, get_ocr_result_list)
+from magic_pdf.operators.models import InferenceResult
+
+YOLO_LAYOUT_BASE_BATCH_SIZE = 4
+MFD_BASE_BATCH_SIZE = 1
+MFR_BASE_BATCH_SIZE = 16
+
+
+class BatchAnalyze:
+    def __init__(self, model: CustomPEKModel, batch_ratio: int):
+        self.model = model
+        self.batch_ratio = batch_ratio
+
+    def __call__(self, images: list) -> list:
+        images_layout_res = []
+
+        layout_start_time = time.time()
+        if self.model.layout_model_name == MODEL_NAME.LAYOUTLMv3:
+            # layoutlmv3
+            for image in images:
+                layout_res = self.model.layout_model(image, ignore_catids=[])
+                images_layout_res.append(layout_res)
+        elif self.model.layout_model_name == MODEL_NAME.DocLayout_YOLO:
+            # doclayout_yolo
+            layout_images = []
+            modified_images = []
+            for image_index, image in enumerate(images):
+                pil_img = Image.fromarray(image)
+                width, height = pil_img.size
+                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
+                    )
+                    layout_images.append(new_image)
+                    modified_images.append([image_index, useful_list])
+                else:
+                    layout_images.append(pil_img)
+
+            images_layout_res += self.model.layout_model.batch_predict(
+                layout_images, self.batch_ratio * YOLO_LAYOUT_BASE_BATCH_SIZE
+            )
+
+            for image_index, useful_list in modified_images:
+                for res in images_layout_res[image_index]:
+                    for i in range(len(res['poly'])):
+                        if i % 2 == 0:
+                            res['poly'][i] = (
+                                res['poly'][i] - useful_list[0] + useful_list[2]
+                            )
+                        else:
+                            res['poly'][i] = (
+                                res['poly'][i] - useful_list[1] + useful_list[3]
+                            )
+        logger.info(
+            f'layout time: {round(time.time() - layout_start_time, 2)}, image num: {len(images)}'
+        )
+
+        if self.model.apply_formula:
+            # 公式检测
+            mfd_start_time = time.time()
+            images_mfd_res = self.model.mfd_model.batch_predict(
+                images, self.batch_ratio * MFD_BASE_BATCH_SIZE
+            )
+            logger.info(
+                f'mfd time: {round(time.time() - mfd_start_time, 2)}, image num: {len(images)}'
+            )
+
+            # 公式识别
+            mfr_start_time = time.time()
+            images_formula_list = self.model.mfr_model.batch_predict(
+                images_mfd_res,
+                images,
+                batch_size=self.batch_ratio * MFR_BASE_BATCH_SIZE,
+            )
+            for image_index in range(len(images)):
+                images_layout_res[image_index] += images_formula_list[image_index]
+            logger.info(
+                f'mfr time: {round(time.time() - mfr_start_time, 2)}, image num: {len(images)}'
+            )
+
+        # 清理显存
+        clean_vram(self.model.device, vram_threshold=8)
+
+        ocr_time = 0
+        ocr_count = 0
+        table_time = 0
+        table_count = 0
+        # reference: magic_pdf/model/doc_analyze_by_custom_model.py:doc_analyze
+        for index in range(len(images)):
+            layout_res = images_layout_res[index]
+            pil_img = Image.fromarray(images[index])
+
+            ocr_res_list, table_res_list, single_page_mfdetrec_res = (
+                get_res_list_from_layout_res(layout_res)
+            )
+            # ocr识别
+            ocr_start = time.time()
+            # Process each area that requires OCR processing
+            for res in ocr_res_list:
+                new_image, useful_list = crop_img(
+                    res, pil_img, crop_paste_x=50, crop_paste_y=50
+                )
+                adjusted_mfdetrec_res = get_adjusted_mfdetrec_res(
+                    single_page_mfdetrec_res, useful_list
+                )
+
+                # OCR recognition
+                new_image = cv2.cvtColor(np.asarray(new_image), cv2.COLOR_RGB2BGR)
+
+                if self.model.apply_ocr:
+                    ocr_res = self.model.ocr_model.ocr(
+                        new_image, mfd_res=adjusted_mfdetrec_res
+                    )[0]
+                else:
+                    ocr_res = self.model.ocr_model.ocr(
+                        new_image, mfd_res=adjusted_mfdetrec_res, rec=False
+                    )[0]
+
+                # Integration results
+                if ocr_res:
+                    ocr_result_list = get_ocr_result_list(ocr_res, useful_list)
+                    layout_res.extend(ocr_result_list)
+            ocr_time += time.time() - ocr_start
+            ocr_count += len(ocr_res_list)
+
+            # 表格识别 table recognition
+            if self.model.apply_table:
+                table_start = time.time()
+                for res in table_res_list:
+                    new_image, _ = crop_img(res, pil_img)
+                    single_table_start_time = time.time()
+                    html_code = None
+                    if self.model.table_model_name == MODEL_NAME.STRUCT_EQTABLE:
+                        with torch.no_grad():
+                            table_result = self.model.table_model.predict(
+                                new_image, 'html'
+                            )
+                            if len(table_result) > 0:
+                                html_code = table_result[0]
+                    elif self.model.table_model_name == MODEL_NAME.TABLE_MASTER:
+                        html_code = self.model.table_model.img2html(new_image)
+                    elif self.model.table_model_name == MODEL_NAME.RAPID_TABLE:
+                        html_code, table_cell_bboxes, elapse = (
+                            self.model.table_model.predict(new_image)
+                        )
+                    run_time = time.time() - single_table_start_time
+                    if run_time > self.model.table_max_time:
+                        logger.warning(
+                            f'table recognition processing exceeds max time {self.model.table_max_time}s'
+                        )
+                    # 判断是否返回正常
+                    if html_code:
+                        expected_ending = html_code.strip().endswith(
+                            '</html>'
+                        ) or html_code.strip().endswith('</table>')
+                        if expected_ending:
+                            res['html'] = html_code
+                        else:
+                            logger.warning(
+                                'table recognition processing fails, not found expected HTML table end'
+                            )
+                    else:
+                        logger.warning(
+                            'table recognition processing fails, not get html return'
+                        )
+                table_time += time.time() - table_start
+                table_count += len(table_res_list)
+
+        if self.model.apply_ocr:
+            logger.info(f'ocr time: {round(ocr_time, 2)}, image num: {ocr_count}')
+        else:
+            logger.info(f'det time: {round(ocr_time, 2)}, image num: {ocr_count}')
+        if self.model.apply_table:
+            logger.info(f'table time: {round(table_time, 2)}, image num: {table_count}')
+
+        return images_layout_res
+
+
+def doc_batch_analyze(
+    dataset: Dataset,
+    ocr: bool = False,
+    show_log: bool = False,
+    start_page_id=0,
+    end_page_id=None,
+    lang=None,
+    layout_model=None,
+    formula_enable=None,
+    table_enable=None,
+    batch_ratio: int | None = None,
+) -> InferenceResult:
+    """Perform batch analysis on a document dataset.
+
+    Args:
+        dataset (Dataset): The dataset containing document pages to be analyzed.
+        ocr (bool, optional): Flag to enable OCR (Optical Character Recognition). Defaults to False.
+        show_log (bool, optional): Flag to enable logging. Defaults to False.
+        start_page_id (int, optional): The starting page ID for analysis. Defaults to 0.
+        end_page_id (int, optional): The ending page ID for analysis. Defaults to None, which means analyze till the last page.
+        lang (str, optional): Language for OCR. Defaults to None.
+        layout_model (optional): Layout model to be used for analysis. Defaults to None.
+        formula_enable (optional): Flag to enable formula detection. Defaults to None.
+        table_enable (optional): Flag to enable table detection. Defaults to None.
+        batch_ratio (int | None, optional): Ratio for batch processing. Defaults to None, which sets it to 1.
+
+    Raises:
+        CUDA_NOT_AVAILABLE: If CUDA is not available, raises an exception as batch analysis is not supported in CPU mode.
+
+    Returns:
+        InferenceResult: The result of the batch analysis containing the analyzed data and the dataset.
+    """
+
+    if not torch.cuda.is_available():
+        raise CUDA_NOT_AVAILABLE('batch analyze not support in CPU mode')
+
+    lang = None if lang == '' else lang
+    # TODO: auto detect batch size
+    batch_ratio = 1 if batch_ratio is None else batch_ratio
+    end_page_id = end_page_id if end_page_id else len(dataset)
+
+    model_manager = ModelSingleton()
+    custom_model: CustomPEKModel = model_manager.get_model(
+        ocr, show_log, lang, layout_model, formula_enable, table_enable
+    )
+    batch_model = BatchAnalyze(model=custom_model, batch_ratio=batch_ratio)
+
+    model_json = []
+
+    # batch analyze
+    images = []
+    for index in range(len(dataset)):
+        if start_page_id <= index <= end_page_id:
+            page_data = dataset.get_page(index)
+            img_dict = page_data.get_image()
+            images.append(img_dict['img'])
+    analyze_result = batch_model(images)
+
+    for index in range(len(dataset)):
+        page_data = dataset.get_page(index)
+        img_dict = page_data.get_image()
+        page_width = img_dict['width']
+        page_height = img_dict['height']
+        if start_page_id <= index <= end_page_id:
+            result = analyze_result.pop(0)
+        else:
+            result = []
+
+        page_info = {'page_no': index, 'height': page_height, 'width': page_width}
+        page_dict = {'layout_dets': result, 'page_info': page_info}
+        model_json.append(page_dict)
+
+    # TODO: clean memory when gpu memory is not enough
+    clean_memory_start_time = time.time()
+    clean_memory(get_device())
+    logger.info(f'clean memory time: {round(time.time() - clean_memory_start_time, 2)}')
+
+    return InferenceResult(model_json, dataset)