切图逻辑修复

magic_pdf/libs/pdf_image_tools.py

@@ -1,4 +1,3 @@
-from loguru import logger
 
 from magic_pdf.io.AbsReaderWriter import AbsReaderWriter
 from magic_pdf.libs.commons import fitz
@@ -20,10 +19,6 @@ def cut_image(bbox: tuple, page_num: int, page: fitz.Page, return_path, imageWri
     # 新版本生成平铺路径
     img_hash256_path = f"{compute_sha256(img_path)}.jpg"
 
-    if any([bbox[0] >= bbox[2], bbox[1] >= bbox[3]]):
-        logger.warning(f"image_bboxes: 错误的box, {bbox}")
-        return img_hash256_path
-
     # 将坐标转换为fitz.Rect对象
     rect = fitz.Rect(*bbox)
     # 配置缩放倍数为3倍

magic_pdf/pre_proc/cut_image.py

@@ -1,3 +1,5 @@
+from loguru import logger
+
 from magic_pdf.libs.commons import join_path
 from magic_pdf.libs.ocr_content_type import ContentType
 from magic_pdf.libs.pdf_image_tools import cut_image
@@ -10,9 +12,13 @@ def ocr_cut_image_and_table(spans, page, page_id, pdf_bytes_md5, imageWriter):
     for span in spans:
         span_type = span['type']
         if span_type == ContentType.Image:
+            if not check_img_bbox(span['bbox']):
+                continue
             span['image_path'] = cut_image(span['bbox'], page_id, page, return_path=return_path('images'),
                                            imageWriter=imageWriter)
         elif span_type == ContentType.Table:
+            if not check_img_bbox(span['bbox']):
+                continue
             span['image_path'] = cut_image(span['bbox'], page_id, page, return_path=return_path('tables'),
                                            imageWriter=imageWriter)
 
@@ -38,15 +44,28 @@ def txt_save_images_by_bboxes(page_num: int, page, pdf_bytes_md5: str,
         return join_path(pdf_bytes_md5, type)
 
     for bbox in image_bboxes:
+        if not check_img_bbox(bbox):
+            continue
         image_path = cut_image(bbox, page_num, page, return_path("images"), imageWriter)
         image_info.append({"bbox": bbox, "image_path": image_path})
 
     for bbox in images_overlap_backup:
+        if not check_img_bbox(bbox):
+            continue
         image_path = cut_image(bbox, page_num, page, return_path("images"), imageWriter)
         image_backup_info.append({"bbox": bbox, "image_path": image_path})
 
     for bbox in table_bboxes:
+        if not check_img_bbox(bbox):
+            continue
         image_path = cut_image(bbox, page_num, page, return_path("tables"), imageWriter)
         table_info.append({"bbox": bbox, "image_path": image_path})
 
     return image_info, image_backup_info, table_info, inline_eq_info, interline_eq_info
+
+
+def check_img_bbox(bbox) -> bool:
+    if any([bbox[0] >= bbox[2], bbox[1] >= bbox[3]]):
+        logger.warning(f"image_bboxes: 错误的box, {bbox}")
+        return False
+    return True