From 4d9355dec7ebf9cbe78bccacdd00e7bd58422263 Mon Sep 17 00:00:00 2001
From: CaptainB <bin@fit2cloud.com>
Date: Fri, 16 Aug 2024 12:57:06 +0800
Subject: [PATCH] =?UTF-8?q?refactor:=20=E9=80=90=E9=A1=B5=E5=8A=A0?=
 =?UTF-8?q?=E8=BD=BDpdf,=20=E5=9B=BE=E7=89=87=E7=B1=BB=E5=9E=8B=E5=8D=95?=
 =?UTF-8?q?=E7=8B=AC=E4=BF=9D=E5=AD=98=E6=88=90=E6=96=87=E4=BB=B6=E5=8A=A0?=
 =?UTF-8?q?=E8=BD=BD?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 apps/common/handle/impl/pdf_split_handle.py | 54 +++++++++++----------
 1 file changed, 28 insertions(+), 26 deletions(-)

diff --git a/apps/common/handle/impl/pdf_split_handle.py b/apps/common/handle/impl/pdf_split_handle.py
index e4bac8d08ca..ecf4880bd59 100644
--- a/apps/common/handle/impl/pdf_split_handle.py
+++ b/apps/common/handle/impl/pdf_split_handle.py
@@ -9,7 +9,7 @@
 import re
 from typing import List
 
-import pypdf
+from pypdf import PdfReader, PdfWriter
 import os
 import tempfile
 from langchain_community.document_loaders import PyPDFLoader
@@ -17,6 +17,8 @@
 from common.handle.base_split_handle import BaseSplitHandle
 from common.util.split_model import SplitModel
 
+import time
+
 default_pattern_list = [re.compile('(?<=^)# .*|(?<=\\n)# .*'),
                         re.compile('(?<=\\n)(?<!#)## (?!#).*|(?<=^)(?<!#)## (?!#).*'),
                         re.compile("(?<=\\n)(?<!#)### (?!#).*|(?<=^)(?<!#)### (?!#).*"),
@@ -26,26 +28,6 @@
                         re.compile("(?<!\n)\n\n+")]
 
 
-def check_pdf_is_image(pdf_path):
-    try:
-        # 打开PDF文件
-        with open(pdf_path, "rb") as f:
-            reader = pypdf.PdfReader(f)
-            for page_num in range(len(reader.pages)):
-                page = reader.pages[page_num]
-
-                # 尝试提取文本
-                text = page.extract_text()
-                if text and text.strip():  # 如果页面中有文本内容
-                    return False  # 不是纯图片
-                else:
-                    return True  # 可能是图片或扫描件
-
-    except Exception as e:
-        print(f"Error: {e}")
-        return None
-
-
 class PdfSplitHandle(BaseSplitHandle):
     def handle(self, file, pattern_list: List, with_filter: bool, limit: int, get_buffer, save_image):
         with tempfile.NamedTemporaryFile(delete=False) as temp_file:
@@ -56,17 +38,37 @@ def handle(self, file, pattern_list: List, with_filter: bool, limit: int, get_bu
             temp_file_path = temp_file.name
 
         try:
-            if check_pdf_is_image(temp_file_path):
-                loader = PyPDFLoader(temp_file_path, extract_images=True)
-            else:
-                loader = PyPDFLoader(temp_file_path, extract_images=False)
+            content = ""
+            reader = PdfReader(temp_file_path)
+            for page_num in range(len(reader.pages)):
+                start_time = time.time()
+                page = reader.pages[page_num]
+                text = page.extract_text()
+
+                if text and text.strip():  # 如果页面中有文本内容
+                    page_content = text
+                else:
+                    try:
+                        writer = PdfWriter()
+                        writer.add_page(page)
+                        with tempfile.NamedTemporaryFile(delete=False) as output_pdf:
+                            writer.write(output_pdf)
+                        loader = PyPDFLoader(output_pdf.name, extract_images=True)
+                        page_content = "\n".join([page.page_content for page in loader.lazy_load()])
+                    finally:
+                        os.remove(output_pdf.name)
+
+                content += page_content
 
-            content = "\n".join([page.page_content for page in loader.lazy_load()])
+                elapsed_time = time.time() - start_time
+                print(f"Time : {elapsed_time: .3f}s,  Page: {page_num + 1},  content-length: {len(page_content)}")
 
             if pattern_list is not None and len(pattern_list) > 0:
                 split_model = SplitModel(pattern_list, with_filter, limit)
             else:
                 split_model = SplitModel(default_pattern_list, with_filter=with_filter, limit=limit)
+
+
         except BaseException as e:
             return {'name': file.name,
                     'content': []}