Upload app.py

app.py

@@ -519,15 +519,30 @@ def generate_kg_image(entities, relations):
 
     # 创建网络图
     G = nx.DiGraph()
+
+    # 定义实体类型颜色
+    entity_colors = {
+        "Person": "red",
+        "Location": "green",
+        "Organization": "blue",
+        "Industry": "purple"
+    }
+
+    # 添加实体节点
     for entity in entities:
-        G.add_node(entity["text"], label=entity["type"])
+        if "text" in entity and "type" in entity:
+            G.add_node(entity["text"], label=entity["type"], color=entity_colors.get(entity["type"], "gray"))
+
+    # 添加关系边
     for relation in relations:
-        G.add_edge(relation["head"], relation["tail"], label=relation["relation"])
+        if "head" in relation and "tail" in relation and "relation" in relation:
+            G.add_edge(relation["head"], relation["tail"], label=relation["relation"])
 
     # 绘制网络图
-    pos = nx.spring_layout(G)
-    plt.figure(figsize=(10, 8))
-    nx.draw(G, pos, with_labels=True, node_color="lightblue")
+    pos = nx.spring_layout(G)  # 使用 spring 布局
+    plt.figure(figsize=(12, 8))
+    node_colors = [G.nodes[node].get("color", "gray") for node in G.nodes]
+    nx.draw(G, pos, with_labels=True, node_color=node_colors, font_size=10, font_weight="bold")
     nx.draw_networkx_edge_labels(G, pos, edge_labels={(u, v): d["label"] for u, v, d in G.edges(data=True)})
 
     # 确保保存目录存在