PDF_Processor.py 6.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160
  1. from .... import logger,fitz,BytesIO,sys
  2. from .nlp_Processor import SpacyNLP,NLTKNLP
  3. class PDFProcessor:
  4. """
  5. PDF文件处理器,用于解析PDF文件并处理每一页的内容。
  6. Attributes:
  7. file_content (BytesIO): PDF文件的内容。
  8. document_id (str): 文档的唯一标识符。
  9. max_tokens (int): 处理文本时的最大令牌数限制。
  10. bucket_name (str): 用于存储处理结果的存储桶名称。
  11. KbmService: KbmService对象,用于提供文本处理和存储服务。
  12. Methods:
  13. process_pdf(): 解析PDF文件并处理每一页的内容,包括文本和图片。
  14. """
  15. def __init__(self, nlp_type=1):
  16. """
  17. 初始化PDF处理器,根据指定的NLP类型选择合适的NLP处理器。
  18. Args:
  19. nlp_type (int, optional): 指定NLP处理器的类型。1为SpacyNLP,2为NLTKNLP。默认为1。
  20. """
  21. # self.nlp = SpacyNLP() if nlp_type == 1 else NLTKNLP() if nlp_type == 2 else None
  22. self.nlp = SpacyNLP()
  23. def get_results(self, file_content, document_id, max_tokens, bucket_name, KbmService):
  24. """
  25. 解析PDF文件并处理每一页的内容,包括文本和图片。
  26. Args:
  27. file_content (BytesIO): PDF文件的内容。
  28. document_id (str): 文档的ID。
  29. max_tokens (int): 最大令牌数。
  30. bucket_name (str): 存储桶的名称。
  31. KbmService: KbmService对象。
  32. Returns:
  33. tuple: 包含处理后的句子列表和页面图片列表的元组。
  34. """
  35. self.file_content = file_content
  36. self.document_id = document_id
  37. self.max_tokens = max_tokens
  38. self.bucket_name = bucket_name
  39. self.KbmService = KbmService
  40. # 初始化日志信息,表示开始解析PDF
  41. sentences = None
  42. logger.info("开始解析pdf")
  43. processed_pages, current_pages = self._process_pdf() # 解析PDF文件并处理每一页的内容
  44. try:
  45. sentences = self._split_sentences(processed_pages,max_tokens) # 将处理后的页面列表分割成句子列表
  46. except Exception as e:
  47. logger.error(f"错误{str(e)}")
  48. logger.info('PDFProcessor处理完成')
  49. return sentences, current_pages
  50. def _process_pdf(self):
  51. """
  52. 解析PDF文件并处理每一页的内容。
  53. Returns:
  54. list: 处理后的页面列表。
  55. """
  56. processed_pages = [] # 初始化处理后的页面列表
  57. current_pages =[]
  58. pdf_reader = None
  59. try:
  60. # 检查文件内容是否为BytesIO对象
  61. if isinstance(self.file_content, BytesIO):
  62. # 如果是BytesIO对象,则打开PDF文件
  63. pdf_reader = fitz.open(stream=self.file_content, filetype="pdf")
  64. else:
  65. # 如果不是BytesIO对象,则抛出ValueError异常
  66. raise ValueError(f"文件内容不是BytesIO对象。类型:{type(self.file_content)}")
  67. except Exception as e:
  68. # 如果在打开PDF时发生异常,则记录错误日志并返回空列表
  69. logger.error(f"打开PDF失败:{str(e)}")
  70. raise ValueError(f"打开PDF失败:{str(e)} {sys._getframe().f_lineno}")
  71. # 遍历PDF中的每一页
  72. for page_index in range(len(pdf_reader)):
  73. current_page = pdf_reader[page_index] # 获取当前页
  74. page_imag = self._process_page_images(current_page) # 处理页面图片
  75. page_text = self._extract_page_text(current_page, pdf_reader) # 提取当前页的文本内容
  76. processed_pages.append(page_text) # 将当前页的文本内容添加到处理后的页面列表中
  77. current_pages.append(page_imag)
  78. pdf_reader.close()
  79. return processed_pages,current_pages
  80. def _extract_page_text(self, page, pdf_reader):
  81. """
  82. 提取PDF页面的文本内容。
  83. Args:
  84. page (fitz.Page): 需要提取文本的PDF页面对象。
  85. pdf_reader (fitz.Document): PDF文件的阅读器对象。
  86. Returns:
  87. str: 页面的文本内容。
  88. """
  89. page_text = page.get_text() # 获取页面的文本内容
  90. for image_info in page.get_images():
  91. image_xref = image_info[0] # 获取图片的引用
  92. extracted_image = pdf_reader.extract_image(image_xref) # 提取图片
  93. raw_image_data = extracted_image["image"] # 获取图片的raw数据
  94. if isinstance(raw_image_data, bytes):
  95. extracted_text = self.KbmService.extract_text_from_image(raw_image_data) # 从图片中提取文本
  96. if extracted_text and len(extracted_text) > 5:
  97. page_text += extracted_text
  98. return page_text
  99. def _process_page_images(self, processed_pages):
  100. """
  101. 处理页面图片,渲染页面为图像,保存到MinIO,并记录图像名称。
  102. Args:
  103. processed_pages (list): 处理后的页面列表。
  104. Returns:
  105. list: 包含页面号和图像名称的元组列表。
  106. """
  107. page_images =[]
  108. try:
  109. # 遍历处理后的页面列表
  110. page_image = self.KbmService.render_page_to_image(processed_pages)
  111. # 保存图像到MinIO,并获取图像名称
  112. image_name = self.KbmService.save_image_to_minio(page_image, self.bucket_name)
  113. # 将页面号和图像名称作为元组添加到page_images列表中
  114. page_images.append(image_name)
  115. except Exception as e:
  116. logger.error(f'错误')
  117. # 返回包含所有页面图像信息的列表
  118. return page_images
  119. def _split_sentences(self, processed_pages,cmax_tokens=250):
  120. """
  121. 使用NLP将处理后的内容分割成句子。
  122. Args:
  123. processed_pages (list): 处理后的页面列表。
  124. Returns:
  125. list: 包含分割后的句子的列表。
  126. """
  127. # 初始化一个空列表,用于存储分割后的句子
  128. sentence_list = []
  129. # 遍历处理后的页面列表
  130. for page_text in processed_pages:
  131. # 使用SNLP模型处理文本,
  132. sentence_list.append( self.nlp.split_text_into_chunks(page_text,cmax_tokens))
  133. # 返回包含所有分割后的句子的列表
  134. return sentence_list