388 lines
16 KiB
Python
388 lines
16 KiB
Python
import pandas as pd
|
||
from io import StringIO # 1. 导入 StringIO
|
||
import json
|
||
import requests
|
||
import json
|
||
import os
|
||
from extract_util import safe_json_loads,highlight_pdf_file,get_content_bbox
|
||
import re
|
||
import asyncio
|
||
from bs4 import BeautifulSoup
|
||
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 save_to_file(data, filepath, is_json=True):
|
||
"""通用保存函数"""
|
||
try:
|
||
with open(filepath, "w", encoding="utf-8") as f:
|
||
if is_json:
|
||
json.dump(data, f, ensure_ascii=False, indent=2)
|
||
else:
|
||
f.write(data)
|
||
print(f"✅ 已保存: {filepath}")
|
||
return True
|
||
except Exception as e:
|
||
print(f"❌ 保存失败 {filepath}: {e}")
|
||
return False
|
||
def is_fault_table_by_first_tr(table_body: str) -> bool:
|
||
if not isinstance(table_body, str):
|
||
return False
|
||
|
||
# 非贪婪匹配第一个 <tr> ... </tr>
|
||
match = re.search(r'<tr[^>]*>(.*?)</tr>', table_body, re.IGNORECASE | re.DOTALL)
|
||
if not match:
|
||
|
||
return False
|
||
|
||
first_tr_html = match.group(1)
|
||
|
||
# 提取所有 <td> 或 <th> 中的文本（简单去标签）
|
||
# 更健壮的做法是递归去除标签，但这里用正则足够
|
||
text_in_first_tr = re.sub(r'<[^>]+>', '', first_tr_html) # 去除所有 HTML 标签
|
||
|
||
# 统一转为字符串并检查关键词
|
||
text_lower = text_in_first_tr.lower()
|
||
keywords = ["故障现象", "故障原因", "维修项目"]
|
||
return all(kw in text_lower for kw in keywords)
|
||
|
||
# ================= 新增：识别故障表类型 =================
|
||
def get_fault_table_type(table_body: str):
|
||
"""
|
||
返回故障表类型：
|
||
'type1' -> 包含 [故障现象, 故障原因, 维修项目]
|
||
'type2' -> 包含 [故障现象, 可能原因, 排除方法]
|
||
None -> 非故障表
|
||
"""
|
||
if not isinstance(table_body, str):
|
||
return None
|
||
|
||
match = re.search(r'<tr[^>]*>(.*?)</tr>', table_body, re.IGNORECASE | re.DOTALL)
|
||
if not match:
|
||
return None
|
||
|
||
text_lower = re.sub(r'<[^>]+>', '', match.group(1)).lower()
|
||
|
||
if all(kw in text_lower for kw in ["故障现象", "故障原因", "维修项目"]):
|
||
return "type1"
|
||
elif all(kw in text_lower for kw in ["故障现象", "可能原因", "排除方法"]):
|
||
return "type2"
|
||
elif all(kw in text_lower for kw in ["故障外部现象", "故障可能原因", "操作员指导","排除方法"]):
|
||
return "type3"
|
||
return None
|
||
|
||
|
||
# ================= 修改：返回时携带 table_type =================
|
||
def get_guzhang_chunk_bbox_and_removed_data(data):
|
||
results = [] # 元素格式: {"group": [...], "table_type": "type1"/"type2"}
|
||
ids_to_remove = set()
|
||
record_map = {item["id"]: item for item in data}
|
||
|
||
for ins in data:
|
||
if ins.get("type") == "table":
|
||
table_type = get_fault_table_type(ins.get("table_body", ""))
|
||
|
||
if table_type is not None:
|
||
current_id = ins["id"]
|
||
prev_ids = [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": group, "table_type": table_type}) # ← 携带类型
|
||
|
||
filtered_data = [item for item in data if item["id"] not in ids_to_remove]
|
||
return results, filtered_data
|
||
|
||
import re
|
||
|
||
def _expand_rowspans(content: str) -> list[list[str]]:
|
||
"""
|
||
将含 rowspan 的 HTML 表格展开为完整二维列表。
|
||
每个被合并的单元格都会被填充为原始内容。
|
||
"""
|
||
all_rows = re.findall(r"<tr>(.*?)</tr>", content, flags=re.DOTALL)
|
||
if not all_rows:
|
||
return []
|
||
|
||
matrix = []
|
||
pending = {} # {col_index: (remaining_rows, cell_text)}
|
||
|
||
for row_html in all_rows:
|
||
cells = re.findall(r"<td([^>]*)>(.*?)</td>", row_html, flags=re.DOTALL)
|
||
row = []
|
||
col = 0
|
||
cell_idx = 0
|
||
|
||
while cell_idx < len(cells) or any(c >= col for c in pending):
|
||
if col in pending:
|
||
# 该列被上方 rowspan 占用，直接填充
|
||
remaining, text = pending[col]
|
||
row.append(text)
|
||
if remaining - 1 > 0:
|
||
pending[col] = (remaining - 1, text)
|
||
else:
|
||
del pending[col]
|
||
col += 1
|
||
|
||
elif cell_idx < len(cells):
|
||
attrs_str, text = cells[cell_idx]
|
||
text = text.strip()
|
||
|
||
rowspan_m = re.search(r'rowspan\s*=\s*["\']?(\d+)', attrs_str)
|
||
rowspan = int(rowspan_m.group(1)) if rowspan_m else 1
|
||
|
||
row.append(text)
|
||
if rowspan > 1:
|
||
pending[col] = (rowspan - 1, text)
|
||
|
||
col += 1
|
||
cell_idx += 1
|
||
else:
|
||
break
|
||
|
||
matrix.append(row)
|
||
|
||
return matrix
|
||
|
||
|
||
def parse_table_content(content: str, chunk_size: int = 10) -> list[dict]:
|
||
"""
|
||
先展开 rowspan，再按 chunk_size 切分，分批解析后汇总返回。
|
||
"""
|
||
matrix = _expand_rowspans(content)
|
||
if not matrix:
|
||
print("⚠️ 未找到有效表格，返回空列表")
|
||
return []
|
||
|
||
header = matrix[0] # 表头
|
||
data_rows = matrix[1:] # 数据行（rowspan 已展开）
|
||
total = len(data_rows)
|
||
|
||
print(f"📋 展开后共 {total} 条数据行，按每批 {chunk_size} 行切分处理")
|
||
|
||
def _to_html(rows: list[list[str]]) -> str:
|
||
def row_html(cells):
|
||
return "<tr>" + "".join(f"<td>{c}</td>" for c in cells) + "</tr>"
|
||
return "<table>" + row_html(header) + "".join(row_html(r) for r in rows) + "</table>"
|
||
|
||
result_list = []
|
||
|
||
for chunk_start in range(0, total, chunk_size):
|
||
chunk = data_rows[chunk_start: chunk_start + chunk_size]
|
||
chunk_html = _to_html(chunk)
|
||
|
||
try:
|
||
chunk_df_list = pd.read_html(StringIO(chunk_html), header=0)
|
||
if not chunk_df_list:
|
||
print(f" ⚠️ 切片 [{chunk_start}:{chunk_start+chunk_size}] 解析失败，跳过")
|
||
continue
|
||
|
||
chunk_result = json.loads(
|
||
chunk_df_list[0].to_json(orient='records', force_ascii=False)
|
||
)
|
||
result_list.extend(chunk_result)
|
||
print(f" ✅ 切片 [{chunk_start}:{chunk_start+chunk_size}] 解析 {len(chunk_result)} 条")
|
||
|
||
except Exception as e:
|
||
print(f" ❌ 切片 [{chunk_start}:{chunk_start+chunk_size}] 处理异常: {e}")
|
||
continue
|
||
|
||
print(f"共汇总 {len(result_list)} 条记录")
|
||
return result_list
|
||
# ================= 修改：主函数按类型分支处理 =================
|
||
async def get_guzhang_node_relation(input_pdf_path, data, prefix_url="http://192.168.0.46:59085", device="发动机"):
|
||
for idx, item in enumerate(data, start=1):
|
||
item["id"] = idx
|
||
|
||
guzhang_groups, cleaned_data = get_guzhang_chunk_bbox_and_removed_data(data)
|
||
print(f"抽取的故障表总数量为：{len(guzhang_groups)}")
|
||
|
||
if not guzhang_groups:
|
||
print("未找到符合条件的组成表格，跳过处理。")
|
||
all_entity_results, all_relation_results = [], []
|
||
|
||
else:
|
||
all_entity_results, all_relation_results = [], []
|
||
|
||
for i, item in enumerate(guzhang_groups):
|
||
compose = item["group"]
|
||
table_type = item["table_type"]
|
||
|
||
try:
|
||
print(f"开始处理第 {i+1} 个故障表，类型: {table_type}")
|
||
content, highlight_list = get_content_bbox(compose)
|
||
print("故障表切片内容为：", content[:200])
|
||
|
||
# df_list = pd.read_html(StringIO(content), header=0)
|
||
# if not df_list:
|
||
# print(f" ⚠️ 第 {i+1} 组未解析到表格，跳过")
|
||
# continue
|
||
|
||
# df = df_list[0]
|
||
# df.dropna(how='all', inplace=True)
|
||
# result_list = json.loads(df.to_json(orient='records', force_ascii=False))
|
||
result_list = parse_table_content(content, chunk_size=10)
|
||
if not result_list:
|
||
continue
|
||
for ins in result_list:
|
||
print(ins)
|
||
print(1111111111111111111111)
|
||
highlight_filename, most_common_page, filename = highlight_pdf_file(input_pdf_path, highlight_list)
|
||
url = f"{prefix_url}/upload/{highlight_filename}#page={most_common_page}"
|
||
logger.info(f"处理完成，高亮文件：{highlight_filename}，高亮页码：{most_common_page}")
|
||
# ---------- type1: 故障现象 / 故障原因 / 维修项目 ----------
|
||
if table_type == "type1":
|
||
for ins in result_list:
|
||
# 安全获取字段，如果是 None 则转为空字符串
|
||
guzhang_xianxiang = ins.get("故障现象") or ""
|
||
guzhang_yuanyin = ins.get("故障原因") or ""
|
||
weixiu_xiangmu = ins.get("维修项目") or ""
|
||
|
||
# 只有当现象或原因不为空时才构建故障名，避免生成 "导致" 这种无意义字符串
|
||
if guzhang_xianxiang or guzhang_yuanyin:
|
||
guzhangname = f"{guzhang_yuanyin}导致{guzhang_xianxiang}"
|
||
else:
|
||
continue # 如果关键信息缺失，跳过该行
|
||
|
||
source_data_list = [{"filename": filename, "info": content, "url": url}]
|
||
|
||
all_entity_results.append({
|
||
"type": "故障模式",
|
||
"properties": {
|
||
"名称": guzhangname,
|
||
"故障现象": guzhang_xianxiang,
|
||
"故障原因": guzhang_yuanyin,
|
||
"维修方案": weixiu_xiangmu,
|
||
"knowledge_source": json.dumps(source_data_list, ensure_ascii=False)
|
||
}
|
||
})
|
||
all_relation_results.append({
|
||
"type": "包含故障",
|
||
"from_entity": f"{device}的维修工作",
|
||
"to_entity": guzhangname
|
||
})
|
||
|
||
# ---------- type2: 故障现象 / 可能原因 / 排除方法 ----------
|
||
elif table_type == "type2":
|
||
for ins in result_list:
|
||
guzhang_xianxiang = ins.get("故障现象") or ""
|
||
guzhang_yuanyin = ins.get("可能原因") or ""
|
||
weixiu_fangfa = ins.get("排除方法") or ""
|
||
|
||
if guzhang_xianxiang or guzhang_yuanyin:
|
||
guzhangname = f"{guzhang_yuanyin}导致{guzhang_xianxiang}"
|
||
else:
|
||
continue
|
||
|
||
source_data_list = [{"filename": filename, "info": content, "url": url}]
|
||
|
||
all_entity_results.append({
|
||
"type": "故障模式",
|
||
"properties": {
|
||
"名称": guzhangname,
|
||
"故障现象": guzhang_xianxiang,
|
||
"故障原因": guzhang_yuanyin,
|
||
"维修方案": weixiu_fangfa,
|
||
"knowledge_source": json.dumps(source_data_list, ensure_ascii=False)
|
||
}
|
||
})
|
||
all_relation_results.append({
|
||
"type": "包含故障",
|
||
"from_entity": f"{device}的维修工作",
|
||
"to_entity": guzhangname
|
||
})
|
||
|
||
elif table_type == "type3":
|
||
for ins in result_list:
|
||
guzhang_xianxiang = ins.get("故障外部现象") or ""
|
||
guzhang_yuanyin = ins.get("故障可能原因") or ""
|
||
caozuoyuan_zhidao = ins.get("操作员指导") or ""
|
||
weixiu_fangfa = ins.get("排除方法") or ""
|
||
|
||
if guzhang_xianxiang or guzhang_yuanyin:
|
||
guzhangname = f"{guzhang_yuanyin}导致{guzhang_xianxiang}"
|
||
else:
|
||
continue
|
||
|
||
source_data_list = [{"filename": filename, "info": content, "url": url}]
|
||
|
||
all_entity_results.append({
|
||
"type": "故障模式",
|
||
"properties": {
|
||
"名称": guzhangname,
|
||
"故障现象": guzhang_xianxiang,
|
||
"故障原因": guzhang_yuanyin,
|
||
"操作员指导": caozuoyuan_zhidao,
|
||
"维修方案": weixiu_fangfa,
|
||
"knowledge_source": json.dumps(source_data_list, ensure_ascii=False)
|
||
}
|
||
})
|
||
all_relation_results.append({
|
||
"type": "包含故障",
|
||
"from_entity": f"{device}的维修工作",
|
||
"to_entity": guzhangname
|
||
})
|
||
except Exception as e:
|
||
print(f"处理第 {i+1} 组时发生未预期错误: {e}")
|
||
continue
|
||
|
||
# 设备 → 维修工作指导 节点和关系（两种类型共用）
|
||
all_entity_results.append({
|
||
"type": "维修工作",
|
||
"properties": {"名称": f"{device}的维修工作"}
|
||
})
|
||
all_relation_results.append({
|
||
"type": "发生故障使用维修工作",
|
||
"from_entity": device,
|
||
"to_entity": f"{device}维修工作"
|
||
})
|
||
|
||
return all_entity_results, all_relation_results
|
||
|
||
|
||
|
||
"""
|
||
实体：设备、维修工作、故障、维修项目
|
||
关系：设备---发生故障使用维修工作--->维修工作
|
||
维修工作---包含--->故障
|
||
故障---被修复--->维修项目
|
||
"""
|
||
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/122-06A0014-B01001_发动机-维修手册.pdf"
|
||
output_json_path = '/storage01/home/hdf/project/wxkgrag/kgextract/output/output2/维修细目表抽取结果.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 = []
|
||
all_entity_results, all_relation_results = asyncio.run(
|
||
get_guzhang_node_relation(input_pdf_path, data=data)
|
||
)
|
||
for ins in all_entity_results:
|
||
print(ins)
|
||
length_entity = len(all_entity_results)
|
||
length_relation = len(all_relation_results)
|
||
|
||
print(f"抽取{length_entity}个实体，抽取{length_relation}个关系")
|
||
result = {
|
||
"entities": all_entity_results,
|
||
"relations": all_relation_results
|
||
}
|
||
|
||
# 保存为 JSON 文件
|
||
try:
|
||
with open(output_json_path, 'w', encoding='utf-8') as f:
|
||
json.dump(result, f, ensure_ascii=False, indent=4)
|
||
print(f"结果已保存至：{output_json_path}")
|
||
except Exception as e:
|
||
print(f"保存 JSON 文件时出错：{e}")
|
||
|