增加注释,去掉dropout
This commit is contained in:
@@ -101,12 +101,12 @@ class EmbeddingFoodClassifierApp:
|
|||||||
self.faiss_index = faiss.read_index(index_path)
|
self.faiss_index = faiss.read_index(index_path)
|
||||||
print(f"FAISS索引已加载: {self.faiss_index.ntotal} 个向量")
|
print(f"FAISS索引已加载: {self.faiss_index.ntotal} 个向量")
|
||||||
|
|
||||||
# 加载图片路径映射
|
# 加载图片路径映射,就是一个list,一张张图片的路径
|
||||||
paths_path = os.path.join(index_dir, 'image_paths.pkl')
|
paths_path = os.path.join(index_dir, 'image_paths.pkl')
|
||||||
with open(paths_path, 'rb') as f:
|
with open(paths_path, 'rb') as f:
|
||||||
self.image_paths = pickle.load(f)
|
self.image_paths = pickle.load(f)
|
||||||
|
|
||||||
# 加载标签映射
|
# 加载标签映射,就是一个list,一张张图片的分类,不过是0,1,2,3,4这种
|
||||||
labels_path = os.path.join(index_dir, 'labels.pkl')
|
labels_path = os.path.join(index_dir, 'labels.pkl')
|
||||||
with open(labels_path, 'rb') as f:
|
with open(labels_path, 'rb') as f:
|
||||||
self.labels = pickle.load(f)
|
self.labels = pickle.load(f)
|
||||||
@@ -605,8 +605,11 @@ class EmbeddingFoodClassifierApp:
|
|||||||
if i < len(indices[0]) and indices[0][i] < len(self.labels):
|
if i < len(indices[0]) and indices[0][i] < len(self.labels):
|
||||||
idx = indices[0][i]
|
idx = indices[0][i]
|
||||||
score = scores[0][i]
|
score = scores[0][i]
|
||||||
|
# 根据FAISS索引拿到对应的分类id,0,1,2,3,4
|
||||||
class_idx = self.labels[idx]
|
class_idx = self.labels[idx]
|
||||||
|
# 根据分类id拿到对应的中文分类名称
|
||||||
class_name = self.class_names[class_idx]
|
class_name = self.class_names[class_idx]
|
||||||
|
# 根据FAISS的索引,拿到图片原始的路径
|
||||||
image_path = self.image_paths[idx]
|
image_path = self.image_paths[idx]
|
||||||
|
|
||||||
similar_classes.append(class_name)
|
similar_classes.append(class_name)
|
||||||
|
|||||||
@@ -42,11 +42,17 @@ class ResNet50EmbeddingNet(nn.Module):
|
|||||||
self.backbone.fc = nn.Identity()
|
self.backbone.fc = nn.Identity()
|
||||||
|
|
||||||
# 添加embedding层
|
# 添加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(
|
self.embedding_layer = nn.Sequential(
|
||||||
nn.Linear(backbone_output_dim, embedding_dim),
|
nn.Linear(backbone_output_dim, embedding_dim),
|
||||||
nn.BatchNorm1d(embedding_dim),
|
nn.BatchNorm1d(embedding_dim),
|
||||||
nn.ReLU(inplace=True),
|
nn.ReLU(inplace=True)
|
||||||
nn.Dropout(0.2)
|
|
||||||
)
|
)
|
||||||
|
|
||||||
# 图片预处理变换(用于推理)
|
# 图片预处理变换(用于推理)
|
||||||
|
|||||||
Reference in New Issue
Block a user