155 lines
5.3 KiB
Python
155 lines
5.3 KiB
Python
|
||
import json
|
||
import fitz # PyMuPDF
|
||
import os
|
||
import fitz
|
||
from collections import Counter
|
||
import uuid
|
||
import re
|
||
def safe_json_loads(s: str):
|
||
try:
|
||
return json.loads(s)
|
||
except json.JSONDecodeError as e:
|
||
# 尝试修复常见问题
|
||
s = s.replace("\\n", "\\\n").replace("\r", "\\r").replace("\t", "\\t")
|
||
# 或者直接清理非法字符
|
||
s = re.sub(r'[\x00-\x1f\x7f]', '', s) # 移除控制字符
|
||
try:
|
||
return json.loads(s)
|
||
except:
|
||
raise e
|
||
|
||
def highlight_pdf_file(input_pdf_path, highlights, color=(1, 1, 0)):
|
||
"""
|
||
在单个 PDF 文件的多个位置添加高亮注释，并返回处理后路径及最常出现的 page_idx。
|
||
|
||
:param input_pdf_path: 输入 PDF 文件路径（字符串）
|
||
:param highlights: 高亮区域列表，每个元素为 dict，格式：
|
||
{"page_idx": int, "bbox": [x0, y0, x1, y1]}
|
||
:param output_dir: 输出目录路径。若为 None，则使用输入文件所在目录。
|
||
:param color: 高亮颜色，RGB 元组，范围 0~1，例如黄色 (1, 1, 0)
|
||
:return: tuple (output_pdf_path: str, most_common_page: int or None)
|
||
"""
|
||
unique_id = uuid.uuid4()
|
||
|
||
if not highlights:
|
||
print("⚠️ 无高亮区域，跳过处理。")
|
||
return input_pdf_path, None,None
|
||
|
||
# 提取所有有效的 page_idx 用于统计
|
||
valid_page_indices = []
|
||
|
||
# 打开 PDF
|
||
doc = fitz.open(input_pdf_path)
|
||
|
||
for item in highlights:
|
||
page_idx = item.get("page_idx")
|
||
bbox = item.get("bbox")
|
||
|
||
# 基本校验
|
||
if page_idx is None or bbox is None:
|
||
print(f"⚠️ 跳过高亮项（缺少 page_idx 或 bbox）: {item}")
|
||
continue
|
||
if not (0 <= page_idx < doc.page_count):
|
||
print(f"⚠️ 页面索引 {page_idx} 超出范围（总页数: {doc.page_count}），跳过。")
|
||
continue
|
||
if len(bbox) != 4:
|
||
print(f"⚠️ bbox 格式错误（需 [x0,y0,x1,y1]）: {bbox}，跳过。")
|
||
continue
|
||
|
||
try:
|
||
page = doc[page_idx]
|
||
highlight_annot = page.add_highlight_annot(bbox)
|
||
highlight_annot.set_colors(stroke=color)
|
||
highlight_annot.update()
|
||
valid_page_indices.append(page_idx) # 记录成功处理的 page_idx
|
||
except Exception as e:
|
||
print(f"❌ 在 page={page_idx}, bbox={bbox} 添加高亮失败: {e}")
|
||
|
||
# 确定输出路径
|
||
filename = os.path.basename(input_pdf_path)
|
||
base, ext = os.path.splitext(filename)
|
||
output_dir = "/app/files/uploads_v1"
|
||
|
||
output_pdf_path = os.path.join(output_dir, f"{base}_{unique_id}{ext}")
|
||
highlight_filename = f"{base}_{unique_id}{ext}"
|
||
# 保存 PDF
|
||
doc.save(output_pdf_path, garbage=3, deflate=True)
|
||
doc.close()
|
||
|
||
print(f"✅ 高亮完成！已保存至: {output_pdf_path}")
|
||
|
||
# 统计最常出现的 page_idx
|
||
most_common_page = None
|
||
if valid_page_indices:
|
||
counter = Counter(valid_page_indices)
|
||
most_common_page, _ = counter.most_common(1)[0]
|
||
|
||
return highlight_filename, most_common_page,filename
|
||
|
||
|
||
def get_content_bbox(data):
|
||
content = ""
|
||
highlight_list = []
|
||
for ins in data:
|
||
if "type" not in ins:
|
||
continue
|
||
|
||
if ins["type"] == "text":
|
||
text_content = ins.get("text", "")
|
||
content += text_content + "\n"
|
||
bbox = ins.get("bbox", [])
|
||
page_idx = ins.get("page_idx", -1)
|
||
highlight_list.append({
|
||
"page_idx": page_idx,
|
||
"bbox": bbox
|
||
})
|
||
|
||
elif ins["type"] == "table":
|
||
table_caption = ins.get("table_caption")
|
||
table_body = ins.get("table_body")
|
||
|
||
# 处理 table_caption：可能是 list 或 str
|
||
caption_str = ""
|
||
if table_caption:
|
||
if isinstance(table_caption, list):
|
||
caption_str = "\n".join(str(c) for c in table_caption)
|
||
else:
|
||
caption_str = str(table_caption)
|
||
|
||
content += caption_str + "\n" + (table_body or "") + "\n"
|
||
bbox = ins.get("bbox", [])
|
||
page_idx = ins.get("page_idx", -1)
|
||
highlight_list.append({
|
||
"page_idx": page_idx,
|
||
"bbox": bbox
|
||
})
|
||
|
||
elif ins["type"] == "image":
|
||
image_caption = ins.get("image_caption", "")
|
||
img_path = ins.get("img_path", "")
|
||
|
||
# 关键修复：处理 image_caption 可能是 list 的情况
|
||
if isinstance(image_caption, list):
|
||
image_caption = "\n".join(str(item) for item in image_caption)
|
||
elif not isinstance(image_caption, str):
|
||
image_caption = str(image_caption)
|
||
|
||
if isinstance(img_path, list):
|
||
img_path = "\n".join(str(p) for p in img_path)
|
||
elif not isinstance(img_path, str):
|
||
img_path = str(img_path)
|
||
|
||
if image_caption or img_path:
|
||
content += image_caption + "\n" + img_path + "\n"
|
||
bbox = ins.get("bbox", [])
|
||
page_idx = ins.get("page_idx", -1)
|
||
highlight_list.append({
|
||
"page_idx": page_idx,
|
||
"bbox": bbox
|
||
})
|
||
else:
|
||
continue
|
||
|
||
return content, highlight_list
|