PDF_Processor.py 6.4 KB

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