增加是否显示可视化向量的参数。

faiss_vector_db/visualize_embeddings.py

@@ -138,7 +138,17 @@ def reduce_dim(X: np.ndarray, method: str, seed: int, tsne_perplexity: int, umap
         raise ValueError(f"未知降维方法: {method}")
 
 
-def plot_2d(Z: np.ndarray, y: np.ndarray, title: str, out_path: Optional[str] = None) -> None:
+def plot_2d(Z: np.ndarray, y: np.ndarray, title: str, out_path: Optional[str] = None, show: bool = True) -> None:
+    """
+    绘制2D降维散点图
+    
+    Args:
+        Z: 降维后的2D坐标
+        y: 类别标签
+        title: 图表标题
+        out_path: 保存路径(可选)
+        show: 是否显示窗口(默认True,训练脚本中应设为False)
+    """
     plt.figure(figsize=(8, 7), dpi=120)
     classes = np.unique(y)
     # 构建调色板
@@ -157,7 +167,8 @@ def plot_2d(Z: np.ndarray, y: np.ndarray, title: str, out_path: Optional[str] =
     plt.tight_layout()
     if out_path:
         plt.savefig(out_path, bbox_inches="tight")
-    plt.show()
+    if show:
+        plt.show()
 
 
 def visualize_embeddings_from_files(embeddings_path: str,
@@ -227,9 +238,8 @@ def visualize_embeddings_from_files(embeddings_path: str,
     out_png = os.path.join(output_dir, f"embedding_{method}_2d.png")
     title = f"Embedding {method.upper()} 2D (N={Xs.shape[0]})"
     
-    # 为了在训练脚本中调用时不弹出窗口，我们需要关闭交互模式
+    # 训练脚本调用时不显示窗口，只保存图片
-    plt.ioff()  # 关闭交互模式
+    plot_2d(Z, ys, title, out_path=out_png, show=False)
-    plot_2d(Z, ys, title, out_path=out_png)
     plt.close('all')  # 关闭所有图形
     
     print(f"✓ 可视化图已保存: {out_png}")