diff --git a/classifier/embedding_food_classifier_app.py b/classifier/embedding_food_classifier_app.py index 6c5538b..5d7e3d3 100644 --- a/classifier/embedding_food_classifier_app.py +++ b/classifier/embedding_food_classifier_app.py @@ -65,16 +65,16 @@ class EmbeddingFoodClassifierApp: BASE_DIR = os.path.dirname(os.path.abspath(__file__)) # model_path = "../model/embedding_20251011_133653/best_embedding_model.pth" # model_path = os.path.join(BASE_DIR, "../model/embedding_20251011_133653/best_embedding_model.pth") - model_path = os.path.join(BASE_DIR, "../model/ProcessedIngredientRecognition/embedding_20251029_170904/best_embedding_model.pth") + # 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/embedding_20251024_091151/best_embedding_model.pth") - # model_path = os.path.join(BASE_DIR, "../model/DishClassification/embedding_20251022_093635/best_embedding_model.pth") + model_path = os.path.join(BASE_DIR, "../model/DishClassification/cosface_20251105_200551/best_cosface_model.pth") # FAISS索引目录 # index_dir = "../faiss_vector_db/faiss_index" # index_dir = os.path.join(BASE_DIR, "../faiss_vector_db/faiss_index") - index_dir = os.path.join(BASE_DIR, "../faiss_vector_db/ProcessedIngredientRecognition/faiss_index") + # index_dir = os.path.join(BASE_DIR, "../faiss_vector_db/ProcessedIngredientRecognition/faiss_index") # index_dir = os.path.join(BASE_DIR, "../faiss_vector_db/WholeIngredientRecognition/faiss_index") - # index_dir = os.path.join(BASE_DIR, "../faiss_vector_db/DishClassification/faiss_index") + index_dir = os.path.join(BASE_DIR, "../faiss_vector_db/DishClassification/faiss_index") if os.path.exists(model_path) and os.path.exists(index_dir): # 1. 加载embedding模型 @@ -91,6 +91,10 @@ class EmbeddingFoodClassifierApp: if isinstance(checkpoint, dict): if 'model_state_dict' in checkpoint: self.model.load_state_dict(checkpoint['model_state_dict']) + elif 'backbone_state_dict' in checkpoint: + # CosFace格式:使用backbone_state_dict(只加载backbone部分) + self.model.load_state_dict(checkpoint['backbone_state_dict']) + print("检测到CosFace模型格式,使用'backbone_state_dict'加载") elif 'state_dict' in checkpoint: self.model.load_state_dict(checkpoint['state_dict']) else: diff --git a/faiss_vector_db/build_faiss_index.py b/faiss_vector_db/build_faiss_index.py index c105ca0..4b73f67 100644 --- a/faiss_vector_db/build_faiss_index.py +++ b/faiss_vector_db/build_faiss_index.py @@ -75,13 +75,25 @@ class FAISSIndexBuilder: # 处理不同的保存格式 if isinstance(checkpoint, dict): if 'model_state_dict' in checkpoint: + # Triplet格式:使用model_state_dict self.model.load_state_dict(checkpoint['model_state_dict']) + print("检测到Triplet模型格式,使用'model_state_dict'加载") + elif 'backbone_state_dict' in checkpoint: + # CosFace格式:使用backbone_state_dict(只加载backbone部分) + self.model.load_state_dict(checkpoint['backbone_state_dict']) + print("检测到CosFace模型格式,使用'backbone_state_dict'加载") elif 'state_dict' in checkpoint: + # 通用格式:使用state_dict self.model.load_state_dict(checkpoint['state_dict']) + print("检测到通用模型格式,使用'state_dict'加载") else: + # 直接加载整个checkpoint作为state_dict self.model.load_state_dict(checkpoint) + print("使用checkpoint直接加载") else: + # checkpoint本身就是state_dict self.model.load_state_dict(checkpoint) + print("checkpoint为state_dict格式,直接加载") # 移动到设备并设置为评估模式 self.model.to(self.device) @@ -408,6 +420,10 @@ class FAISSSearcher: if isinstance(checkpoint, dict): if 'model_state_dict' in checkpoint: self.model.load_state_dict(checkpoint['model_state_dict']) + elif 'backbone_state_dict' in checkpoint: + # CosFace格式:使用backbone_state_dict(只加载backbone部分) + self.model.load_state_dict(checkpoint['backbone_state_dict']) + print("检测到CosFace模型格式,使用'backbone_state_dict'加载") elif 'state_dict' in checkpoint: self.model.load_state_dict(checkpoint['state_dict']) else: @@ -480,16 +496,17 @@ def main(): """主函数""" # 配置参数 # MODEL_PATH = "../model/embedding_20251011_133653/best_embedding_model.pth" - MODEL_PATH = "../model/ProcessedIngredientRecognition/embedding_20251103_172012/best_embedding_model.pth" + # MODEL_PATH = "../model/ProcessedIngredientRecognition/embedding_20251103_172012/best_embedding_model.pth" # MODEL_PATH = "../model/WholeIngredientRecognition/embedding_20251030_135024/best_embedding_model.pth" - # MODEL_PATH = "../model/DishClassification/embedding_20251022_093635/best_embedding_model.pth" - TRAIN_DIR = "../dataset/ProcessedIngredientRecognition/train" + # MODEL_PATH = "../model/DishClassification/cosface_20251105_200551/best_embedding_model.pth" + MODEL_PATH = "../model/DishClassification/cosface_20251105_200551/best_cosface_model.pth" + # TRAIN_DIR = "../dataset/ProcessedIngredientRecognition/train" # TRAIN_DIR = "../dataset/WholeIngredientRecognition/train" - # TRAIN_DIR = "../dataset/DishClassification/train" + TRAIN_DIR = "../dataset/DishClassification/train" - OUTPUT_DIR = "ProcessedIngredientRecognition/faiss_index" + # OUTPUT_DIR = "ProcessedIngredientRecognition/faiss_index" # OUTPUT_DIR = "WholeIngredientRecognition/faiss_index" - # OUTPUT_DIR = "DishClassification/faiss_index" + OUTPUT_DIR = "DishClassification/faiss_index" BATCH_SIZE = 16 INDEX_TYPE = 'flat' # 'flat', 'ivf', 'hnsw' EMBEDDING_DIM = 512 diff --git a/faiss_vector_db/visualize_embeddings.py b/faiss_vector_db/visualize_embeddings.py index cab4355..aacfecc 100644 --- a/faiss_vector_db/visualize_embeddings.py +++ b/faiss_vector_db/visualize_embeddings.py @@ -162,12 +162,12 @@ def plot_2d(Z: np.ndarray, y: np.ndarray, title: str, out_path: Optional[str] = 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 路径") + parser.add_argument("--embeddings", type=str, default=os.path.join("DishClassification/faiss_index", "embeddings.json"), help="embeddings.json 路径") # parser.add_argument("--embeddings", type=str, default=os.path.join("WholeIngredientRecognition/faiss_index", "embeddings.json"), help="embeddings.json 路径") - parser.add_argument("--embeddings", type=str, default=os.path.join("ProcessedIngredientRecognition/faiss_index", "embeddings.json"), help="embeddings.json 路径") - # parser.add_argument("--labels", type=str, default=os.path.join("DishClassification/faiss_index", "labels.json"), help="labels.json 路径") + # parser.add_argument("--embeddings", type=str, default=os.path.join("ProcessedIngredientRecognition/faiss_index", "embeddings.json"), help="embeddings.json 路径") + parser.add_argument("--labels", type=str, default=os.path.join("DishClassification/faiss_index", "labels.json"), help="labels.json 路径") # parser.add_argument("--labels", type=str, default=os.path.join("WholeIngredientRecognition/faiss_index", "labels.json"), help="labels.json 路径") - parser.add_argument("--labels", type=str, default=os.path.join("ProcessedIngredientRecognition/faiss_index", "labels.json"), help="labels.json 路径") + # parser.add_argument("--labels", type=str, default=os.path.join("ProcessedIngredientRecognition/faiss_index", "labels.json"), help="labels.json 路径") parser.add_argument("--method", type=str, default="pca", choices=["pca", "tsne", "umap"], help="降维方法") parser.add_argument("--seed", type=int, default=42) parser.add_argument("--max_points", type=int, default=None, help="抽样上限,避免t-SNE/UMAP过慢;None为全量")