From 37fbe998ac24b859c709bce4d19f33f97f9508fe Mon Sep 17 00:00:00 2001
From: myhloli <moe@myhloli.com>
Date: Wed, 18 Sep 2024 15:47:36 +0800
Subject: [PATCH] feat(ocr_mkcontent): support drop reason in none_with_reason
 modeEnable the `NONE_WITH_REASON` drop mode in `para_to_standard_format_v2`
 by updating the function signature to include the `drop_reason` parameter and
 handling it within the function logic. This enhancement allows the function
 to convey the reason for dropping content in the output.

---
 magic_pdf/dict2md/ocr_mkcontent.py  | 27 ++++++++++++++++++---------
 magic_pdf/libs/MakeContentConfig.py |  1 +
 2 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/magic_pdf/dict2md/ocr_mkcontent.py b/magic_pdf/dict2md/ocr_mkcontent.py
index 418c2aad..3b0091a1 100644
--- a/magic_pdf/dict2md/ocr_mkcontent.py
+++ b/magic_pdf/dict2md/ocr_mkcontent.py
@@ -265,30 +265,28 @@ def para_to_standard_format(para, img_buket_path):
     return para_content
 
 
-def para_to_standard_format_v2(para_block, img_buket_path, page_idx):
+def para_to_standard_format_v2(para_block, img_buket_path, page_idx, drop_reason=None):
     para_type = para_block['type']
+    para_content = {}
     if para_type == BlockType.Text:
         para_content = {
             'type': 'text',
             'text': merge_para_with_text(para_block),
-            'page_idx': page_idx,
         }
     elif para_type == BlockType.Title:
         para_content = {
             'type': 'text',
             'text': merge_para_with_text(para_block),
             'text_level': 1,
-            'page_idx': page_idx,
         }
     elif para_type == BlockType.InterlineEquation:
         para_content = {
             'type': 'equation',
             'text': merge_para_with_text(para_block),
             'text_format': 'latex',
-            'page_idx': page_idx,
         }
     elif para_type == BlockType.Image:
-        para_content = {'type': 'image', 'page_idx': page_idx}
+        para_content = {'type': 'image'}
         for block in para_block['blocks']:
             if block['type'] == BlockType.ImageBody:
                 para_content['img_path'] = join_path(
@@ -299,7 +297,7 @@ def para_to_standard_format_v2(para_block, img_buket_path, page_idx):
             if block['type'] == BlockType.ImageFootnote:
                 para_content['img_footnote'] = merge_para_with_text(block)
     elif para_type == BlockType.Table:
-        para_content = {'type': 'table', 'page_idx': page_idx}
+        para_content = {'type': 'table'}
         for block in para_block['blocks']:
             if block['type'] == BlockType.TableBody:
                 if block["lines"][0]["spans"][0].get('latex', ''):
@@ -312,6 +310,11 @@ def para_to_standard_format_v2(para_block, img_buket_path, page_idx):
             if block['type'] == BlockType.TableFootnote:
                 para_content['table_footnote'] = merge_para_with_text(block)
 
+    para_content['page_idx'] = page_idx
+
+    if drop_reason is not None:
+        para_content['drop_reason'] = drop_reason
+
     return para_content
 
 
@@ -397,8 +400,9 @@ def union_make(pdf_info_dict: list,
                img_buket_path: str = ''):
     output_content = []
     for page_info in pdf_info_dict:
+        drop_reason = page_info.get('drop_reason', None)
         if page_info.get('need_drop', False):
-            drop_reason = page_info.get('drop_reason')
+            # drop_reason = page_info.get('drop_reason')
             if drop_mode == DropMode.NONE:
                 pass
             elif drop_mode == DropMode.WHOLE_PDF:
@@ -425,8 +429,13 @@ def union_make(pdf_info_dict: list,
             output_content.extend(page_markdown)
         elif make_mode == MakeMode.STANDARD_FORMAT:
             for para_block in paras_of_layout:
-                para_content = para_to_standard_format_v2(
-                    para_block, img_buket_path, page_idx)
+                if drop_mode == DropMode.NONE_WITH_REASON:
+                    para_content = para_to_standard_format_v2(
+                        para_block, img_buket_path, page_idx, drop_reason)
+                else:
+                    para_content = para_to_standard_format_v2(
+                        para_block, img_buket_path, page_idx)
+
                 output_content.append(para_content)
     if make_mode in [MakeMode.MM_MD, MakeMode.NLP_MD]:
         return '\n\n'.join(output_content)
diff --git a/magic_pdf/libs/MakeContentConfig.py b/magic_pdf/libs/MakeContentConfig.py
index b1650aff..896edf62 100644
--- a/magic_pdf/libs/MakeContentConfig.py
+++ b/magic_pdf/libs/MakeContentConfig.py
@@ -8,3 +8,4 @@ class DropMode:
     WHOLE_PDF = "whole_pdf"
     SINGLE_PAGE = "single_page"
     NONE = "none"
+    NONE_WITH_REASON = "none_with_reason"