diff --git a/faiss_vector_db/visualize_embeddings.py b/faiss_vector_db/visualize_embeddings.py index 5ddf248..42f9a10 100644 --- a/faiss_vector_db/visualize_embeddings.py +++ b/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) + # 训练脚本调用时不显示窗口,只保存图片 + plot_2d(Z, ys, title, out_path=out_png, show=False) plt.close('all') # 关闭所有图形 print(f"✓ 可视化图已保存: {out_png}")