magic-pdf 0.7.1__py3-none-any.whl → 0.8.1__py3-none-any.whl

magic_pdf/libs/draw_bbox.py

@@ -1,6 +1,6 @@
-from magic_pdf.libs.Constants import CROSS_PAGE
 from magic_pdf.libs.commons import fitz  # PyMuPDF
-from magic_pdf.libs.ocr_content_type import ContentType, BlockType, CategoryId
+from magic_pdf.libs.Constants import CROSS_PAGE
+from magic_pdf.libs.ocr_content_type import BlockType, CategoryId, ContentType
 from magic_pdf.model.magic_model import MagicModel
 
 
@@ -65,11 +65,13 @@ def draw_bbox_with_number(i, bbox_list, page, rgb_config, fill_config):
         )  # Insert the index in the top left corner of the rectangle
 
 
-def draw_layout_bbox(pdf_info, pdf_bytes, out_path):
+def draw_layout_bbox(pdf_info, pdf_bytes, out_path, filename):
     layout_bbox_list = []
     dropped_bbox_list = []
-    tables_list, tables_body_list, tables_caption_list, tables_footnote_list = [], [], [], []
+    tables_list, tables_body_list = [], []
+    tables_caption_list, tables_footnote_list = [], []
     imgs_list, imgs_body_list, imgs_caption_list = [], [], []
+    imgs_footnote_list = []
     titles_list = []
     texts_list = []
     interequations_list = []
@@ -77,41 +79,43 @@ def draw_layout_bbox(pdf_info, pdf_bytes, out_path):
         page_layout_list = []
         page_dropped_list = []
         tables, tables_body, tables_caption, tables_footnote = [], [], [], []
-        imgs, imgs_body, imgs_caption = [], [], []
+        imgs, imgs_body, imgs_caption, imgs_footnote = [], [], [], []
         titles = []
         texts = []
         interequations = []
-        for layout in page["layout_bboxes"]:
-            page_layout_list.append(layout["layout_bbox"])
+        for layout in page['layout_bboxes']:
+            page_layout_list.append(layout['layout_bbox'])
         layout_bbox_list.append(page_layout_list)
-        for dropped_bbox in page["discarded_blocks"]:
-            page_dropped_list.append(dropped_bbox["bbox"])
+        for dropped_bbox in page['discarded_blocks']:
+            page_dropped_list.append(dropped_bbox['bbox'])
         dropped_bbox_list.append(page_dropped_list)
-        for block in page["para_blocks"]:
-            bbox = block["bbox"]
-            if block["type"] == BlockType.Table:
+        for block in page['para_blocks']:
+            bbox = block['bbox']
+            if block['type'] == BlockType.Table:
                 tables.append(bbox)
-                for nested_block in block["blocks"]:
-                    bbox = nested_block["bbox"]
-                    if nested_block["type"] == BlockType.TableBody:
+                for nested_block in block['blocks']:
+                    bbox = nested_block['bbox']
+                    if nested_block['type'] == BlockType.TableBody:
                         tables_body.append(bbox)
-                    elif nested_block["type"] == BlockType.TableCaption:
+                    elif nested_block['type'] == BlockType.TableCaption:
                         tables_caption.append(bbox)
-                    elif nested_block["type"] == BlockType.TableFootnote:
+                    elif nested_block['type'] == BlockType.TableFootnote:
                         tables_footnote.append(bbox)
-            elif block["type"] == BlockType.Image:
+            elif block['type'] == BlockType.Image:
                 imgs.append(bbox)
-                for nested_block in block["blocks"]:
-                    bbox = nested_block["bbox"]
-                    if nested_block["type"] == BlockType.ImageBody:
+                for nested_block in block['blocks']:
+                    bbox = nested_block['bbox']
+                    if nested_block['type'] == BlockType.ImageBody:
                         imgs_body.append(bbox)
-                    elif nested_block["type"] == BlockType.ImageCaption:
+                    elif nested_block['type'] == BlockType.ImageCaption:
                         imgs_caption.append(bbox)
-            elif block["type"] == BlockType.Title:
+                    elif nested_block['type'] == BlockType.ImageFootnote:
+                        imgs_footnote.append(bbox)
+            elif block['type'] == BlockType.Title:
                 titles.append(bbox)
-            elif block["type"] == BlockType.Text:
+            elif block['type'] == BlockType.Text:
                 texts.append(bbox)
-            elif block["type"] == BlockType.InterlineEquation:
+            elif block['type'] == BlockType.InterlineEquation:
                 interequations.append(bbox)
         tables_list.append(tables)
         tables_body_list.append(tables_body)
@@ -120,30 +124,40 @@ def draw_layout_bbox(pdf_info, pdf_bytes, out_path):
         imgs_list.append(imgs)
         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)
         interequations_list.append(interequations)
 
-    pdf_docs = fitz.open("pdf", pdf_bytes)
+    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)
-        draw_bbox_without_number(i, dropped_bbox_list, page, [158, 158, 158], True)
-        draw_bbox_without_number(i, tables_list, page, [153, 153, 0], True)  # color !
-        draw_bbox_without_number(i, tables_body_list, page, [204, 204, 0], True)
-        draw_bbox_without_number(i, tables_caption_list, page, [255, 255, 102], True)
-        draw_bbox_without_number(i, tables_footnote_list, page, [229, 255, 204], True)
+        draw_bbox_without_number(i, dropped_bbox_list, page, [158, 158, 158],
+                                 True)
+        draw_bbox_without_number(i, tables_list, page, [153, 153, 0],
+                                 True)  # color !
+        draw_bbox_without_number(i, tables_body_list, page, [204, 204, 0],
+                                 True)
+        draw_bbox_without_number(i, tables_caption_list, page, [255, 255, 102],
+                                 True)
+        draw_bbox_without_number(i, tables_footnote_list, page,
+                                 [229, 255, 204], True)
         draw_bbox_without_number(i, imgs_list, page, [51, 102, 0], True)
         draw_bbox_without_number(i, imgs_body_list, page, [153, 255, 51], True)
-        draw_bbox_without_number(i, imgs_caption_list, page, [102, 178, 255], True)
+        draw_bbox_without_number(i, imgs_caption_list, page, [102, 178, 255],
+                                 True)
+        draw_bbox_with_number(i, imgs_footnote_list, page, [255, 178, 102],
+                              True),
         draw_bbox_without_number(i, titles_list, page, [102, 102, 255], True)
         draw_bbox_without_number(i, texts_list, page, [153, 0, 76], True)
-        draw_bbox_without_number(i, interequations_list, page, [0, 255, 0], True)
+        draw_bbox_without_number(i, interequations_list, page, [0, 255, 0],
+                                 True)
 
     # Save the PDF
-    pdf_docs.save(f"{out_path}/layout.pdf")
+    pdf_docs.save(f'{out_path}/{filename}_layout.pdf')
 
 
-def draw_span_bbox(pdf_info, pdf_bytes, out_path):
+def draw_span_bbox(pdf_info, pdf_bytes, out_path, filename):
     text_list = []
     inline_equation_list = []
     interline_equation_list = []
@@ -154,22 +168,22 @@ def draw_span_bbox(pdf_info, pdf_bytes, out_path):
     next_page_inline_equation_list = []
 
     def get_span_info(span):
-        if span["type"] == ContentType.Text:
+        if span['type'] == ContentType.Text:
             if span.get(CROSS_PAGE, False):
-                next_page_text_list.append(span["bbox"])
+                next_page_text_list.append(span['bbox'])
             else:
-                page_text_list.append(span["bbox"])
-        elif span["type"] == ContentType.InlineEquation:
+                page_text_list.append(span['bbox'])
+        elif span['type'] == ContentType.InlineEquation:
             if span.get(CROSS_PAGE, False):
-                next_page_inline_equation_list.append(span["bbox"])
+                next_page_inline_equation_list.append(span['bbox'])
             else:
-                page_inline_equation_list.append(span["bbox"])
-        elif span["type"] == ContentType.InterlineEquation:
-            page_interline_equation_list.append(span["bbox"])
-        elif span["type"] == ContentType.Image:
-            page_image_list.append(span["bbox"])
-        elif span["type"] == ContentType.Table:
-            page_table_list.append(span["bbox"])
+                page_inline_equation_list.append(span['bbox'])
+        elif span['type'] == ContentType.InterlineEquation:
+            page_interline_equation_list.append(span['bbox'])
+        elif span['type'] == ContentType.Image:
+            page_image_list.append(span['bbox'])
+        elif span['type'] == ContentType.Table:
+            page_table_list.append(span['bbox'])
 
     for page in pdf_info:
         page_text_list = []
@@ -188,84 +202,89 @@ def draw_span_bbox(pdf_info, pdf_bytes, out_path):
             next_page_inline_equation_list.clear()
 
         # 构造dropped_list
-        for block in page["discarded_blocks"]:
-            if block["type"] == BlockType.Discarded:
-                for line in block["lines"]:
-                    for span in line["spans"]:
-                        page_dropped_list.append(span["bbox"])
+        for block in page['discarded_blocks']:
+            if block['type'] == BlockType.Discarded:
+                for line in block['lines']:
+                    for span in line['spans']:
+                        page_dropped_list.append(span['bbox'])
         dropped_list.append(page_dropped_list)
         # 构造其余useful_list
-        for block in page["para_blocks"]:
-            if block["type"] in [
-                BlockType.Text,
-                BlockType.Title,
-                BlockType.InterlineEquation,
+        for block in page['para_blocks']:
+            if block['type'] in [
+                    BlockType.Text,
+                    BlockType.Title,
+                    BlockType.InterlineEquation,
             ]:
-                for line in block["lines"]:
-                    for span in line["spans"]:
+                for line in block['lines']:
+                    for span in line['spans']:
                         get_span_info(span)
-            elif block["type"] in [BlockType.Image, BlockType.Table]:
-                for sub_block in block["blocks"]:
-                    for line in sub_block["lines"]:
-                        for span in line["spans"]:
+            elif block['type'] in [BlockType.Image, BlockType.Table]:
+                for sub_block in block['blocks']:
+                    for line in sub_block['lines']:
+                        for span in line['spans']:
                             get_span_info(span)
         text_list.append(page_text_list)
         inline_equation_list.append(page_inline_equation_list)
         interline_equation_list.append(page_interline_equation_list)
         image_list.append(page_image_list)
         table_list.append(page_table_list)
-    pdf_docs = fitz.open("pdf", pdf_bytes)
+    pdf_docs = fitz.open('pdf', pdf_bytes)
     for i, page in enumerate(pdf_docs):
         # 获取当前页面的数据
         draw_bbox_without_number(i, text_list, page, [255, 0, 0], False)
-        draw_bbox_without_number(i, inline_equation_list, page, [0, 255, 0], False)
-        draw_bbox_without_number(i, interline_equation_list, page, [0, 0, 255], False)
+        draw_bbox_without_number(i, inline_equation_list, page, [0, 255, 0],
+                                 False)
+        draw_bbox_without_number(i, interline_equation_list, page, [0, 0, 255],
+                                 False)
         draw_bbox_without_number(i, image_list, page, [255, 204, 0], False)
         draw_bbox_without_number(i, table_list, page, [204, 0, 255], False)
         draw_bbox_without_number(i, dropped_list, page, [158, 158, 158], False)
 
     # Save the PDF
-    pdf_docs.save(f"{out_path}/spans.pdf")
+    pdf_docs.save(f'{out_path}/{filename}_spans.pdf')
 
 
-def drow_model_bbox(model_list: list, pdf_bytes, out_path):
+def drow_model_bbox(model_list: list, pdf_bytes, out_path, filename):
     dropped_bbox_list = []
     tables_body_list, tables_caption_list, tables_footnote_list = [], [], []
-    imgs_body_list, imgs_caption_list = [], []
+    imgs_body_list, imgs_caption_list, imgs_footnote_list = [], [], []
     titles_list = []
     texts_list = []
     interequations_list = []
-    pdf_docs = fitz.open("pdf", pdf_bytes)
+    pdf_docs = fitz.open('pdf', pdf_bytes)
     magic_model = MagicModel(model_list, pdf_docs)
     for i in range(len(model_list)):
         page_dropped_list = []
         tables_body, tables_caption, tables_footnote = [], [], []
-        imgs_body, imgs_caption = [], []
+        imgs_body, imgs_caption, imgs_footnote = [], [], []
         titles = []
         texts = []
         interequations = []
         page_info = magic_model.get_model_list(i)
-        layout_dets = page_info["layout_dets"]
+        layout_dets = page_info['layout_dets']
         for layout_det in layout_dets:
-            bbox = layout_det["bbox"]
-            if layout_det["category_id"] == CategoryId.Text:
+            bbox = layout_det['bbox']
+            if layout_det['category_id'] == CategoryId.Text:
                 texts.append(bbox)
-            elif layout_det["category_id"] == CategoryId.Title:
+            elif layout_det['category_id'] == CategoryId.Title:
                 titles.append(bbox)
-            elif layout_det["category_id"] == CategoryId.TableBody:
+            elif layout_det['category_id'] == CategoryId.TableBody:
                 tables_body.append(bbox)
-            elif layout_det["category_id"] == CategoryId.TableCaption:
+            elif layout_det['category_id'] == CategoryId.TableCaption:
                 tables_caption.append(bbox)
-            elif layout_det["category_id"] == CategoryId.TableFootnote:
+            elif layout_det['category_id'] == CategoryId.TableFootnote:
                 tables_footnote.append(bbox)
-            elif layout_det["category_id"] == CategoryId.ImageBody:
+            elif layout_det['category_id'] == CategoryId.ImageBody:
                 imgs_body.append(bbox)
-            elif layout_det["category_id"] == CategoryId.ImageCaption:
+            elif layout_det['category_id'] == CategoryId.ImageCaption:
                 imgs_caption.append(bbox)
-            elif layout_det["category_id"] == CategoryId.InterlineEquation_YOLO:
+            elif layout_det[
+                    'category_id'] == CategoryId.InterlineEquation_YOLO:
                 interequations.append(bbox)
-            elif layout_det["category_id"] == CategoryId.Abandon:
+            elif layout_det['category_id'] == CategoryId.Abandon:
                 page_dropped_list.append(bbox)
+            elif layout_det['category_id'] == CategoryId.ImageFootnote:
+                imgs_footnote.append(bbox)
 
         tables_body_list.append(tables_body)
         tables_caption_list.append(tables_caption)
@@ -276,17 +295,24 @@ def drow_model_bbox(model_list: list, pdf_bytes, out_path):
         texts_list.append(texts)
         interequations_list.append(interequations)
         dropped_bbox_list.append(page_dropped_list)
+        imgs_footnote_list.append(imgs_footnote)
 
     for i, page in enumerate(pdf_docs):
-        draw_bbox_with_number(i, dropped_bbox_list, page, [158, 158, 158], True) # color !
+        draw_bbox_with_number(i, dropped_bbox_list, page, [158, 158, 158],
+                              True)  # color !
         draw_bbox_with_number(i, tables_body_list, page, [204, 204, 0], True)
-        draw_bbox_with_number(i, tables_caption_list, page, [255, 255, 102], True)
-        draw_bbox_with_number(i, tables_footnote_list, page, [229, 255, 204], True)
+        draw_bbox_with_number(i, tables_caption_list, page, [255, 255, 102],
+                              True)
+        draw_bbox_with_number(i, tables_footnote_list, page, [229, 255, 204],
+                              True)
         draw_bbox_with_number(i, imgs_body_list, page, [153, 255, 51], True)
-        draw_bbox_with_number(i, imgs_caption_list, page, [102, 178, 255], True)
+        draw_bbox_with_number(i, imgs_caption_list, page, [102, 178, 255],
+                              True)
+        draw_bbox_with_number(i, imgs_footnote_list, page, [255, 178, 102],
+                              True)
         draw_bbox_with_number(i, titles_list, page, [102, 102, 255], True)
         draw_bbox_with_number(i, texts_list, page, [153, 0, 76], True)
         draw_bbox_with_number(i, interequations_list, page, [0, 255, 0], True)
 
     # Save the PDF
-    pdf_docs.save(f"{out_path}/model.pdf")
+    pdf_docs.save(f'{out_path}/{filename}_model.pdf')

magic_pdf/libs/ocr_content_type.py

@@ -1,23 +1,25 @@
 class ContentType:
-    Image = "image"
-    Table = "table"
-    Text = "text"
-    InlineEquation = "inline_equation"
-    InterlineEquation = "interline_equation"
-    
+    Image = 'image'
+    Table = 'table'
+    Text = 'text'
+    InlineEquation = 'inline_equation'
+    InterlineEquation = 'interline_equation'
+
+
 class BlockType:
-    Image = "image"
-    ImageBody = "image_body"
-    ImageCaption = "image_caption"
-    Table = "table"
-    TableBody = "table_body"
-    TableCaption = "table_caption"
-    TableFootnote = "table_footnote"
-    Text = "text"
-    Title = "title"
-    InterlineEquation = "interline_equation"
-    Footnote = "footnote"
-    Discarded = "discarded"
+    Image = 'image'
+    ImageBody = 'image_body'
+    ImageCaption = 'image_caption'
+    ImageFootnote = 'image_footnote'
+    Table = 'table'
+    TableBody = 'table_body'
+    TableCaption = 'table_caption'
+    TableFootnote = 'table_footnote'
+    Text = 'text'
+    Title = 'title'
+    InterlineEquation = 'interline_equation'
+    Footnote = 'footnote'
+    Discarded = 'discarded'
 
 
 class CategoryId:
@@ -33,3 +35,4 @@ class CategoryId:
     InlineEquation = 13
     InterlineEquation_YOLO = 14
     OcrText = 15
+    ImageFootnote = 101

magic_pdf/libs/version.py

@@ -1 +1 @@
-__version__ = "0.7.1"
+__version__ = "0.8.1"

magic_pdf/model/doc_analyze_by_custom_model.py

@@ -103,20 +103,32 @@ def custom_model_init(ocr: bool = False, show_log: bool = False):
     return custom_model
 
 
-def doc_analyze(pdf_bytes: bytes, ocr: bool = False, show_log: bool = False):
+def doc_analyze(pdf_bytes: bytes, ocr: bool = False, show_log: bool = False,
+                start_page_id=0, end_page_id=None):
 
     model_manager = ModelSingleton()
     custom_model = model_manager.get_model(ocr, show_log)
 
     images = load_images_from_pdf(pdf_bytes)
 
+    # end_page_id = end_page_id if end_page_id else len(images) - 1
+    end_page_id = end_page_id if end_page_id is not None and end_page_id >= 0 else len(images) - 1
+
+    if end_page_id > len(images) - 1:
+        logger.warning("end_page_id is out of range, use images length")
+        end_page_id = len(images) - 1
+
     model_json = []
     doc_analyze_start = time.time()
+
     for index, img_dict in enumerate(images):
         img = img_dict["img"]
         page_width = img_dict["width"]
         page_height = img_dict["height"]
-        result = custom_model(img)
+        if start_page_id <= index <= end_page_id:
+            result = custom_model(img)
+        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)