407 lines
15 KiB
Python
407 lines
15 KiB
Python
import json
|
||
import fitz # PyMuPDF
|
||
import os
|
||
import fitz
|
||
from collections import Counter
|
||
import uuid
|
||
import re
|
||
import logging
|
||
# 配置 logger
|
||
logging.basicConfig(
|
||
level=logging.INFO,
|
||
format="%(asctime)s - %(levelname)s - %(message)s",
|
||
datefmt="%Y-%m-%d %H:%M:%S"
|
||
)
|
||
logger = logging.getLogger(__name__)
|
||
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 文件的多个位置添加高亮注释。
|
||
逻辑融合：
|
||
1. 输入参数与保存路径遵循参考代码 B (输出到 /app/files/uploads_v1)。
|
||
2. 坐标处理遵循参考代码 A (假设 bbox 为 0-1000 归一化坐标，需转换为实际像素坐标)。
|
||
3. 返回值格式遵循参考代码 B (highlight_filename, most_common_page_1based, original_filename)。
|
||
|
||
:param input_pdf_path: 输入 PDF 文件路径（字符串）
|
||
:param highlights: 高亮区域列表，每个元素为 dict，格式：
|
||
{"page_idx": int, "bbox": [x0, y0, x1, y1]}
|
||
注意：bbox 应为 0-1000 归一化坐标
|
||
:param color: 高亮颜色，RGB 元组，范围 0~1，例如黄色 (1, 1, 0)
|
||
:return: tuple (highlight_filename: str, most_common_page: int or None, original_filename: str)
|
||
"""
|
||
unique_id = uuid.uuid4()
|
||
|
||
# 调试打印 (参照参考代码 B)
|
||
print(input_pdf_path)
|
||
logger.info(input_pdf_path)
|
||
|
||
if not highlights:
|
||
logger.info("⚠️ 无高亮区域，跳过处理。")
|
||
filename = os.path.basename(input_pdf_path)
|
||
return filename, None, filename
|
||
|
||
# 提取所有有效的 page_idx 用于统计
|
||
valid_page_indices = []
|
||
|
||
# 打开 PDF
|
||
try:
|
||
doc = fitz.open(input_pdf_path)
|
||
logger.info(f"已打开 PDF: {input_pdf_path}")
|
||
except Exception as e:
|
||
logger.info(f"❌ 无法打开 PDF: {e}")
|
||
filename = os.path.basename(input_pdf_path)
|
||
return filename, None, filename
|
||
|
||
for item in highlights:
|
||
page_idx = item.get("page_idx")
|
||
bbox = item.get("bbox")
|
||
|
||
# 基本校验 (参照参考代码 B)
|
||
if page_idx is None or bbox is None:
|
||
logger.info(f"⚠️ 跳过高亮项（缺少 page_idx 或 bbox）: {item}")
|
||
continue
|
||
if not (0 <= page_idx < doc.page_count):
|
||
logger.info(f"⚠️ 页面索引 {page_idx} 超出范围（总页数: {doc.page_count}），跳过。")
|
||
continue
|
||
if len(bbox) != 4:
|
||
logger.info(f"⚠️ bbox 格式错误（需 [x0,y0,x1,y1]）: {bbox}，跳过。")
|
||
continue
|
||
|
||
try:
|
||
page = doc[page_idx]
|
||
|
||
# --- 核心逻辑：坐标转换 (参照参考代码 A) ---
|
||
# 假设输入 bbox 是 0-1000 的归一化坐标
|
||
page_rect = page.rect
|
||
page_width = page_rect.width
|
||
page_height = page_rect.height
|
||
|
||
# 公式：实际坐标 = (归一化坐标 / 1000) * 页面实际尺寸
|
||
x0 = (bbox[0] / 1000) * page_width
|
||
y0 = (bbox[1] / 1000) * page_height
|
||
x1 = (bbox[2] / 1000) * page_width
|
||
y1 = (bbox[3] / 1000) * page_height
|
||
|
||
rect = fitz.Rect(x0, y0, x1, y1)
|
||
|
||
# 添加高亮标注
|
||
highlight_annot = page.add_highlight_annot(rect)
|
||
highlight_annot.set_colors(stroke=color)
|
||
highlight_annot.update()
|
||
|
||
valid_page_indices.append(page_idx) # 记录成功处理的 page_idx
|
||
except Exception as e:
|
||
logger.info(f"❌ 在 page={page_idx}, bbox={bbox} 添加高亮失败: {e}")
|
||
|
||
# 确定输出路径 (参照参考代码 B 的固定路径)
|
||
filename = os.path.basename(input_pdf_path)
|
||
base, ext = os.path.splitext(filename)
|
||
output_dir = "/app/files/uploads_v1"
|
||
|
||
# 确保目录存在
|
||
os.makedirs(output_dir, exist_ok=True)
|
||
|
||
output_pdf_path = os.path.join(output_dir, f"{base}_{unique_id}{ext}")
|
||
highlight_filename = f"{base}_{unique_id}{ext}"
|
||
|
||
# 保存 PDF (参照参考代码 B 的参数)
|
||
doc.save(output_pdf_path, garbage=3, deflate=True)
|
||
doc.close()
|
||
|
||
logger.info(f"✅ 高亮完成！已保存至: {output_pdf_path}")
|
||
|
||
# 统计最常出现的 page_idx (参照参考代码 B 的逻辑，返回 1-based)
|
||
most_common_page = None
|
||
if valid_page_indices:
|
||
counter = Counter(valid_page_indices)
|
||
most_common_page_0based, _ = counter.most_common(1)[0]
|
||
most_common_page = most_common_page_0based + 1 # 转换为 1-based 页码
|
||
|
||
# 返回值格式：(highlight_filename, most_common_page, original_filename)
|
||
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:
|
||
bbox = ins.get("bbox", [])
|
||
page_idx = ins.get("page_idx", -1)
|
||
highlight_list.append({
|
||
"page_idx": page_idx,
|
||
"bbox": bbox
|
||
})
|
||
|
||
return content, highlight_list
|
||
|
||
#获取指定表格过滤后的文本
|
||
def filter_content_bbox(data, max_length=3000):
|
||
"""
|
||
将 data 中的 text/table/image 内容按字符长度分块（每块不超过 max_length），
|
||
返回 [{"content": str, "highlight_list": list}, ...]
|
||
"""
|
||
chunks = []
|
||
current_content = ""
|
||
current_highlight_list = []
|
||
|
||
def _flush_chunk():
|
||
"""内部函数：将当前内容打包进 chunks，并重置"""
|
||
if current_content or current_highlight_list:
|
||
chunks.append({
|
||
"content": current_content,
|
||
"highlight_list": current_highlight_list.copy()
|
||
})
|
||
|
||
for ins in data:
|
||
if "type" not in ins:
|
||
continue
|
||
|
||
added_content = ""
|
||
bbox = ins.get("bbox", [])
|
||
page_idx = ins.get("page_idx", -1)+1
|
||
|
||
if ins["type"] == "text":
|
||
text_content = ins.get("text", "")
|
||
added_content = text_content + "\n"
|
||
|
||
elif ins["type"] == "table":
|
||
table_caption = ins.get("table_caption")
|
||
table_body = ins.get("table_body")
|
||
|
||
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)
|
||
|
||
added_content = caption_str + "\n" + (table_body or "") + "\n"
|
||
|
||
elif ins["type"] == "image":
|
||
image_caption = ins.get("image_caption", "")
|
||
img_path = ins.get("img_path", "")
|
||
|
||
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:
|
||
added_content = image_caption + "\n" + img_path + "\n"
|
||
else:
|
||
continue # 跳过空图像项
|
||
else:
|
||
continue
|
||
|
||
# 检查加上 added_content 后是否超限
|
||
if len(current_content) + len(added_content) > max_length:
|
||
# 如果当前块非空，先 flush
|
||
if current_content:
|
||
_flush_chunk()
|
||
current_content = ""
|
||
current_highlight_list = []
|
||
|
||
# 如果单个 added_content 超过 max_length，仍需保留（避免丢弃）
|
||
if len(added_content) > max_length:
|
||
# 强制截断或直接放入（这里选择直接放入，避免信息丢失）
|
||
current_content = added_content[:max_length]
|
||
current_highlight_list.append({"page_idx": page_idx, "bbox": bbox})
|
||
_flush_chunk()
|
||
# 剩余部分？此处简化处理：忽略超长部分（或可递归拆分，但复杂）
|
||
continue
|
||
else:
|
||
current_content = added_content
|
||
current_highlight_list.append({"page_idx": page_idx, "bbox": bbox})
|
||
else:
|
||
current_content += added_content
|
||
current_highlight_list.append({"page_idx": page_idx, "bbox": bbox})
|
||
|
||
# 处理最后一块
|
||
if current_content or current_highlight_list:
|
||
chunks.append({
|
||
"content": current_content,
|
||
"highlight_list": current_highlight_list
|
||
})
|
||
|
||
return chunks
|
||
|
||
import json
|
||
|
||
def data_replace(data, host_ip: str,port:int):
|
||
"""
|
||
遍历 data 列表，若元素包含 'img_path' 字段，则为其添加前缀。
|
||
"""
|
||
for ins in data:
|
||
if isinstance(ins, dict) and "img_path" in ins:
|
||
# 原路径可能是相对路径，如 'images/xxx.jpg'
|
||
original_path = ins["img_path"]
|
||
# 拼接完整 URL
|
||
ins["img_path"] = f"http://{host_ip}:{port}/output/{original_path}"
|
||
# 可选：打印调试信息
|
||
return data
|
||
|
||
def get_most_common_page(highlights):
|
||
"""从高亮列表计算最常出现的页码（1-based），不做任何文件 I/O"""
|
||
valid_page_indices = [
|
||
item["page_idx"] for item in highlights
|
||
if item.get("page_idx") is not None
|
||
]
|
||
if not valid_page_indices:
|
||
return None
|
||
counter = Counter(valid_page_indices)
|
||
most_common_0based, _ = counter.most_common(1)[0]
|
||
return most_common_0based + 1
|
||
|
||
|
||
def highlight_pdf_file_batch(input_pdf_path, all_highlights, color=(1, 1, 0), output_filename=None):
|
||
"""
|
||
批量高亮：将所有表格的高亮区域一次性写入 PDF，只打开/保存一次。
|
||
|
||
:param input_pdf_path: 原始 PDF 路径
|
||
:param all_highlights: 合并后的所有高亮列表 List[{"page_idx": int, "bbox": [...]}]
|
||
:param output_filename: 可选，指定输出文件名（不含目录）；不传则自动生成
|
||
:return: (highlight_filename, filename)
|
||
"""
|
||
filename = os.path.basename(input_pdf_path)
|
||
if output_filename is None:
|
||
unique_id = uuid.uuid4()
|
||
base, ext = os.path.splitext(filename)
|
||
output_filename = f"{base}_{unique_id}{ext}"
|
||
|
||
if not all_highlights:
|
||
return filename, filename
|
||
|
||
try:
|
||
doc = fitz.open(input_pdf_path)
|
||
except Exception as e:
|
||
print(f"❌ 无法打开 PDF: {e}")
|
||
return filename, filename
|
||
|
||
for item in all_highlights:
|
||
page_idx = item.get("page_idx")
|
||
bbox = item.get("bbox")
|
||
if page_idx is None or bbox is None:
|
||
continue
|
||
if not (0 <= page_idx < doc.page_count):
|
||
continue
|
||
if len(bbox) != 4:
|
||
continue
|
||
try:
|
||
page = doc[page_idx]
|
||
page_rect = page.rect
|
||
x0 = (bbox[0] / 1000) * page_rect.width
|
||
y0 = (bbox[1] / 1000) * page_rect.height
|
||
x1 = (bbox[2] / 1000) * page_rect.width
|
||
y1 = (bbox[3] / 1000) * page_rect.height
|
||
rect = fitz.Rect(x0, y0, x1, y1)
|
||
annot = page.add_highlight_annot(rect)
|
||
annot.set_colors(stroke=color)
|
||
annot.update()
|
||
except Exception as e:
|
||
print(f"❌ 添加高亮失败 page={page_idx}: {e}")
|
||
|
||
output_dir = "/app/files/uploads_v1"
|
||
os.makedirs(output_dir, exist_ok=True)
|
||
highlight_filename = output_filename
|
||
output_pdf_path = os.path.join(output_dir, highlight_filename)
|
||
doc.save(output_pdf_path, garbage=3, deflate=True)
|
||
doc.close()
|
||
print(f"✅ 批量高亮完成，已保存至: {output_pdf_path}")
|
||
return highlight_filename, filename
|
||
|
||
|
||
# if __name__ == "__main__":
|
||
# host_ip = "192.168.0.46" # 👈 替换为你实际的主机 IP 或域名
|
||
|
||
# file_path = '/storage01/home/hdf/project/wxkgrag/kgextract/output/122-06A0014-B01001_发动机-维修手册_content_list.json'
|
||
# try:
|
||
# with open(file_path, 'r', encoding='utf-8') as file:
|
||
# data = json.load(file)
|
||
# except Exception as e:
|
||
# print(f"读取 JSON 文件时出错：{e}")
|
||
# data = []
|
||
|
||
# # 修改 img_path 字段
|
||
# data_replace(data=data, host_ip=host_ip)
|
||
|
||
# # 可选：保存回文件或打印验证
|
||
# # 示例：打印所有包含 img_path 的项
|
||
# for item in data:
|
||
# if "img_path" in item:
|
||
# print(item["img_path"])
|