From 5504e735a19d0c89abfda4042e4dae9335b4d2ba Mon Sep 17 00:00:00 2001 From: zhangpu <1250681871@qq.com> Date: Thu, 4 Dec 2025 14:02:34 +0800 Subject: [PATCH] =?UTF-8?q?=E5=A2=9E=E5=8A=A0=E6=98=AF=E5=90=A6=E6=98=BE?= =?UTF-8?q?=E7=A4=BA=E5=8F=AF=E8=A7=86=E5=8C=96=E5=90=91=E9=87=8F=E7=9A=84?= =?UTF-8?q?=E5=8F=82=E6=95=B0=E3=80=82?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- faiss_vector_db/visualize_embeddings.py | 20 +++++++++++++++----- 1 file changed, 15 insertions(+), 5 deletions(-) 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}")