diff --git a/classifier/embedding_food_classifier_app.py b/classifier/embedding_food_classifier_app.py index 0a7a20b..6900ea2 100644 --- a/classifier/embedding_food_classifier_app.py +++ b/classifier/embedding_food_classifier_app.py @@ -101,12 +101,12 @@ class EmbeddingFoodClassifierApp: self.faiss_index = faiss.read_index(index_path) print(f"FAISS索引已加载: {self.faiss_index.ntotal} 个向量") - # 加载图片路径映射 + # 加载图片路径映射,就是一个list,一张张图片的路径 paths_path = os.path.join(index_dir, 'image_paths.pkl') with open(paths_path, 'rb') as f: self.image_paths = pickle.load(f) - # 加载标签映射 + # 加载标签映射,就是一个list,一张张图片的分类,不过是0,1,2,3,4这种 labels_path = os.path.join(index_dir, 'labels.pkl') with open(labels_path, 'rb') as f: self.labels = pickle.load(f) @@ -605,8 +605,11 @@ class EmbeddingFoodClassifierApp: if i < len(indices[0]) and indices[0][i] < len(self.labels): idx = indices[0][i] score = scores[0][i] + # 根据FAISS索引拿到对应的分类id,0,1,2,3,4 class_idx = self.labels[idx] + # 根据分类id拿到对应的中文分类名称 class_name = self.class_names[class_idx] + # 根据FAISS的索引,拿到图片原始的路径 image_path = self.image_paths[idx] similar_classes.append(class_name) diff --git a/net/resnet_embedding.py b/net/resnet_embedding.py index 46c9752..af0b938 100644 --- a/net/resnet_embedding.py +++ b/net/resnet_embedding.py @@ -42,11 +42,17 @@ class ResNet50EmbeddingNet(nn.Module): self.backbone.fc = nn.Identity() # 添加embedding层 + # self.embedding_layer = nn.Sequential( + # nn.Linear(backbone_output_dim, embedding_dim), + # nn.BatchNorm1d(embedding_dim), + # nn.ReLU(inplace=True), + # nn.Dropout(0.2) + # ) + # 暂时不用加Dropout self.embedding_layer = nn.Sequential( nn.Linear(backbone_output_dim, embedding_dim), nn.BatchNorm1d(embedding_dim), - nn.ReLU(inplace=True), - nn.Dropout(0.2) + nn.ReLU(inplace=True) ) # 图片预处理变换(用于推理)