From dedf72f53c28d85149c7f4eeceec601a601d7385 Mon Sep 17 00:00:00 2001 From: zhangpu <1250681871@qq.com> Date: Tue, 3 Mar 2026 17:07:06 +0800 Subject: [PATCH] =?UTF-8?q?=E5=AF=B9=E4=BA=8E=E7=B1=BB=E5=88=AB=E4=B8=AD?= =?UTF-8?q?=E6=A0=B7=E6=9C=AC=E8=B6=85=E8=BF=8730=E4=B8=AA=E7=9A=84?= =?UTF-8?q?=E9=87=87=E7=94=A8K-Means=E7=AE=97=E6=B3=95=E6=88=AA=E5=8F=96?= =?UTF-8?q?=E5=89=8D30=E4=B8=AA=EF=BC=8C=E5=AD=98=E5=85=A5=E5=90=91?= =?UTF-8?q?=E9=87=8F=E6=95=B0=E6=8D=AE=E5=BA=93?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- classifier/embedding_food_classifier_app.py | 2 +- faiss_vector_db/build_faiss_index.py | 136 ++++++++++++++++---- toAndroid/toAndroidEmbedding.py | 10 +- train/grid_search_cosface.py | 4 +- 4 files changed, 120 insertions(+), 32 deletions(-) diff --git a/classifier/embedding_food_classifier_app.py b/classifier/embedding_food_classifier_app.py index 607254c..82fb85e 100644 --- a/classifier/embedding_food_classifier_app.py +++ b/classifier/embedding_food_classifier_app.py @@ -71,7 +71,7 @@ class EmbeddingFoodClassifierApp: # model_path = os.path.join(BASE_DIR, "../model/ProcessedIngredientRecognition/embedding_20251029_170904/best_embedding_model.pth") # model_path = os.path.join(BASE_DIR, "../model/WholeIngredientRecognition/cosface_20251106_134718/best_cosface_model.pth") # model_path = os.path.join(BASE_DIR, "../model/WholeIngredientRecognition/grid_search_20251204_140816/model_s64.0_m0.4.pth") - model_path = os.path.join(BASE_DIR, "../model/WholeIngredientRecognition/grid_search_20260302_092600/best_model_s64.0_m0.4.pth") + model_path = os.path.join(BASE_DIR, "../model/WholeIngredientRecognition/grid_search_20260302_182102/best_model_s60.0_m0.42.pth") # model_path = os.path.join(BASE_DIR, "../model/DishClassification/cosface_20251111_153649/best_cosface_model.pth") # FAISS索引目录 diff --git a/faiss_vector_db/build_faiss_index.py b/faiss_vector_db/build_faiss_index.py index 607d0ab..b6d00cc 100644 --- a/faiss_vector_db/build_faiss_index.py +++ b/faiss_vector_db/build_faiss_index.py @@ -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) diff --git a/toAndroid/toAndroidEmbedding.py b/toAndroid/toAndroidEmbedding.py index 43195f8..f025c1e 100644 --- a/toAndroid/toAndroidEmbedding.py +++ b/toAndroid/toAndroidEmbedding.py @@ -19,10 +19,8 @@ def main(): base_model = create_mobile_resnet50_embedding(embedding_dim=512, pretrained=True) # model_path = "../model/DishClassification/grid_search_20251121_102723/model_s56.0_m0.4.pth" # model_path = "../model/WholeIngredientRecognition/cosface_20251106_134718/best_cosface_model.pth" - #原本路径,原本模型 - #model_path = "../model/WholeIngredientRecognition/grid_search_20251204_140816/model_s64.0_m0.4.pth" - #新路径,新模型 (模型不可用) - model_path = "../model/WholeIngredientRecognition/cosface_20260121_133207/best_cosface_model.pth" + model_path = "../model/WholeIngredientRecognition/grid_search_20260302_182102/best_model_s60.0_m0.42.pth" + # model_path = "../model/WholeIngredientRecognition/cosface_20260121_133207/best_cosface_model.pth" # model_path = "../model/ProcessedIngredientRecognition/embedding_20251029_173607/best_embedding_model.pth" if not os.path.exists(model_path): @@ -98,9 +96,9 @@ def main(): # 保存模型,一定不要乱改,防止模型覆盖!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! # output_path = "../model/DishClassification/grid_search_20251121_102723/best_embedding_model_mobile.pt" #原本路径 - #output_path = "../model/WholeIngredientRecognition/grid_search_20251204_140816/best_embedding_model_mobile.pt" + output_path = "../model/WholeIngredientRecognition/grid_search_20260302_182102/best_embedding_model_mobile.pt" #新路径 - output_path = "../model/WholeIngredientRecognition/cosface_20260121_133207/best_cosface_model_mobile.pt" + # output_path = "../model/WholeIngredientRecognition/cosface_20260121_133207/best_cosface_model_mobile.pt" # output_path = "../model/WholeIngredientRecognition/cosface_20251106_134718/best_embedding_model_mobile.pt" # output_path = "../model/ProcessedIngredientRecognition/embedding_20251029_173607/best_embedding_model_mobile.pt" traced_model.save(output_path) diff --git a/train/grid_search_cosface.py b/train/grid_search_cosface.py index a85c6b4..6e7d18a 100644 --- a/train/grid_search_cosface.py +++ b/train/grid_search_cosface.py @@ -57,8 +57,8 @@ GRID_PARAMS = { 'm': [0.32, 0.35, 0.38, 0.40,0.45,0.50], # margin参数 }, 'whole_ingredient': { - 's': [64.0], - 'm': [0.40, 0.42, 0.43, 0.45], + 's': [60.0,64.0], + 'm': [0.42, 0.45], }, 'processed_ingredient': { 's': [56.0, 60.0, 64.0, 68.0],