259 lines
11 KiB
Python
259 lines
11 KiB
Python
import json
|
||
import fitz # PyMuPDF
|
||
import os
|
||
from collections import Counter
|
||
from Prompt import Prompt_guzhang
|
||
from openai import OpenAI
|
||
from extract_util import safe_json_loads, highlight_pdf_file, get_content_bbox,get_filter_content_bbox
|
||
|
||
|
||
|
||
|
||
def get_compose_chunk_bbox_and_removed_data(data):
|
||
"""
|
||
返回：
|
||
- 符合条件的故障表分组（每组 [前前, 前, 当前]）
|
||
- 从原始 data 中移除了所有相关记录后的新 data 列表
|
||
"""
|
||
title = "组成"
|
||
|
||
# 支持多个 label 模式（可扩展）
|
||
labels = [
|
||
'<table><tr><td rowspan=1 colspan=1>序号</td><td rowspan=1 colspan=1>组成部分</td><td rowspan=1 colspan=1>功能</td><td rowspan=1 colspan=1>数量</td></tr><tr><td rowspan=1 colspan=1>1</td>',
|
||
# 可在此添加更多变体，例如：
|
||
# '<table><tr><th>序号</th><th>故障现象</th><th>故障原因</th><th>维修项目</th></tr>',
|
||
# '<table><tr><td>序号</td><td>故障现象</td><td>故障原因</td><td>维修项目</td></tr>',
|
||
]
|
||
|
||
# 构建 id -> record 映射
|
||
record_map = {item["id"]: item for item in data}
|
||
|
||
results = []
|
||
ids_to_remove = set()
|
||
|
||
for ins in data:
|
||
if ins.get("type") == "table":
|
||
# 检查 caption 是否包含 "故障表"
|
||
caption_has_title = any(title in caption for caption in ins.get("table_caption", []))
|
||
|
||
# 检查 body 是否匹配任意一个 label
|
||
body = ins.get("table_body", "")
|
||
body_matches_any_label = any(label in body for label in labels)
|
||
|
||
if caption_has_title and body_matches_any_label:
|
||
current_id = ins["id"]
|
||
prev_ids = [current_id - 2, current_id - 1]
|
||
|
||
# 收集要移除的 ID（前两条 + 当前）
|
||
group_ids = [pid for pid in prev_ids if pid >= 1] + [current_id]
|
||
ids_to_remove.update(group_ids)
|
||
|
||
# 构建分组
|
||
group = [record_map.get(pid) for pid in prev_ids if pid >= 1] + [ins]
|
||
results.append(group)
|
||
|
||
# 过滤原始 data
|
||
filtered_data = [item for item in data if item["id"] not in ids_to_remove]
|
||
|
||
return results, filtered_data
|
||
|
||
def get_tujiecompose_chunk_bbox_and_removed_data(data):
|
||
"""
|
||
返回：
|
||
- 符合条件的故障表分组（每组 [前前, 前, 当前]）
|
||
- 从原始 data 中移除了所有相关记录后的新 data 列表
|
||
"""
|
||
title = "组成"
|
||
|
||
# 支持多个 label 模式（可扩展）
|
||
labels = [
|
||
'<table><tr><td colspan=\"1\" rowspan=\"1\">组成编码</td><td colspan=\"1\" rowspan=\"1\">名称</td><td colspan=\"1\" rowspan=\"1\">是否为关重件</td><td colspan=\"1\" rowspan=\"1\">是否为寿命件</td>',
|
||
'<table><tr><td rowspan=1 colspan=1>序号</td><td rowspan=1 colspan=1>组成部分</td><td rowspan=1 colspan=1>功能</td><td rowspan=1 colspan=1>数量</td></tr><tr><td rowspan=1 colspan=1>1</td>'
|
||
# 可在此添加更多变体，例如：
|
||
# '<table><tr><th>序号</th><th>故障现象</th><th>故障原因</th><th>维修项目</th></tr>',
|
||
# '<table><tr><td>序号</td><td>故障现象</td><td>故障原因</td><td>维修项目</td></tr>',
|
||
]
|
||
|
||
# 构建 id -> record 映射
|
||
record_map = {item["id"]: item for item in data}
|
||
|
||
results = []
|
||
ids_to_remove = set()
|
||
|
||
for ins in data:
|
||
if ins.get("type") == "table":
|
||
# 检查 caption 是否包含 "故障表"
|
||
caption_has_title = any(title in caption for caption in ins.get("table_caption", []))
|
||
|
||
# 检查 body 是否匹配任意一个 label
|
||
body = ins.get("table_body", "")
|
||
body_matches_any_label = any(label in body for label in labels)
|
||
|
||
if caption_has_title and body_matches_any_label:
|
||
current_id = ins["id"]
|
||
prev_ids = [current_id - 2, current_id - 1]
|
||
|
||
# 收集要移除的 ID（前两条 + 当前）
|
||
group_ids = [pid for pid in prev_ids if pid >= 1] + [current_id]
|
||
ids_to_remove.update(group_ids)
|
||
|
||
# 构建分组
|
||
group = [record_map.get(pid) for pid in prev_ids if pid >= 1] + [ins]
|
||
results.append(group)
|
||
|
||
# 过滤原始 data
|
||
filtered_data = [item for item in data if item["id"] not in ids_to_remove]
|
||
|
||
return results, filtered_data
|
||
|
||
def get_guzhang_chunk_bbox_and_removed_data(data):
|
||
title = "故障表"
|
||
label = "<table><tr><td colspan=\"1\" rowspan=\"1\">序号</td><td colspan=\"1\" rowspan=\"1\">故障现象</td><td colspan=\"1\" rowspan=\"1\">故障原因</td><td colspan=\"1\" rowspan=\"1\">维修项目</td></tr>"
|
||
|
||
record_map = {item["id"]: item for item in data}
|
||
results = []
|
||
ids_to_remove = set()
|
||
|
||
for ins in data:
|
||
if ins.get("type") == "table":
|
||
caption_has_title = any(title in caption for caption in ins.get("table_caption", []))
|
||
body_has_label = label in ins.get("table_body", "")
|
||
if caption_has_title and body_has_label:
|
||
current_id = ins["id"]
|
||
prev_ids = [current_id - 2, current_id - 1]
|
||
group_ids = [pid for pid in prev_ids if pid >= 1] + [current_id]
|
||
ids_to_remove.update(group_ids)
|
||
group = [record_map.get(pid) for pid in prev_ids if pid >= 1] + [ins]
|
||
results.append(group)
|
||
|
||
filtered_data = [item for item in data if item["id"] not in ids_to_remove]
|
||
return results, filtered_data
|
||
|
||
|
||
START_MARKERS = [
|
||
# 标准格式（带 rowspan/colspan）
|
||
"<table><tr><td rowspan=1 colspan=1>维修项目编号</td><td rowspan=1 colspan=1>组成编码</td><td rowspan=1 colspan=1>名称</td><td rowspan=1 colspan=1>维修项目</td><td rowspan=1 colspan=1>维修间隔期</td>",
|
||
|
||
# 可在此添加其他格式，例如：
|
||
# "<table><tr><td>维修项目编号</td><td>组成编码</td><td>名称</td><td>维修项目</td><td>维修间隔期</td>",
|
||
]
|
||
|
||
def extract_maintenance_groups_and_remaining(records, markers=START_MARKERS, max_following_for_last=None):
|
||
"""
|
||
普通起始标记：组 = [start, ..., next_start - 1]
|
||
最后一个起始标记：组 = [start, ..., start + max_following_for_last] （若指定）
|
||
"""
|
||
groups = []
|
||
current_group = None
|
||
used_ids = set()
|
||
|
||
for record in records:
|
||
body = record.get("table_body", "")
|
||
is_start = (record.get("type") == "table" and
|
||
any(marker in body for marker in markers))
|
||
|
||
if is_start:
|
||
if current_group is not None:
|
||
groups.append(current_group)
|
||
used_ids.update(r["id"] for r in current_group)
|
||
current_group = [record]
|
||
else:
|
||
if current_group is not None:
|
||
# 如果指定了 max_following_for_last，且这是最后一个组（暂时不知道是不是最后一个）
|
||
# 我们无法在遍历时知道“这是不是最后一个起始标记”
|
||
# 所以更好的办法是：先按原逻辑分组，再处理最后一个组
|
||
current_group.append(record)
|
||
|
||
if current_group is not None:
|
||
groups.append(current_group)
|
||
used_ids.update(r["id"] for r in current_group)
|
||
|
||
# —————— 关键：后处理最后一个组 ——————
|
||
if groups and max_following_for_last is not None:
|
||
last_group = groups[-1]
|
||
if len(last_group) > max_following_for_last + 1: # +1 包含起始记录
|
||
truncated = last_group[:max_following_for_last + 1]
|
||
# 更新 used_ids：移除被截掉的 ID
|
||
removed_ids = {r["id"] for r in last_group[max_following_for_last + 1:]}
|
||
used_ids -= removed_ids
|
||
groups[-1] = truncated
|
||
|
||
remaining_records = [r for r in records if r["id"] not in used_ids]
|
||
return groups, remaining_records
|
||
|
||
|
||
START_MARKERS_ope = [
|
||
'<table><tr><td>项目编号</td><td>操作项目</td></tr><tr>',
|
||
'<table><tr><td rowspan=1 colspan=1>项目编号</td><td rowspan=1 colspan=1>操作项目</td></tr><tr>',
|
||
'<table><tr><td colspan="1" rowspan="1">项目编号</td><td colspan="1" rowspan="1">操作项目</td></tr><tr>'
|
||
]
|
||
|
||
def extract_operation_groups_and_remaining(records, markers=START_MARKERS_ope, max_following_for_last=None):
|
||
"""
|
||
普通起始标记：组 = [start, ..., next_start - 1]
|
||
最后一个起始标记：组 = [start, ..., start + max_following_for_last] （若指定）
|
||
"""
|
||
groups = []
|
||
current_group = None
|
||
used_ids = set()
|
||
|
||
for record in records:
|
||
body = record.get("table_body", "")
|
||
is_start = (record.get("type") == "table" and
|
||
any(marker in body for marker in markers))
|
||
|
||
if is_start:
|
||
if current_group is not None:
|
||
groups.append(current_group)
|
||
used_ids.update(r["id"] for r in current_group)
|
||
current_group = [record]
|
||
else:
|
||
if current_group is not None:
|
||
# 如果指定了 max_following_for_last，且这是最后一个组（暂时不知道是不是最后一个）
|
||
# 我们无法在遍历时知道“这是不是最后一个起始标记”
|
||
# 所以更好的办法是：先按原逻辑分组，再处理最后一个组
|
||
current_group.append(record)
|
||
|
||
if current_group is not None:
|
||
groups.append(current_group)
|
||
used_ids.update(r["id"] for r in current_group)
|
||
|
||
# —————— 关键：后处理最后一个组 ——————
|
||
if groups and max_following_for_last is not None:
|
||
last_group = groups[-1]
|
||
if len(last_group) > max_following_for_last + 1: # +1 包含起始记录
|
||
truncated = last_group[:max_following_for_last + 1]
|
||
# 更新 used_ids：移除被截掉的 ID
|
||
removed_ids = {r["id"] for r in last_group[max_following_for_last + 1:]}
|
||
used_ids -= removed_ids
|
||
groups[-1] = truncated
|
||
|
||
remaining_records = [r for r in records if r["id"] not in used_ids]
|
||
return groups, remaining_records
|
||
|
||
|
||
|
||
def get_chunk_filter(data,max_length=7000):
|
||
for idx, item in enumerate(data, start=1):
|
||
item["id"] = idx
|
||
groups, remaining_records_repair = extract_maintenance_groups_and_remaining(data)
|
||
groups, remaining_records_operation = extract_operation_groups_and_remaining(remaining_records_repair)
|
||
groups, remaining_records_compose = get_compose_chunk_bbox_and_removed_data(remaining_records_operation)
|
||
groups, remaining_records_composetujie =get_tujiecompose_chunk_bbox_and_removed_data(remaining_records_compose)
|
||
results, filtered_data=get_guzhang_chunk_bbox_and_removed_data(remaining_records_composetujie)
|
||
res = get_filter_content_bbox(filtered_data,max_length=max_length)
|
||
return res
|
||
|
||
# if __name__ == "__main__":
|
||
# file_path = '/storage01/home/hdf/project/wxkgrag/kgextract/output/122-06A0014-B01001_发动机-维修手册_content_list.json'
|
||
# input_pdf_path = "/storage01/home/hdf/project/wxkgrag/kgextract/know/122-06A0014-B01001_发动机-维修手册_highlighted_fdb9e9db-98d2-4ef6-bc72-52f875453f85.pdf"
|
||
# try:
|
||
# with open(file_path, 'r', encoding='utf-8') as file:
|
||
# data = json.load(file)
|
||
# except Exception as e:
|
||
# print(f"读取 JSON 文件时出错：{e}")
|
||
# data = []
|
||
# res = get_chunk_filter(data=data)
|
||
# print(res)
|
||
|