在网格搜索中加入了可视化,但是训练脚本好像给改错了。
This commit is contained in:
@@ -358,6 +358,64 @@ class FAISSIndexBuilder:
|
||||
return index
|
||||
|
||||
|
||||
def extract_embeddings_only(model_path: str,
|
||||
train_dir: str,
|
||||
output_dir: str,
|
||||
embedding_dim: int = 512,
|
||||
batch_size: int = 16) -> Tuple[np.ndarray, List[int]]:
|
||||
"""
|
||||
仅提取特征向量并保存为JSON(用于可视化)
|
||||
不构建完整的FAISS索引,节省时间
|
||||
|
||||
Args:
|
||||
model_path: 模型路径
|
||||
train_dir: 训练数据目录
|
||||
output_dir: 输出目录
|
||||
embedding_dim: 特征向量维度
|
||||
batch_size: 批处理大小
|
||||
|
||||
Returns:
|
||||
(embeddings, labels): 特征向量数组和标签列表
|
||||
"""
|
||||
print("=" * 60)
|
||||
print("开始提取特征向量用于可视化")
|
||||
print("=" * 60)
|
||||
|
||||
builder = FAISSIndexBuilder(model_path, embedding_dim)
|
||||
|
||||
# 扫描数据
|
||||
image_paths, class_names, labels = builder.scan_training_data(train_dir)
|
||||
builder.image_paths = image_paths
|
||||
builder.labels = labels
|
||||
|
||||
# 提取特征
|
||||
embeddings = builder.extract_features_batch(image_paths, batch_size)
|
||||
builder.embeddings = embeddings
|
||||
|
||||
# 仅保存 embeddings.json 和 labels.json(可视化需要的)
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
embeddings_json_path = os.path.join(output_dir, 'embeddings.json')
|
||||
labels_json_path = os.path.join(output_dir, 'labels.json')
|
||||
|
||||
print(f"保存特征向量到: {embeddings_json_path}")
|
||||
with open(embeddings_json_path, "w", encoding="utf-8") as f:
|
||||
json.dump(embeddings.tolist(), f, ensure_ascii=False)
|
||||
|
||||
print(f"保存标签到: {labels_json_path}")
|
||||
with open(labels_json_path, "w", encoding="utf-8") as f:
|
||||
json.dump(labels, f, ensure_ascii=False)
|
||||
|
||||
print("=" * 60)
|
||||
print(f"✓ 特征向量提取完成")
|
||||
print(f" 输出目录: {output_dir}")
|
||||
print(f" 样本数: {len(embeddings)}")
|
||||
print(f" 特征维度: {embedding_dim}")
|
||||
print("=" * 60)
|
||||
|
||||
return embeddings, labels
|
||||
|
||||
|
||||
class FAISSSearcher:
|
||||
"""FAISS相似度检索器"""
|
||||
|
||||
|
||||
@@ -160,6 +160,84 @@ def plot_2d(Z: np.ndarray, y: np.ndarray, title: str, out_path: Optional[str] =
|
||||
plt.show()
|
||||
|
||||
|
||||
def visualize_embeddings_from_files(embeddings_path: str,
|
||||
labels_path: str,
|
||||
output_dir: str,
|
||||
method: str = "pca",
|
||||
max_points: Optional[int] = None,
|
||||
seed: int = 42,
|
||||
**kwargs) -> str:
|
||||
"""
|
||||
从文件加载并可视化embeddings(用于训练脚本调用)
|
||||
|
||||
Args:
|
||||
embeddings_path: embeddings.json路径
|
||||
labels_path: labels.json路径
|
||||
output_dir: 输出目录
|
||||
method: 降维方法 (pca/tsne/umap)
|
||||
max_points: 抽样上限
|
||||
seed: 随机种子
|
||||
**kwargs: 其他降维参数
|
||||
|
||||
Returns:
|
||||
输出的PNG图片路径
|
||||
"""
|
||||
print("=" * 60)
|
||||
print(f"开始可视化 embeddings")
|
||||
print(f"降维方法: {method}")
|
||||
print("=" * 60)
|
||||
|
||||
# 加载数据
|
||||
print(f"加载数据: {embeddings_path}")
|
||||
X = load_embeddings_json(embeddings_path)
|
||||
y = load_labels_json(labels_path)
|
||||
|
||||
# 数据对齐
|
||||
if X.shape[0] != y.shape[0]:
|
||||
n = min(X.shape[0], y.shape[0])
|
||||
print(f"警告: 数据不一致,截断到 {n}")
|
||||
X, y = X[:n], y[:n]
|
||||
|
||||
# 抽样
|
||||
Xs, ys, _ = subsample(X, y, max_points, seed=seed)
|
||||
if Xs.shape[0] < X.shape[0]:
|
||||
print(f"已抽样: {Xs.shape[0]}/{X.shape[0]}")
|
||||
|
||||
# 降维参数
|
||||
tsne_perplexity = kwargs.get('tsne_perplexity', 30)
|
||||
umap_n_neighbors = kwargs.get('umap_n_neighbors', 15)
|
||||
umap_min_dist = kwargs.get('umap_min_dist', 0.1)
|
||||
|
||||
# 降维
|
||||
print(f"执行降维: {method}")
|
||||
Z = reduce_dim(Xs, method, seed, tsne_perplexity, umap_n_neighbors, umap_min_dist)
|
||||
|
||||
# 诊断
|
||||
print("\n==== 诊断信息 ====")
|
||||
diag_info = diagnostics(Xs, ys, reduced2d=Z, method=method)
|
||||
print(diag_info)
|
||||
|
||||
# 保存诊断信息
|
||||
diag_path = os.path.join(output_dir, f"embedding_{method}_diagnostics.txt")
|
||||
with open(diag_path, 'w', encoding='utf-8') as f:
|
||||
f.write(diag_info)
|
||||
print(f"✓ 诊断信息已保存: {diag_path}")
|
||||
|
||||
# 绘图
|
||||
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)
|
||||
plt.close('all') # 关闭所有图形
|
||||
|
||||
print(f"✓ 可视化图已保存: {out_png}")
|
||||
print("=" * 60)
|
||||
|
||||
return out_png
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="可视化高维 embedding 并进行坍塌诊断")
|
||||
# parser.add_argument("--embeddings", type=str, default=os.path.join("DishClassification/faiss_index", "embeddings.json"), help="embeddings.json 路径")
|
||||
|
||||
Reference in New Issue
Block a user