Upload app.py

app.py

@@ -508,34 +508,38 @@ def process_text(text, model_type="bert"):
     return ent_text, rel_text, kg_text, f"{total_duration:.2f} 秒"
 
 
-from matplotlib.figure import Figure
-from io import BytesIO
+import os
 
 def generate_kg_image(entities, relations):
     """
-    生成知识图谱的图片
+    生成知识图谱的图片并保存到文件
     """
-    fig = Figure(figsize=(10, 8))
-    ax = fig.add_subplot(1, 1, 1)
-    ax.axis("off")
-    
+    import matplotlib.pyplot as plt
+    import networkx as nx
+
     # 创建网络图
     G = nx.DiGraph()
     for entity in entities:
         G.add_node(entity["text"], label=entity["type"])
     for relation in relations:
         G.add_edge(relation["head"], relation["tail"], label=relation["relation"])
-    
+
     # 绘制网络图
     pos = nx.spring_layout(G)
-    nx.draw(G, pos, with_labels=True, node_color="lightblue", ax=ax)
-    nx.draw_networkx_edge_labels(G, pos, edge_labels={(u, v): d["label"] for u, v, d in G.edges(data=True)}, ax=ax)
-    
-    # 保存为图片
-    buf = BytesIO()
-    fig.savefig(buf, format="png")
-    buf.seek(0)
-    return buf
+    plt.figure(figsize=(10, 8))
+    nx.draw(G, pos, with_labels=True, node_color="lightblue")
+    nx.draw_networkx_edge_labels(G, pos, edge_labels={(u, v): d["label"] for u, v, d in G.edges(data=True)})
+
+    # 确保保存目录存在
+    save_dir = "D:/image"
+    if not os.path.exists(save_dir):
+        os.makedirs(save_dir)  # 自动创建文件夹
+
+    # 保存图片到指定目录
+    file_path = os.path.join(save_dir, f"knowledge_graph_{int(time.time())}.png")
+    plt.savefig(file_path)
+    plt.close()
+    return file_path
 
 def process_file(file, model_type="bert"):
     try:
@@ -566,9 +570,9 @@ def process_file(file, model_type="bert"):
         # 生成知识图谱图片
         entities, _ = ner(text, model_type)
         relations, _ = re_extract(entities, text)
-        kg_image = generate_kg_image(entities, relations)
+        kg_image_path = generate_kg_image(entities, relations)  # 返回文件路径
         
-        return ent_text, rel_text, kg_text, duration, kg_image
+        return ent_text, rel_text, kg_text, duration, kg_image_path
         
     except Exception as e:
         logging.error(f"文件处理错误: {str(e)}")