from pyvis.network import Network
import json
from selenium import webdriver
import time

# 读取您的JSON知识图谱数据
with open(r'D:\dsWork\dsProject\dsLightRag\Doc\史校长资料\技术一-知识图谱源文件\middle_school_math_graph.json',encoding='utf-8') as f:
    data = json.load(f)

# 创建网络图
net = Network(height='800px', width='100%', notebook=False)

# 添加节点
for node in data['nodes']:
    net.add_node(node['id'], label=node['id'][:20]+'...' if len(node['id'])>20 else node['id'])

# 添加边（如果有的话）
if 'edges' in data:
    for edge in data['edges']:
        net.add_edge(edge['source'], edge['target'])

# 保存为HTML文件
html_path = 'knowledge_graph_json.html'
net.show(html_path)

# 使用selenium截图保存为PNG
options = webdriver.ChromeOptions()
options.add_argument('--headless')
options.add_argument('--disable-gpu')
options.add_argument('--window-size=1200,900')

driver = webdriver.Chrome(options=options)
driver.get(f'file:///{os.path.abspath(html_path)}')
time.sleep(3)  # 等待页面加载

driver.save_screenshot('knowledge_graph_json.png')
driver.quit()