Upload 2 files

graph_visualization.py

@@ -0,0 +1,61 @@
+import json
+from pyvis.network import Network
+
+
+def create_graph(nodes, edges, physics_enabled=True):
+    net = Network(
+        notebook=True,
+        height="100vh",
+        width="100vw",
+        bgcolor="#222222",
+        font_color="white",
+        cdn_resources="remote",
+    )
+
+    for node in nodes:
+        net.add_node(
+            node["id"],
+            label=node["label"],
+            title=node["label"],
+            color="blue" if node["label"] == "OOP" else "green",
+        )
+
+    for edge in edges:
+        net.add_edge(edge["source"], edge["target"], title=edge["type"])
+
+    net.force_atlas_2based(
+        gravity=-50,
+        central_gravity=0.01,
+        spring_length=100,
+        spring_strength=0.08,
+        damping=0.4,
+    )
+
+    options = {
+        "nodes": {"physics": physics_enabled},
+        "edges": {"smooth": True},
+        "interaction": {"hover": True, "zoomView": True},
+        "physics": {
+            "enabled": physics_enabled,
+            "stabilization": {"enabled": True, "iterations": 200},
+        },
+    }
+
+    net.set_options(json.dumps(options))
+    return net
+
+
+def visualize_graph(json_data, physics_enabled=True):
+    if isinstance(json_data, str):
+        data = json.loads(json_data)
+    else:
+        data = json_data
+    nodes = data["nodes"]
+    edges = data["edges"]
+    net = create_graph(nodes, edges, physics_enabled)
+    html = net.generate_html()
+    html = html.replace("'", '"')
+    html = html.replace(
+        '<div id="mynetwork"', '<div id="mynetwork" style="height: 100vh; width: 100%;"'
+    )
+    return f"""<iframe style="width: 100%; height: 100vh; border: none; margin: 0; padding: 0;" srcdoc='{html}'></iframe>"""

main.py

@@ -0,0 +1,84 @@
+import gradio as gr
+from datasets import load_from_disk
+from PIL import Image
+import io
+import base64
+import json
+
+from graph_visualization import visualize_graph
+
+# Load the dataset
+dataset = load_from_disk(
+    "/home/zaher/Projects/ZAKA_CAPSTONE_PROJECT/notebooks/OOP_KG_Dataset"
+)
+
+
+def reshape_json_data_to_fit_visualize_graph(graph_data):
+
+    nodes = graph_data["nodes"]
+    edges = graph_data["edges"]
+
+    transformed_nodes = [
+        {"id": nodes["id"][idx], "label": nodes["label"][idx]}
+        for idx in range(len(nodes["id"]))
+    ]
+
+    transformed_edges = [
+        {"source": edges["source"][idx], "target": edges["target"][idx], "type": "->"}
+        for idx in range(len(edges["source"]))
+    ]
+
+    # print(f"transformed nodes = {transformed_nodes}")
+
+    graph_data = {"nodes": transformed_nodes, "edges": transformed_edges}
+
+    return graph_data
+
+
+def display_example(index):
+    example = dataset[index]
+    # print("This is the example: ")
+    # print(example)
+    # Get the image
+    img = example["image"]
+
+    # Prepare the graph data
+    graph_data = {"nodes": example["nodes"], "edges": example["edges"]}
+
+    # # Convert graph_data to JSON string
+    # json_data = json.dumps(graph_data)
+    transformed_graph_data = reshape_json_data_to_fit_visualize_graph(graph_data)
+
+    # print(json_data)
+    # Generate the graph visualization
+    graph_html = visualize_graph(transformed_graph_data)
+
+    return img, graph_html
+
+
+def create_interface():
+    with gr.Blocks() as demo:
+        gr.Markdown("# Knowledge Graph Visualizer")
+
+        with gr.Row():
+            index_slider = gr.Slider(
+                minimum=0, maximum=len(dataset) - 1, step=1, label="Example Index"
+            )
+
+        with gr.Row():
+            image_output = gr.Image(type="pil", label="Image")
+            graph_output = gr.HTML(label="Knowledge Graph")
+
+        index_slider.change(
+            fn=display_example,
+            inputs=[index_slider],
+            outputs=[image_output, graph_output],
+        )
+
+    return demo
+
+
+# Create and launch the interface
+if __name__ == "__main__":
+    demo = create_interface()
+    demo.launch(debug=True)