docvortex 0.2.1__py3-none-any.whl

docvortex/document/pdf/constants.py

@@ -0,0 +1,132 @@
+"""PDF 分析各领域阶段共享的阈值、类型集合与映射。"""
+
+from __future__ import annotations
+
+from ...schema import RAW_ALGORITHM, RAW_CAPTION, RAW_FOOTNOTE, RAW_FORMULA_NUMBER, RAW_PHONETIC, BlockType
+
+LAYOUT_BASE_BATCH_SIZE = 1
+MFR_BASE_BATCH_SIZE = 16
+OCR_DET_BASE_BATCH_SIZE = 8
+LAYOUT_TITLE_SPLIT_OVERLAP_THRESHOLD = 0.8
+IMAGE_BLOCK_CONTAINMENT_THRESHOLD = 0.8
+IMAGE_BLOCK_LAYOUT_COVERAGE_THRESHOLD = 0.9
+IMAGE_BLOCK_LAYOUT_MIN_VISUAL_COUNT = 2
+BATCH_RATIO = 2
+TABLE_TEXT_LINE_OVERLAP_THRESHOLD = 0.5
+TABLE_TEXT_ORIENTATION_MIN_VALID_LINES = 3
+TABLE_TEXT_ORIENTATION_MIN_DOMINANCE_RATIO = 0.6
+TABLE_TEXT_ORIENTATION_ANGLES = frozenset({0, 90, 180, 270})
+_VLM_UNCLASSIFIED_TITLE_TYPE = "title"
+
+TITLE_BLOCK_TYPES = {
+    _VLM_UNCLASSIFIED_TITLE_TYPE,
+    BlockType.DOC_TITLE,
+    BlockType.PARAGRAPH_TITLE,
+}
+CODE_CONTENT_BLOCK_TYPES = {BlockType.CODE, RAW_ALGORITHM}
+LINE_METADATA_BLOCK_TYPES = {
+    BlockType.TEXT,
+    BlockType.REF_TEXT,
+    BlockType.DOC_TITLE,
+    BlockType.PARAGRAPH_TITLE,
+    RAW_CAPTION,
+    RAW_FOOTNOTE,
+}
+NATURAL_LANGUAGE_CONTENT_BLOCK_TYPES = {
+    BlockType.TEXT,
+    BlockType.DOC_TITLE,
+    BlockType.PARAGRAPH_TITLE,
+    BlockType.ASIDE_TEXT,
+    BlockType.HEADER,
+    BlockType.FOOTER,
+    BlockType.PAGE_NUMBER,
+    BlockType.PAGE_FOOTNOTE,
+    BlockType.REF_TEXT,
+    BlockType.LIST,
+    BlockType.INDEX,
+    RAW_CAPTION,
+    RAW_FOOTNOTE,
+}
+VLM_VISUAL_ANNOTATION_TYPE_MAP = {
+    BlockType.TABLE_CAPTION: RAW_CAPTION,
+    BlockType.IMAGE_CAPTION: RAW_CAPTION,
+    BlockType.CODE_CAPTION: RAW_CAPTION,
+    BlockType.TABLE_FOOTNOTE: RAW_FOOTNOTE,
+    BlockType.IMAGE_FOOTNOTE: RAW_FOOTNOTE,
+}
+VLM_MODEL_LIST_FIELDS = frozenset({"type", "bbox", "content", "angle", "score", "sub_type", "cell_merge"})
+MODEL_JSON_VISUAL_BLOCK_TYPES = {
+    BlockType.IMAGE,
+    BlockType.CHART,
+    BlockType.TABLE,
+    BlockType.EQUATION,
+}
+LOCAL_LAYOUT_IMAGE_BLOCK_BODY_TYPES = {BlockType.IMAGE, BlockType.CHART}
+LOCAL_LAYOUT_IMAGE_BLOCK_AREA_TYPES = {
+    *LOCAL_LAYOUT_IMAGE_BLOCK_BODY_TYPES,
+    BlockType.IMAGE_CAPTION,
+    RAW_CAPTION,
+    BlockType.IMAGE_FOOTNOTE,
+    RAW_FOOTNOTE,
+}
+VLM_LAYOUT_LABEL_MAP = {
+    "abstract": BlockType.TEXT,
+    "algorithm": BlockType.CODE,
+    "aside_text": BlockType.ASIDE_TEXT,
+    "chart": BlockType.CHART,
+    "content": BlockType.INDEX,
+    "display_formula": BlockType.EQUATION,
+    "doc_title": BlockType.DOC_TITLE,
+    "figure_title": RAW_CAPTION,
+    "footer": BlockType.FOOTER,
+    "footer_image": BlockType.FOOTER,
+    "footnote": BlockType.PAGE_FOOTNOTE,
+    "formula_number": RAW_FORMULA_NUMBER,
+    "header": BlockType.HEADER,
+    "header_image": BlockType.HEADER,
+    "image": BlockType.IMAGE,
+    "number": BlockType.PAGE_NUMBER,
+    "paragraph_title": BlockType.PARAGRAPH_TITLE,
+    "reference_content": BlockType.REF_TEXT,
+    "seal": BlockType.IMAGE,
+    "table": BlockType.TABLE,
+    "text": BlockType.TEXT,
+    "vertical_text": BlockType.TEXT,
+    "vision_footnote": RAW_FOOTNOTE,
+}
+PIPELINE_DET_TYPE = {
+    BlockType.TEXT,
+    BlockType.CODE,
+    BlockType.ASIDE_TEXT,
+    BlockType.INDEX,
+    BlockType.DOC_TITLE,
+    RAW_CAPTION,
+    BlockType.FOOTER,
+    BlockType.PAGE_FOOTNOTE,
+    BlockType.HEADER,
+    BlockType.PAGE_NUMBER,
+    BlockType.PARAGRAPH_TITLE,
+    BlockType.REF_TEXT,
+    RAW_FOOTNOTE,
+}
+NOT_EXTRACT_TYPES = {
+    BlockType.TEXT,
+    _VLM_UNCLASSIFIED_TITLE_TYPE,
+    BlockType.DOC_TITLE,
+    BlockType.PARAGRAPH_TITLE,
+    BlockType.HEADER,
+    BlockType.FOOTER,
+    BlockType.PAGE_NUMBER,
+    BlockType.PAGE_FOOTNOTE,
+    BlockType.REF_TEXT,
+    BlockType.TABLE_CAPTION,
+    BlockType.IMAGE_CAPTION,
+    RAW_CAPTION,
+    BlockType.TABLE_FOOTNOTE,
+    BlockType.IMAGE_FOOTNOTE,
+    RAW_FOOTNOTE,
+    BlockType.CODE_CAPTION,
+    RAW_PHONETIC,
+}
+VLM_TXT_DET_TYPE = NOT_EXTRACT_TYPES
+VLM_OCR_DET_TYPE = set(LINE_METADATA_BLOCK_TYPES)

docvortex/document/pdf/diagnostics.py

@@ -0,0 +1,350 @@
+from __future__ import annotations
+
+import json
+from io import BytesIO
+from typing import Any
+
+from loguru import logger
+from pypdf import PageObject, PdfReader, PdfWriter
+from reportlab.pdfgen import canvas
+
+from ...schema import BBox, RAW_PHONETIC, Block, BlockType, PageInfo
+
+# 文本类 block 共用 text bbox 样式，避免新增文本形态时遗漏多个绘制入口。
+TEXT_LIKE_BLOCK_TYPES_FOR_BBOX = {
+    BlockType.TEXT,
+    BlockType.REF_TEXT,
+    RAW_PHONETIC,
+}
+
+# layout.pdf 中这些 block 直接使用自身 bbox，复合 block 则使用子 block bbox。
+DIRECT_LAYOUT_BBOX_BLOCK_TYPES = TEXT_LIKE_BLOCK_TYPES_FOR_BBOX | {
+    BlockType.DOC_TITLE,
+    BlockType.PARAGRAPH_TITLE,
+    BlockType.EQUATION,
+    BlockType.LIST,
+    BlockType.INDEX,
+}
+
+
+def _get_layout_source_blocks(page: PageInfo) -> list[Block]:
+    """获取 layout.pdf 的页内原始布局块，避免段落合并后跨页子项串页绘制。"""
+    return page.preproc_blocks
+
+
+def cal_canvas_rect(page: Any, bbox: BBox) -> list[float]:
+    """
+    Calculate the rectangle coordinates on the canvas based on the original PDF page and bounding box.
+
+    Args:
+        page: A PyPDF2 Page object representing a single page in the PDF.
+        bbox: [x0, y0, x1, y1] representing the bounding box coordinates.
+
+    Returns:
+        rect: [x0, y0, width, height] representing the rectangle coordinates on the canvas.
+    """
+    page_width, page_height = float(page.cropbox[2]), float(page.cropbox[3])
+
+    actual_width = page_width  # The width of the final PDF display
+    actual_height = page_height  # The height of the final PDF display
+
+    rotation_obj = page.get("/Rotate", 0)
+    try:
+        rotation = int(rotation_obj) % 360  # cast rotation to int to handle IndirectObject
+    except (ValueError, TypeError) as e:
+        logger.warning(f"Invalid /Rotate value {rotation_obj!r} on page; defaulting to 0. Error: {e}")
+        rotation = 0
+
+    if rotation in [90, 270]:
+        # PDF is rotated 90 degrees or 270 degrees, and the width and height need to be swapped
+        actual_width, actual_height = actual_height, actual_width
+
+    x0, y0, x1, y1 = bbox
+    rect_w = abs(x1 - x0)
+    rect_h = abs(y1 - y0)
+
+    if rotation == 270:
+        rect_w, rect_h = rect_h, rect_w
+        x0 = actual_height - y1
+        y0 = actual_width - x1
+    elif rotation == 180:
+        x0 = page_width - x1
+        # y0 stays the same
+    elif rotation == 90:
+        rect_w, rect_h = rect_h, rect_w
+        x0, y0 = y0, x0
+    else:
+        # rotation == 0
+        y0 = page_height - y1
+
+    rect = [x0, y0, rect_w, rect_h]
+    return rect
+
+
+def draw_bbox_without_number(
+    i: int,
+    bbox_list: list[list[BBox]],
+    page: Any,
+    c: Any,
+    rgb_config: list[int],
+    fill_config: bool,
+) -> Any:
+    new_rgb = [float(color) / 255 for color in rgb_config]
+    page_data = bbox_list[i]
+
+    for bbox in page_data:
+        rect = cal_canvas_rect(page, bbox)  # Define the rectangle
+
+        if fill_config:  # filled rectangle
+            c.setFillColorRGB(new_rgb[0], new_rgb[1], new_rgb[2], 0.3)
+            c.rect(rect[0], rect[1], rect[2], rect[3], stroke=0, fill=1)
+        else:  # bounding box
+            c.setStrokeColorRGB(new_rgb[0], new_rgb[1], new_rgb[2])
+            c.rect(rect[0], rect[1], rect[2], rect[3], stroke=1, fill=0)
+    return c
+
+
+def draw_bbox_with_number(
+    i: int,
+    bbox_list: list[list[BBox]],
+    page: Any,
+    c: Any,
+    rgb_config: list[int],
+    fill_config: bool,
+    draw_bbox: bool = True,
+) -> Any:
+    new_rgb = [float(color) / 255 for color in rgb_config]
+    page_data = bbox_list[i]
+    # 强制转换为 float
+    # page_width, page_height = float(page.cropbox[2]), float(page.cropbox[3])
+
+    for j, bbox in enumerate(page_data):
+        # 确保bbox的每个元素都是float
+        rect = cal_canvas_rect(page, bbox)  # Define the rectangle
+
+        if draw_bbox:
+            if fill_config:
+                c.setFillColorRGB(*new_rgb, 0.3)
+                c.rect(rect[0], rect[1], rect[2], rect[3], stroke=0, fill=1)
+            else:
+                c.setStrokeColorRGB(*new_rgb)
+                c.rect(rect[0], rect[1], rect[2], rect[3], stroke=1, fill=0)
+        c.setFillColorRGB(*new_rgb, 1.0)
+        c.setFontSize(size=10)
+
+        c.saveState()
+        rotation_obj = page.get("/Rotate", 0)
+        try:
+            rotation = int(rotation_obj) % 360  # cast rotation to int to handle IndirectObject
+        except (ValueError, TypeError):
+            logger.warning(f"Invalid /Rotate value: {rotation_obj!r}, defaulting to 0")
+            rotation = 0
+
+        if rotation == 0:
+            c.translate(rect[0] + rect[2] + 2, rect[1] + rect[3] - 10)
+        elif rotation == 90:
+            c.translate(rect[0] + 10, rect[1] + rect[3] + 2)
+        elif rotation == 180:
+            c.translate(rect[0] - 2, rect[1] + 10)
+        elif rotation == 270:
+            c.translate(rect[0] + rect[2] - 10, rect[1] - 2)
+
+        c.rotate(rotation)
+        c.drawString(0, 0, str(j + 1))
+        c.restoreState()
+
+    return c
+
+
+def draw_layout_bbox(
+    pages: list[PageInfo],
+    pdf_bytes: bytes,
+    out_path: str,
+    filename: str,
+) -> None:
+    dropped_bbox_list = []
+    tables_body_list, tables_caption_list, tables_footnote_list = [], [], []
+    imgs_body_list, imgs_caption_list, imgs_footnote_list = [], [], []
+    codes_body_list, codes_caption_list, codes_footnote_list = [], [], []
+    titles_list = []
+    texts_list = []
+    equations_list = []
+    lists_list = []
+    list_items_list = []
+    indexs_list = []
+
+    for page in pages:
+        page_dropped_list = []
+        tables_body, tables_caption, tables_footnote = [], [], []
+        imgs_body, imgs_caption, imgs_footnote = [], [], []
+        codes_body, codes_caption, codes_footnote = [], [], []
+        titles = []
+        texts = []
+        equations = []
+        lists = []
+        list_items = []
+        indices = []
+
+        for dropped_bbox in page.discarded_blocks:
+            page_dropped_list.append(dropped_bbox.bbox)
+        dropped_bbox_list.append(page_dropped_list)
+        for block in _get_layout_source_blocks(page):
+            bbox = block.bbox
+            if block.type == BlockType.TABLE:
+                for nested_block in block.blocks:
+                    bbox = nested_block.bbox
+                    if nested_block.type == BlockType.TABLE_BODY:
+                        tables_body.append(bbox)
+                    elif nested_block.type == BlockType.TABLE_CAPTION:
+                        tables_caption.append(bbox)
+                    elif nested_block.type == BlockType.TABLE_FOOTNOTE:
+                        if nested_block._cross_page:
+                            continue
+                        tables_footnote.append(bbox)
+            elif block.type == BlockType.IMAGE:
+                for nested_block in block.blocks:
+                    bbox = nested_block.bbox
+                    if nested_block.type == BlockType.IMAGE_BODY:
+                        imgs_body.append(bbox)
+                    elif nested_block.type == BlockType.IMAGE_CAPTION:
+                        imgs_caption.append(bbox)
+                    elif nested_block.type == BlockType.IMAGE_FOOTNOTE:
+                        imgs_footnote.append(bbox)
+            elif block.type == BlockType.CODE:
+                for nested_block in block.blocks:
+                    if nested_block.type == BlockType.CODE_BODY:
+                        bbox = nested_block.bbox
+                        codes_body.append(bbox)
+                    elif nested_block.type == BlockType.CODE_CAPTION:
+                        bbox = nested_block.bbox
+                        codes_caption.append(bbox)
+                    elif nested_block.type == BlockType.CODE_FOOTNOTE:
+                        bbox = nested_block.bbox
+                        codes_footnote.append(bbox)
+            elif block.type == BlockType.CHART:
+                for nested_block in block.blocks:
+                    if nested_block.type == BlockType.CHART_BODY:
+                        bbox = nested_block.bbox
+                        imgs_body.append(bbox)
+                    elif nested_block.type == BlockType.CHART_CAPTION:
+                        bbox = nested_block.bbox
+                        imgs_caption.append(bbox)
+                    elif nested_block.type == BlockType.CHART_FOOTNOTE:
+                        bbox = nested_block.bbox
+                        imgs_footnote.append(bbox)
+            elif block.type in {BlockType.DOC_TITLE, BlockType.PARAGRAPH_TITLE}:
+                titles.append(bbox)
+            elif block.type in TEXT_LIKE_BLOCK_TYPES_FOR_BBOX:
+                texts.append(bbox)
+            elif block.type == BlockType.EQUATION:
+                equations.append(bbox)
+            elif block.type == BlockType.LIST:
+                lists.append(bbox)
+                if block.blocks:
+                    for sub_block in block.blocks:
+                        list_items.append(sub_block.bbox)
+            elif block.type == BlockType.INDEX:
+                indices.append(bbox)
+
+        tables_body_list.append(tables_body)
+        tables_caption_list.append(tables_caption)
+        tables_footnote_list.append(tables_footnote)
+        imgs_body_list.append(imgs_body)
+        imgs_caption_list.append(imgs_caption)
+        imgs_footnote_list.append(imgs_footnote)
+        titles_list.append(titles)
+        texts_list.append(texts)
+        equations_list.append(equations)
+        lists_list.append(lists)
+        list_items_list.append(list_items)
+        indexs_list.append(indices)
+        codes_body_list.append(codes_body)
+        codes_caption_list.append(codes_caption)
+        codes_footnote_list.append(codes_footnote)
+
+    layout_bbox_list = []
+
+    for page in pages:
+        page_block_list = []
+        for block in _get_layout_source_blocks(page):
+            if block.type in DIRECT_LAYOUT_BBOX_BLOCK_TYPES:
+                bbox = block.bbox
+                page_block_list.append(bbox)
+            elif block.type in [BlockType.IMAGE, BlockType.CHART, BlockType.CODE, BlockType.TABLE]:
+                for sub_block in block.blocks:
+                    if sub_block._cross_page:
+                        continue
+                    bbox = sub_block.bbox
+                    page_block_list.append(bbox)
+
+        layout_bbox_list.append(page_block_list)
+
+    pdf_bytes_io = BytesIO(pdf_bytes)
+    pdf_docs = PdfReader(pdf_bytes_io)
+    output_pdf = PdfWriter()
+
+    for i, page in enumerate(pdf_docs.pages):
+        # 获取原始页面尺寸
+        page_width, page_height = float(page.cropbox[2]), float(page.cropbox[3])
+        custom_page_size = (page_width, page_height)
+
+        packet = BytesIO()
+        # 使用原始PDF的尺寸创建canvas
+        c = canvas.Canvas(packet, pagesize=custom_page_size)
+
+        c = draw_bbox_without_number(i, codes_body_list, page, c, [102, 0, 204], True)
+        c = draw_bbox_without_number(i, codes_caption_list, page, c, [204, 153, 255], True)
+        c = draw_bbox_without_number(i, codes_footnote_list, page, c, [229, 204, 255], True)
+        c = draw_bbox_without_number(i, dropped_bbox_list, page, c, [158, 158, 158], True)
+        c = draw_bbox_without_number(i, tables_body_list, page, c, [204, 204, 0], True)
+        c = draw_bbox_without_number(i, tables_caption_list, page, c, [255, 255, 102], True)
+        c = draw_bbox_without_number(i, tables_footnote_list, page, c, [229, 255, 204], True)
+        c = draw_bbox_without_number(i, imgs_body_list, page, c, [153, 255, 51], True)
+        c = draw_bbox_without_number(i, imgs_caption_list, page, c, [102, 178, 255], True)
+        c = draw_bbox_without_number(i, imgs_footnote_list, page, c, [255, 178, 102], True)
+        c = draw_bbox_without_number(i, titles_list, page, c, [102, 102, 255], True)
+        c = draw_bbox_without_number(i, texts_list, page, c, [153, 0, 76], True)
+        c = draw_bbox_without_number(i, equations_list, page, c, [0, 255, 0], True)
+        c = draw_bbox_without_number(i, lists_list, page, c, [40, 169, 92], True)
+        c = draw_bbox_without_number(i, list_items_list, page, c, [40, 169, 92], False)
+        c = draw_bbox_without_number(i, indexs_list, page, c, [40, 169, 92], True)
+        c = draw_bbox_with_number(i, layout_bbox_list, page, c, [255, 0, 0], False, draw_bbox=False)
+
+        c.save()
+        packet.seek(0)
+        overlay_pdf = PdfReader(packet)
+
+        # 添加检查确保overlay_pdf.pages不为空
+        if len(overlay_pdf.pages) > 0:
+            new_page = PageObject(pdf=None)
+            new_page.update(page)
+            page = new_page
+            page.merge_page(overlay_pdf.pages[0])
+        else:
+            # 记录日志并继续处理下一个页面
+            # logger.warning(f"layout.pdf: 第{i + 1}页未能生成有效的overlay PDF")
+            pass
+
+        output_pdf.add_page(page)
+
+    # 保存结果
+    with open(f"{out_path}/{filename}", "wb") as f:
+        output_pdf.write(f)
+
+
+if __name__ == "__main__":
+    # 读取PDF文件
+    pdf_path = "examples/demo1.pdf"
+    with open(pdf_path, "rb") as f:
+        pdf_bytes = f.read()
+
+    # 从json文件读取pdf_info
+
+    json_path = "examples/demo1_1746005777.0863056_middle.json"
+    with open(json_path, "r", encoding="utf-8") as f:
+        pdf_ann = json.load(f)
+    pdf_info = pdf_ann["pdf_info"]
+    from .document import PDFDocument
+
+    doc = PDFDocument(pdf_bytes)
+    doc.draw_layout_bbox(pdf_info, "examples/output_with_layout.pdf")