123 lines
5.6 KiB
Python
123 lines
5.6 KiB
Python
|
import json
|
|||
|
|
import os
|
|||
|
|
import fitz
|
|||
|
|
import PyPDF2
|
|||
|
|
import tempfile
|
|||
|
|
from flask_app.general.clean_pdf import extract_common_header, clean_page_content
|
|||
|
|
from flask_app.general.llm.doubao import doubao_model
|
|||
|
|
from flask_app.old_version.table_ocr_old import CommonOcr
|
|||
|
|
|
|||
|
|
def extract_img_table_text(ocr_result_pages):
|
|||
|
|
# 调用豆包对json形式的表格数据进行重构
|
|||
|
|
# print(ocr_result_pages)
|
|||
|
|
base_prompt = '''
|
|||
|
|
任务：你负责解析以json形式输入的表格信息，根据提供的文件内容来恢复表格信息并输出，不要遗漏任何一个文字。
|
|||
|
|
|
|||
|
|
要求与指南：
|
|||
|
|
1. 请运用文档表格理解能力，根据文件内容定位表格的每列和每行，列与列之间用丨分隔，若某列或者某行没有信息则用/填充。
|
|||
|
|
2. 请不要遗漏任何一个文字，同时不要打乱行与行之间的顺序，也不要打乱列与列之间的顺序，严格按照文字的位置信息来恢复表格信息。
|
|||
|
|
示例输出：
|
|||
|
|
表格标题：
|
|||
|
|
|序号|名称|数量|单位|单价(元)|总价(元)|技术参数|备注|
|
|||
|
|
|形象展示区|
|
|||
|
|
|1|公园主E题雕塑|1|套|/|/|根据江夏体育与文化元素定制|/|
|
|||
|
|
|
|||
|
|
..........
|
|||
|
|
'''
|
|||
|
|
base_prompt += f"\n\n文件内容:\n{json.dumps(ocr_result_pages, ensure_ascii=False, indent=4)}"
|
|||
|
|
model_res = doubao_model(base_prompt)
|
|||
|
|
return model_res
|
|||
|
|
|
|||
|
|
# 判断pdf中是否有图片, 并输出含有图片的页面列表
|
|||
|
|
def has_images(pdf_path):
|
|||
|
|
# 打开PDF文件
|
|||
|
|
pdf_document = fitz.open(pdf_path)
|
|||
|
|
# 存储包含图片的页面页数
|
|||
|
|
pages_with_imgs = {}
|
|||
|
|
# 遍历PDF的每一页
|
|||
|
|
for page_num in range(pdf_document.page_count):
|
|||
|
|
page = pdf_document.load_page(page_num)
|
|||
|
|
# 获取页面的图片列表
|
|||
|
|
images = page.get_images(full=True)
|
|||
|
|
# 如果页面中有图片，返回True
|
|||
|
|
if images:
|
|||
|
|
pages_with_imgs[page_num + 1] = images
|
|||
|
|
# 如果遍历了所有页面，都没有图片，则返回False
|
|||
|
|
return pages_with_imgs
|
|||
|
|
|
|||
|
|
# 调用通用表格识别对图片中的表格进行提取，放回json形式的表格结构
|
|||
|
|
def table_ocr_extract(image_path):
|
|||
|
|
table_ocr = CommonOcr(img_path=image_path) # 创建时传递 img_path
|
|||
|
|
return table_ocr.recognize()
|
|||
|
|
|
|||
|
|
# 提取pdf中某一页的所有图片
|
|||
|
|
def extract_images_from_page(pdf_path, image_list, page_num):
|
|||
|
|
images = []
|
|||
|
|
try:
|
|||
|
|
doc = fitz.open(pdf_path)
|
|||
|
|
|
|||
|
|
for img in image_list:
|
|||
|
|
xref = img[0]
|
|||
|
|
base_image = doc.extract_image(xref)
|
|||
|
|
image_bytes = base_image['image']
|
|||
|
|
image_ext = base_image.get('ext', 'png')
|
|||
|
|
images.append({'data': image_bytes, 'ext': image_ext, 'page_num': page_num + 1})
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"提取图片时出错: {e}")
|
|||
|
|
return images
|
|||
|
|
def pdf_image2txt(file_path, img_pdf_list):
|
|||
|
|
common_header = extract_common_header(file_path)
|
|||
|
|
# print(f"公共抬头：{common_header}")
|
|||
|
|
# print("--------------------正文开始-------------------")
|
|||
|
|
result = ""
|
|||
|
|
pdf_document = fitz.open(file_path)
|
|||
|
|
with open(file_path, 'rb') as file:
|
|||
|
|
reader = PyPDF2.PdfReader(file)
|
|||
|
|
num_pages = len(reader.pages)
|
|||
|
|
# print(f"Total pages: {num_pages}")
|
|||
|
|
for page_num in range(num_pages):
|
|||
|
|
page = reader.pages[page_num]
|
|||
|
|
# text = page.extract_text()
|
|||
|
|
text = page.extract_text() or ""
|
|||
|
|
cleaned_text = clean_page_content(text, common_header)
|
|||
|
|
# # print(f"--------第{page_num}页-----------")
|
|||
|
|
if (page_num + 1) in img_pdf_list:
|
|||
|
|
print(f"第 {page_num + 1} 页含有图片，开始提取图片并OCR")
|
|||
|
|
images = extract_images_from_page(file_path, img_pdf_list[page_num + 1], page_num)
|
|||
|
|
for img in images:
|
|||
|
|
try:
|
|||
|
|
with tempfile.NamedTemporaryFile(delete=False, suffix='.' + img['ext']) as temp_image:
|
|||
|
|
temp_image.write(img['data'])
|
|||
|
|
temp_image.flush()
|
|||
|
|
# 调用OCR函数
|
|||
|
|
ocr_result = table_ocr_extract(temp_image.name)
|
|||
|
|
ocr_result = json.loads(ocr_result)
|
|||
|
|
# 判断是否提取成功并且 pages 中有数据
|
|||
|
|
if ocr_result['code'] == 200 and len(ocr_result['result']['pages']) > 0:
|
|||
|
|
print("提取成功，图片数据已提取。")
|
|||
|
|
ocr_result_pages = ocr_result['result']['pages']
|
|||
|
|
table_text = extract_img_table_text(ocr_result_pages)
|
|||
|
|
if table_text.strip():
|
|||
|
|
cleaned_text += "\n" + table_text
|
|||
|
|
else:
|
|||
|
|
print("提取失败或没有页面数据。")
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"OCR处理失败: {e}")
|
|||
|
|
finally:
|
|||
|
|
try:
|
|||
|
|
os.remove(temp_image.name)
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"删除临时文件失败: {e}")
|
|||
|
|
result += cleaned_text
|
|||
|
|
|
|||
|
|
directory = os.path.dirname(os.path.abspath(file_path))
|
|||
|
|
output_path = os.path.join(directory, 'extract.txt')
|
|||
|
|
# 将结果保存到 extract.txt 文件中
|
|||
|
|
try:
|
|||
|
|
with open(output_path, 'w', encoding='utf-8') as output_file:
|
|||
|
|
output_file.write(result)
|
|||
|
|
print(f"提取内容已保存到: {output_path}")
|
|||
|
|
except IOError as e:
|
|||
|
|
print(f"写入文件时发生错误: {e}")
|
|||
|
|
# 返回保存的文件路径
|
|||
|
|
return output_path
|