Files
KG_generation/test.py
T
chentianrui 9609bb67b4 上传文件
2025-08-01 15:31:56 +08:00

240 lines
8.7 KiB
Python

import os
import logging
import shutil
import argparse
from project2json.project_converter import convert_project_to_json
from supplement_kg import costsummary_upwards
from equipment_calculation.main import process_BCL_calculate
from unified_structure import batch_write_BCLresult_into_json
from build_kg_ontolo import create_KGs_from_folder, connect_to_neo4j, read_config, clear_database
# 设置日志
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
logger = logging.getLogger(__name__)
def create_and_clean_folder(folder_path):
"""创建并清理文件夹"""
os.makedirs(folder_path, exist_ok=True)
# 清空文件夹,避免之前的文件干扰
for file in os.listdir(folder_path):
file_path = os.path.join(folder_path, file)
if os.path.isfile(file_path):
os.unlink(file_path)
elif os.path.isdir(file_path):
shutil.rmtree(file_path)
return folder_path
def step1_convert_project_to_json(input_folder, base_output_folder):
"""步骤1: 工程文件转JSON"""
logger.info("步骤1: 工程文件转JSON")
# 创建并清理输出文件夹
json_output_folder = os.path.join(base_output_folder, "json")
create_and_clean_folder(json_output_folder)
# 执行转换
result = convert_project_to_json(input_folder, json_output_folder)
if not result:
logger.error("工程文件转JSON失败")
return False, None
logger.info(f"工程文件转JSON完成,结果保存在 {json_output_folder}")
return True, json_output_folder
def step2_calculate_and_merge(input_folder, base_output_folder):
"""步骤2: 工程量费用预览计算和将BCL计算结果补充到JSON文件中"""
logger.info("步骤2: 工程量费用预览计算和将BCL计算结果补充到JSON文件中")
# 创建并清理输出文件夹
bcl_output_folder = os.path.join(base_output_folder, "bcl_results")
merged_output_folder = os.path.join(base_output_folder, "merged")
create_and_clean_folder(bcl_output_folder)
create_and_clean_folder(merged_output_folder)
# 确保输入文件夹存在
if not os.path.exists(input_folder):
logger.error(f"输入文件夹 {input_folder} 不存在")
return False, None
# 2.1 工程量费用预览计算和人材机合价计算
logger.info("2.1: 工程量费用预览计算和人材机合价计算")
results = process_BCL_calculate(input_folder, bcl_output_folder)
if not results:
logger.error("工程量费用预览计算和人材机合价计算失败")
return False, None
logger.info(f"工程量费用预览计算和人材机合价计算完成,结果保存在 {bcl_output_folder}")
# 2.2 将BCL计算结果补充到JSON文件中
logger.info("2.2: 将BCL计算结果补充到JSON文件中")
count = batch_write_BCLresult_into_json(input_folder, bcl_output_folder, merged_output_folder)
if count == 0:
logger.error("将BCL计算结果补充到JSON文件中失败")
return False, None
logger.info(f"BCL计算结果已补充到JSON文件中,成功处理了 {count} 个文件,结果保存在 {merged_output_folder}")
return True, merged_output_folder
def step3_summarize_costs(input_folder, base_output_folder):
"""步骤3: 文件内项目划分和清单节点费用预览向上汇总"""
logger.info("步骤3: 文件内项目划分和清单节点费用预览向上汇总")
# 创建并清理输出文件夹
final_output_folder = os.path.join(base_output_folder, "final")
create_and_clean_folder(final_output_folder)
# 确保输入文件夹存在
if not os.path.exists(input_folder):
logger.error(f"输入文件夹 {input_folder} 不存在")
return False, None
# 执行费用向上汇总
result = costsummary_upwards(input_folder, final_output_folder)
if not result:
logger.error("文件内项目划分和清单节点费用预览向上汇总失败")
return False, None
logger.info(f"文件内项目划分和清单节点费用预览向上汇总完成,结果保存在 {final_output_folder}")
return True, final_output_folder
def step4_create_knowledge_graph(input_folder):
"""步骤4: 将JSON文件转换为知识图谱"""
logger.info("步骤4: 将JSON文件转换为知识图谱")
# 确保输入文件夹存在
if not os.path.exists(input_folder):
logger.error(f"输入文件夹 {input_folder} 不存在")
return False
# 读取配置文件
config = read_config()
# 连接到Neo4j数据库
if not connect_to_neo4j(config["neo4j"]["uri"], config["neo4j"]["user"], config["neo4j"]["password"]):
logger.error("连接Neo4j数据库失败")
return False
# 清空数据库
if not clear_database():
logger.error("清空Neo4j数据库失败")
return False
# 处理输入文件夹中的所有JSON文件
success_count, total_count = create_KGs_from_folder(input_folder)
if success_count == 0:
logger.error("将JSON文件转换为知识图谱失败")
return False
logger.info(f"知识图谱创建完成,成功处理 {success_count}/{total_count} 个文件")
return True
def parse_args():
"""解析命令行参数"""
parser = argparse.ArgumentParser(description="工程文件处理和知识图谱生成工具")
parser.add_argument(
"--step",
choices=["all", "1", "2", "3", "4"],
default="all",
help="选择执行的步骤: 1=工程文件转JSON, 2=BCL计算和合并, 3=费用向上汇总, 4=创建知识图谱, all=执行所有步骤",
)
parser.add_argument("--input", default="project2json/uploads", help="输入文件夹路径,包含工程文件")
parser.add_argument("--output", default="project2json/outputs", help="输出文件夹路径")
return parser.parse_args()
def main():
"""主函数"""
# 解析命令行参数
args = parse_args()
# 设置输入输出路径
input_folder = args.input
base_output_folder = args.output
# 确保基础输出文件夹存在
os.makedirs(base_output_folder, exist_ok=True)
# 根据选择的步骤执行相应的功能
if args.step == "all":
# 执行所有步骤
success, json_output_folder = step1_convert_project_to_json(input_folder, base_output_folder)
if not success:
return
success, merged_output_folder = step2_calculate_and_merge(json_output_folder, base_output_folder)
if not success:
return
success, final_output_folder = step3_summarize_costs(merged_output_folder, base_output_folder)
if not success:
return
if not step4_create_knowledge_graph(final_output_folder):
return
elif args.step == "1":
# 只执行步骤1
step1_convert_project_to_json(input_folder, base_output_folder)
elif args.step == "2":
# 只执行步骤2
# 使用json文件夹作为输入
json_input_folder = os.path.join(base_output_folder, "json")
if not os.path.exists(json_input_folder) or not os.listdir(json_input_folder):
logger.error(f"错误: {json_input_folder} 文件夹不存在或为空,请先执行步骤1")
return
step2_calculate_and_merge(json_input_folder, base_output_folder)
elif args.step == "3":
# 只执行步骤3
# 使用merged文件夹作为输入
merged_input_folder = os.path.join(base_output_folder, "merged")
if not os.path.exists(merged_input_folder) or not os.listdir(merged_input_folder):
logger.error(f"错误: {merged_input_folder} 文件夹不存在或为空,请先执行步骤2")
return
step3_summarize_costs(merged_input_folder, base_output_folder)
elif args.step == "4":
# 只执行步骤4
# 使用final文件夹作为输入
final_input_folder = os.path.join(base_output_folder, "final")
if not os.path.exists(final_input_folder) or not os.listdir(final_input_folder):
logger.error(f"错误: {final_input_folder} 文件夹不存在或为空,请先执行步骤3")
return
step4_create_knowledge_graph(final_input_folder)
logger.info("所有选定的步骤已完成")
if __name__ == "__main__":
main()
"""
# 执行所有步骤
python test.py --step all --input project2json/uploads --output project2json/outputs
# 只执行步骤1
python test.py --step 1 --input project2json/uploads --output project2json/outputs
# 只执行步骤2
python test.py --step 2 --input project2json/uploads --output project2json/outputs
# 只执行步骤3
python test.py --step 3 --input project2json/uploads --output project2json/outputs
# 只执行步骤4
python test.py --step 4 --output project2json/outputs
"""