diff --git a/apps/common/handle/impl/pdf_split_handle.py b/apps/common/handle/impl/pdf_split_handle.py index b242292ff2..6b60abd090 100644 --- a/apps/common/handle/impl/pdf_split_handle.py +++ b/apps/common/handle/impl/pdf_split_handle.py @@ -42,46 +42,23 @@ def handle(self, file, pattern_list: List, with_filter: bool, limit: int, get_bu pdf_document = fitz.open(temp_file_path) try: - content = "" - for page_num in range(len(pdf_document)): - start_time = time.time() - page = pdf_document.load_page(page_num) - text = page.get_text() - - if text and text.strip(): # 如果页面中有文本内容 - page_content = text - else: - try: - new_doc = fitz.open() - new_doc.insert_pdf(pdf_document, from_page=page_num, to_page=page_num) - page_num_pdf = tempfile.gettempdir() + f"/{file.name}_{page_num}.pdf" - new_doc.save(page_num_pdf) - new_doc.close() - - loader = PyPDFLoader(page_num_pdf, extract_images=True) - page_content = "\n" + loader.load()[0].page_content - except NotImplementedError as e: - # 文件格式不支持，直接退出 - raise e - except BaseException as e: - # 当页出错继续进行下一页，防止一个页面出错导致整个文件解析失败 - max_kb.error(f"File: {file.name}, Page: {page_num + 1}, error: {e}") - continue - finally: - os.remove(page_num_pdf) - - content += page_content - - elapsed_time = time.time() - start_time - # todo 实现进度条代替下面的普通输出 - max_kb.debug( - f"File: {file.name}, Page: {page_num + 1}, Time : {elapsed_time: .3f}s, content-length: {len(page_content)}") + # 处理有目录的pdf + result = self.handle_toc(pdf_document) + if result is not None: + return { 'name': file.name, 'content': result } + + # 没目录但是有链接的pdf + result = self.handle_links(pdf_document) + if result is not None and len(result) > 0: + return { 'name': file.name, 'content': result } + + # 没有目录的pdf + content = self.handle_pdf_content(file, pdf_document) + 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: max_kb.error(f"File: {file.name}, error: {e}") return {'name': file.name, @@ -95,6 +72,130 @@ def handle(self, file, pattern_list: List, with_filter: bool, limit: int, get_bu 'content': split_model.parse(content) } + @staticmethod + def handle_pdf_content(file, pdf_document): + content = "" + for page_num in range(len(pdf_document)): + start_time = time.time() + page = pdf_document.load_page(page_num) + text = page.get_text() + + if text and text.strip(): # 如果页面中有文本内容 + page_content = text + else: + try: + new_doc = fitz.open() + new_doc.insert_pdf(pdf_document, from_page=page_num, to_page=page_num) + page_num_pdf = tempfile.gettempdir() + f"/{file.name}_{page_num}.pdf" + new_doc.save(page_num_pdf) + new_doc.close() + + loader = PyPDFLoader(page_num_pdf, extract_images=True) + page_content = "\n" + loader.load()[0].page_content + except NotImplementedError as e: + # 文件格式不支持，直接退出 + raise e + except BaseException as e: + # 当页出错继续进行下一页，防止一个页面出错导致整个文件解析失败 + max_kb.error(f"File: {file.name}, Page: {page_num + 1}, error: {e}") + continue + finally: + os.remove(page_num_pdf) + + content += page_content + + elapsed_time = time.time() - start_time + max_kb.debug(f"File: {file.name}, Page: {page_num + 1}, Time : {elapsed_time: .3f}s, content-length: {len(page_content)}") + + return content + + @staticmethod + def handle_toc(doc): + # 找到目录 + toc = doc.get_toc() + if toc is None or len(toc) == 0: + return None + + # 创建存储章节内容的数组 + chapters = [] + + # 遍历目录并按章节提取文本 + for i, entry in enumerate(toc): + level, title, start_page = entry + start_page -= 1 # PyMuPDF 页码从 0 开始，书签页码从 1 开始 + + # 确定结束页码，如果是最后一个章节则到文档末尾 + if i + 1 < len(toc): + end_page = toc[i + 1][2] - 1 + else: + end_page = doc.page_count - 1 + + # 提取该章节的文本内容 + chapter_text = "" + for page_num in range(start_page, end_page + 1): + page = doc.load_page(page_num) # 加载页面 + chapter_text += page.get_text("text") # 提取文本 + + # 保存章节内容和章节标题 + chapters.append({"title": title, "content": chapter_text}) + return chapters + + @staticmethod + def handle_links(doc): + # 创建存储章节内容的数组 + chapters = [] + + # 遍历 PDF 的每一页，查找带有目录链接的页 + for page_num in range(doc.page_count): + page = doc.load_page(page_num) + links = page.get_links() + + # 检查该页是否包含内部链接（即指向文档内部的页面） + for num in range(len(links)): + link = links[num] + if link['kind'] == 1: # 'kind' 为 1 表示内部链接 + # 获取链接目标的页面 + dest_page = link['page'] + rect = link['from'] # 获取链接的矩形区域 + + # 提取链接区域的文本作为标题 + link_title = page.get_text("text", clip=rect).strip().split("\n")[0] + print(f'link_title: {link_title}') + # 提取目标页面内容作为章节开始 + start_page = dest_page + end_page = dest_page + # 下一个link + next_link = links[num + 1] if num + 1 < len(links) else None + next_link_title = None + if next_link is not None and next_link['kind'] == 1: + rect = next_link['from'] + next_link_title = page.get_text("text", clip=rect).strip().split("\n")[0] + print(f'next_link_title: {next_link_title}') + end_page = next_link['page'] + + # 提取章节内容 + chapter_text = "" + for p_num in range(start_page, end_page + 1): + p = doc.load_page(p_num) + text = p.get_text("text") + + if link_title in text: + text = text[text.find(link_title) + len(link_title):] + + if next_link_title is not None: + idx = text.replace('\n', '').find(next_link_title) + if idx > -1: + text = text[:idx] + chapter_text += text + + # 保存章节信息 + chapters.append({ + "title": link_title, + "content": chapter_text + }) + + return chapters + def support(self, file, get_buffer): file_name: str = file.name.lower() if file_name.endswith(".pdf") or file_name.endswith(".PDF"):