visualize.py

import networkx as nx
import rapidjson
import warnings
import os

from pyvis.network import Network

warnings.filterwarnings("ignore")

# Load the GraphML file
file_path = "./cache/graph_chunk_entity_relation.graphml"

assert os.path.exists(file_path), f"File {file_path} does not exist."
G = nx.read_graphml(file_path)

def create_graph(json_data):
    """
    Create interactive knowledge graph using pyvis
    """

    G = nx.Graph()

    # Add nodes with tooltips and error handling for missing keys
    for node in json_data['nodes']:
        # Get node type with fallback
        type = node.get("type", "Entity")

        # Get detailed type with fallback
        detailed_type = node.get("detailed_type", type)
        
        # Use node ID and type info for the tooltip
        G.add_node(node['id'], title=f"{type}: {detailed_type}")

    # Add edges with labels
    for edge in json_data['edges']:
        # Check if the required keys exist
        if 'from' in edge and 'to' in edge:
            label = edge.get('label', 'related')
            G.add_edge(edge['from'], edge['to'], title=label, label=label)

    # Create network visualization
    network = Network(
        width="100%",
        height="100vh",
        notebook=False,
        bgcolor="#f8fafc", 
        font_color="#1e293b"
    )
    
    # Configure network display
    network.from_nx(G)
    
    # Customize node appearance
    for node in network.nodes:
        node['color'] = {'background': '#e0e7ff', 'border': '#6366f1', 'highlight': {'background': '#c7d2fe', 'border': '#4f46e5'}}
        node['font'] = {'size': 14, 'color': '#1e293b'}
        node['shape'] = 'dot'
        node['size'] = 20
    
    # Customize edge appearance
    for edge in network.edges:
        edge['width'] = 4
        edge['color'] = {'color': '#6366f1', 'highlight': '#4f46e5'}
        edge['font'] = {'size': 12, 'color': '#4b5563', 'face': 'Arial'}
    
    # Save and display the network
    filename_out = "knowledge_graph.html"
    network.show(filename_out)
    print(f"Knowledge graph saved to {filename_out}")

# Convert the graph to node-link data format
js_graph = nx.node_link_data(G)
js_data = rapidjson.loads(rapidjson.dumps(js_graph))
# print(js_data)

create_graph(js_data)

# # Create a Pyvis network
# network = Network(width="100%",
#                   height="100vh",
#                   notebook=True,
#                   bgcolor="#f8fafc",
#                   font_color="#1e293b")

# # Convert NetworkX graph to Pyvis network
# network.from_nx(G)

# # Add colors and title to nodes
# for node in network.nodes:
#     if "description" in node:
#         node["title"] = node["description"]

#     node['color'] = {'background': '#e0e7ff', 'border': '#6366f1', 'highlight': {'background': '#c7d2fe', 'border': '#4f46e5'}}
#     node['font'] = {'size': 14, 'color': '#1e293b'}
#     node['shape'] = 'dot'
#     node['size'] = 20

# # Add title to edges
# for edge in network.edges:
#     if "description" in edge:
#         edge["title"] = edge["description"]

#     edge['width'] = 4
#     edge['color'] = {'color': '#6366f1', 'highlight': '#4f46e5'}
#     edge['font'] = {'size': 12, 'color': '#4b5563', 'face': 'Arial'}

# # Save and display the network
# filename_out = "knowledge_graph.html"
# network.show(filename_out)
# print(f"Knowledge graph saved to {filename_out}")