yixiong/FoodClassifier
3
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

对于类别中样本超过30个的采用K-Means算法截取前30个，存入向量数据库

Browse Source
This commit is contained in:
zhangpu
2026-03-03 17:07:06 +08:00
parent 8a50e2973b
commit dedf72f53c
4 changed files with 120 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files
faiss_vector_db/build_faiss_index.py
+113 -23
View File
@@ -317,8 +317,89 @@ class FAISSIndexBuilder:
print(f"特征向量已保存到: {embeddings_path}")
def build_complete_index(self, train_dir: str, output_dir: str,
batch_size: int = 32, index_type: str = 'flat'):
def select_representative_samples(
self,
embeddings: np.ndarray,
image_paths: List[str],
labels: List[int],
max_per_class: int = 30) -> Tuple[np.ndarray, List[str], List[int]]:
"""
使用 K-Means 为每个类别筛选代表性样本，消除样本数量不均衡带来的密度偏差。
原理：
- 样本数 <= max_per_class：全部保留
- 样本数 > max_per_class：对该类向量做 K-Means(k=max_per_class)
然后为每个聚类中心找最近的真实图片，只保留这些代表性样本
Args:
embeddings : 全量嵌入向量数组 (N, dim)
image_paths : 图片路径列表
labels : 标签索引列表
max_per_class: 每类最多保留的样本数（默认 30）
Returns:
Tuple[np.ndarray, List[str], List[int]]: 筛选后的 (embeddings, paths, labels)
"""
print(f"\n{'='*60}")
print(f"K-Means 代表性样本筛选 (max_per_class={max_per_class})")
print(f"{'='*60}")
embeddings_array = embeddings if isinstance(embeddings, np.ndarray) else np.array(embeddings)
labels_array = np.array(labels)
selected_global_indices: List[int] = []
for class_idx in np.unique(labels_array):
class_name = self.idx_to_class[int(class_idx)]
global_indices = np.where(labels_array == class_idx)[0]
count = len(global_indices)
if count <= max_per_class:
# 样本数不超过上限，全部保留
selected_global_indices.extend(global_indices.tolist())
print(f" [{class_name}]: {count} 张, 全部保留")
else:
# 样本数超过上限 -> K-Means 聚类后选代表
class_embs = embeddings_array[global_indices].astype(np.float32)
dim = class_embs.shape[1]
k = max_per_class
# 1. 用 FAISS 内置 K-Means 聚类（无需 sklearn）
kmeans = faiss.Kmeans(dim, k, niter=30, verbose=False)
kmeans.train(class_embs)
centroids = kmeans.centroids # shape: (k, dim)
# 2. 为每个幽灵中心，找离它最近的真实图片（幽灵变真身）
temp_index = faiss.IndexFlatL2(dim)
temp_index.add(class_embs)
_, nearest_local = temp_index.search(centroids, 1) # (k, 1)
nearest_local = nearest_local[:, 0].tolist()
# 3. 去重（极少情况下两个中心可能最近邻相同）
unique_local = list(set(nearest_local))
# 4. 映射回全局索引
selected_global = global_indices[unique_local].tolist()
selected_global_indices.extend(selected_global)
print(f" [{class_name}]: {count} 张 -> K-Means 筛选为 {len(selected_global)} 张")
# 按索引排序，确保顺序稳定
selected_global_indices = sorted(selected_global_indices)
filtered_embeddings = embeddings_array[selected_global_indices]
filtered_paths = [image_paths[i] for i in selected_global_indices]
filtered_labels = [labels[i] for i in selected_global_indices]
total_before = len(image_paths)
total_after = len(filtered_paths)
print(f"\nK-Means 筛选完成: {total_before} -> {total_after} 个向量"
f"（减少了 {total_before - total_after} 个）")
print(f"{'='*60}\n")
return filtered_embeddings, filtered_paths, filtered_labels
def build_complete_index(self, train_dir: str, output_dir: str,
batch_size: int = 32, index_type: str = 'flat',
max_per_class: int = 30):
"""
完整的索引构建流程
@@ -329,32 +410,38 @@ class FAISSIndexBuilder:
index_type: 索引类型
"""
print("=" * 60)
print("开始构建FAISS索引")
print("\u5f00始构建 FAISS 索引")
print("=" * 60)
# 1. 扫描训练数据
# 1. 扫描训练数据（全量，不做任何过滤）
image_paths, class_names, labels = self.scan_training_data(train_dir)
# 2. 提取全量特征向量（K-Means 聚类需要看到所有向量）
embeddings = self.extract_features_batch(image_paths, batch_size)
# 3. K-Means 代表性样本筛选（每类最多保留 max_per_class 个向量）
embeddings, image_paths, labels = self.select_representative_samples(
embeddings, image_paths, labels, max_per_class
)
self.image_paths = image_paths
self.labels = labels
# 2. 提取特征向量
embeddings = self.extract_features_batch(image_paths, batch_size)
self.embeddings = embeddings
# 3. 构建FAISS索引
# 4. 构建 FAISS 索引（使用筛选后的向量）
index = self.build_faiss_index(embeddings, index_type)
# 4. 保存索引和元数据
# 5. 保存索引和元数据
self.save_index_and_metadata(index, output_dir)
print("=" * 60)
print("FAISS索引构建完成!")
print(f"输出目录: {output_dir}")
print(f"总图片数: {len(image_paths)}")
print(f"类别数: {len(self.class_names)}")
print(f"特征维度: {self.embedding_dim}")
print("FAISS 索引构建完成!")
print(f"输出目录 : {output_dir}")
print(f"入库向量数 : {len(image_paths)}（K-Means 筛选后）")
print(f"类别数 : {len(self.class_names)}")
print(f"每类上限 : {max_per_class} 张")
print(f"特征维度 : {self.embedding_dim}")
print("=" * 60)
return index
@@ -564,7 +651,8 @@ def main():
#MODEL_PATH = "../model/WholeIngredientRecognition/cosface_20260121_133207/best_cosface_model.pth"
#grid_search最佳模型 s=68, m=0.42 1.22
# MODEL_PATH = "../model/WholeIngredientRecognition/grid_search_20260122_112545/best_model_s68.0_m0.42.pth"
MODEL_PATH = "../model/WholeIngredientRecognition/grid_search_20260302_092600/best_model_s64.0_m0.4.pth"
# MODEL_PATH = "../model/WholeIngredientRecognition/grid_search_20260302_092600/best_model_s64.0_m0.4.pth"
MODEL_PATH = "../model/WholeIngredientRecognition/grid_search_20260302_182102/best_model_s60.0_m0.42.pth"
# MODEL_PATH = "../model/DishClassification/grid_search_20251121_102723/model_s56.0_m0.4.pth"
# MODEL_PATH = "../model/DishClassification/cosface_20251111_153649/best_cosface_model.pth"
# TRAIN_DIR = "../dataset/ProcessedIngredientRecognition/train"
@@ -575,9 +663,10 @@ def main():
OUTPUT_DIR = "WholeIngredientRecognition/faiss_index"
# OUTPUT_DIR = "DishClassification/faiss_index"
BATCH_SIZE = 16
INDEX_TYPE = 'flat' # 'flat', 'ivf', 'hnsw'
INDEX_TYPE = 'flat' # 'flat', 'ivf', 'hnsw'
EMBEDDING_DIM = 512
MAX_PER_CLASS = 30 # 每类最多保留的代表性样本数（K-Means 筛选上限）
try:
# 构建索引
builder = FAISSIndexBuilder(MODEL_PATH, EMBEDDING_DIM)
@@ -585,7 +674,8 @@ def main():
train_dir=TRAIN_DIR,
output_dir=OUTPUT_DIR,
batch_size=BATCH_SIZE,
index_type=INDEX_TYPE
index_type=INDEX_TYPE,
max_per_class=MAX_PER_CLASS
)
print("\n" + "=" * 60)