555 lines
20 KiB
Python
555 lines
20 KiB
Python
"""
|
|
食物分类器 - FAISS索引构建器
|
|
将训练集图片转换为特征向量并构建FAISS索引用于快速相似度检索
|
|
|
|
功能：
|
|
1. 加载训练好的ResNet50 Embedding模型
|
|
2. 遍历训练集所有图片，提取特征向量
|
|
3. 构建FAISS索引文件
|
|
4. 创建标签映射文件
|
|
5. 实现快速相似度检索功能
|
|
|
|
"""
|
|
|
|
import os
|
|
import json
|
|
import pickle
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image
|
|
from tqdm import tqdm
|
|
from datetime import datetime
|
|
from typing import List, Tuple, Dict, Optional
|
|
import faiss
|
|
|
|
# 导入项目模块
|
|
from net.resnet_embedding import create_resnet50_embedding
|
|
from settings import settings
|
|
|
|
|
|
class FAISSIndexBuilder:
|
|
"""FAISS索引构建器"""
|
|
|
|
def __init__(self, model_path: str, embedding_dim: int = 512):
|
|
"""
|
|
初始化FAISS索引构建器
|
|
|
|
Args:
|
|
model_path: 训练好的模型路径
|
|
embedding_dim: 特征向量维度
|
|
"""
|
|
self.model_path = model_path
|
|
self.embedding_dim = embedding_dim
|
|
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
|
|
|
# 初始化模型
|
|
self.model = None
|
|
self.load_model()
|
|
|
|
# 数据存储
|
|
self.embeddings = []
|
|
self.image_paths = []
|
|
self.labels = []
|
|
self.class_names = []
|
|
self.class_to_idx = {}
|
|
self.idx_to_class = {}
|
|
|
|
print(f"使用设备: {self.device}")
|
|
print(f"特征向量维度: {self.embedding_dim}")
|
|
|
|
def load_model(self):
|
|
"""加载训练好的ResNet50 Embedding模型"""
|
|
try:
|
|
print(f"正在加载模型: {self.model_path}")
|
|
|
|
# 创建模型
|
|
self.model = create_resnet50_embedding(
|
|
embedding_dim=self.embedding_dim,
|
|
pretrained=True,
|
|
use_internal_preprocess=False
|
|
)
|
|
|
|
# 加载权重
|
|
checkpoint = torch.load(self.model_path, map_location=self.device)
|
|
|
|
# 处理不同的保存格式
|
|
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)
|
|
self.model.eval()
|
|
|
|
print("模型加载成功!")
|
|
|
|
except Exception as e:
|
|
print(f"模型加载失败: {e}")
|
|
raise
|
|
|
|
def scan_training_data(self, train_dir: str) -> Tuple[List[str], List[str], List[int]]:
|
|
"""
|
|
扫描训练数据目录，获取所有图片路径和标签
|
|
|
|
Args:
|
|
train_dir: 训练数据目录路径
|
|
|
|
Returns:
|
|
Tuple[List[str], List[str], List[int]]: (图片路径列表, 类别名称列表, 标签索引列表)
|
|
"""
|
|
print(f"正在扫描训练数据目录: {train_dir}")
|
|
|
|
image_paths = []
|
|
class_names = []
|
|
labels = []
|
|
|
|
# 获取所有类别目录
|
|
class_dirs = [d for d in os.listdir(train_dir)
|
|
if os.path.isdir(os.path.join(train_dir, d))]
|
|
class_dirs.sort() # 确保顺序一致
|
|
|
|
print(f"发现类别: {class_dirs}")
|
|
|
|
# 创建类别映射
|
|
self.class_names = class_dirs
|
|
self.class_to_idx = {class_name: idx for idx, class_name in enumerate(class_dirs)}
|
|
self.idx_to_class = {idx: class_name for idx, class_name in enumerate(class_dirs)}
|
|
|
|
# 扫描每个类别目录
|
|
for class_name in class_dirs:
|
|
class_dir = os.path.join(train_dir, class_name)
|
|
class_idx = self.class_to_idx[class_name]
|
|
|
|
# 获取该类别下的所有图片
|
|
image_files = [f for f in os.listdir(class_dir)
|
|
if f.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp'))]
|
|
|
|
print(f"类别 '{class_name}': {len(image_files)} 张图片")
|
|
|
|
for image_file in image_files:
|
|
# 可以不添加那些增强的图片
|
|
if image_file.startswith("img"):
|
|
image_path = os.path.join(class_dir, image_file)
|
|
image_paths.append(image_path)
|
|
class_names.append(class_name)
|
|
labels.append(class_idx)
|
|
|
|
print(f"总计扫描到 {len(image_paths)} 张图片")
|
|
return image_paths, class_names, labels
|
|
|
|
def extract_features_batch(self, image_paths: List[str], batch_size: int = 32) -> np.ndarray:
|
|
"""
|
|
批量提取图片特征向量
|
|
|
|
Args:
|
|
image_paths: 图片路径列表
|
|
batch_size: 批处理大小
|
|
|
|
Returns:
|
|
np.ndarray: 特征向量数组，形状为 (N, embedding_dim)
|
|
"""
|
|
print(f"开始提取特征向量，批大小: {batch_size}")
|
|
|
|
all_embeddings = []
|
|
|
|
with torch.no_grad():
|
|
for i in tqdm(range(0, len(image_paths), batch_size), desc="提取特征"):
|
|
batch_paths = image_paths[i:i + batch_size]
|
|
batch_images = []
|
|
valid_indices = []
|
|
|
|
# 加载并预处理批次图片
|
|
for j, img_path in enumerate(batch_paths):
|
|
try:
|
|
# 加载图片
|
|
image = Image.open(img_path).convert('RGB')
|
|
|
|
# 预处理
|
|
tensor = self.model.preprocess_image(image)
|
|
batch_images.append(tensor.squeeze(0)) # 移除batch维度
|
|
valid_indices.append(i + j)
|
|
|
|
except Exception as e:
|
|
print(f"警告: 无法处理图片 {img_path}: {e}")
|
|
continue
|
|
|
|
if not batch_images:
|
|
continue
|
|
|
|
# 堆叠成批次并移动到设备
|
|
batch_tensor = torch.stack(batch_images).to(self.device)
|
|
|
|
# 提取特征
|
|
batch_embeddings = self.model(batch_tensor, normalize=True)
|
|
|
|
# 转换为numpy并添加到结果中
|
|
all_embeddings.append(batch_embeddings.cpu().numpy())
|
|
|
|
# 合并所有批次的结果
|
|
if all_embeddings:
|
|
embeddings = np.vstack(all_embeddings)
|
|
print(f"成功提取 {embeddings.shape[0]} 个特征向量")
|
|
return embeddings
|
|
else:
|
|
raise ValueError("没有成功提取任何特征向量")
|
|
|
|
def build_faiss_index(self, embeddings: np.ndarray, index_type: str = 'flat') -> faiss.Index:
|
|
"""
|
|
构建FAISS索引
|
|
|
|
Args:
|
|
embeddings: 特征向量数组
|
|
index_type: 索引类型 ('flat', 'ivf', 'hnsw')
|
|
|
|
Returns:
|
|
faiss.Index: 构建好的FAISS索引
|
|
"""
|
|
print(f"构建FAISS索引，类型: {index_type}")
|
|
|
|
dimension = embeddings.shape[1]
|
|
n_vectors = embeddings.shape[0]
|
|
|
|
if index_type == 'flat':
|
|
# 精确搜索索引
|
|
index = faiss.IndexFlatIP(dimension) # 内积索引（适合归一化向量）
|
|
|
|
elif index_type == 'ivf':
|
|
# IVF索引（适合大规模数据）
|
|
nlist = min(100, max(1, n_vectors // 10)) # 聚类中心数量
|
|
quantizer = faiss.IndexFlatIP(dimension)
|
|
index = faiss.IndexIVFFlat(quantizer, dimension, nlist)
|
|
|
|
# 训练索引
|
|
print("训练IVF索引...")
|
|
index.train(embeddings.astype(np.float32))
|
|
|
|
elif index_type == 'hnsw':
|
|
# HNSW索引（快速近似搜索）
|
|
M = 16 # 连接数
|
|
index = faiss.IndexHNSWFlat(dimension, M)
|
|
index.hnsw.efConstruction = 200
|
|
|
|
else:
|
|
raise ValueError(f"不支持的索引类型: {index_type}")
|
|
|
|
# 添加向量到索引
|
|
print("添加向量到索引...")
|
|
index.add(embeddings.astype(np.float32))
|
|
|
|
print(f"索引构建完成，包含 {index.ntotal} 个向量")
|
|
return index
|
|
|
|
def save_index_and_metadata(self, index: faiss.Index, output_dir: str):
|
|
"""
|
|
保存FAISS索引和相关元数据
|
|
|
|
Args:
|
|
index: FAISS索引
|
|
output_dir: 输出目录
|
|
"""
|
|
# 创建输出目录
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
|
|
# 保存FAISS索引
|
|
index_path = os.path.join(output_dir, 'faiss_index.bin')
|
|
faiss.write_index(index, index_path)
|
|
print(f"FAISS索引已保存到: {index_path}")
|
|
|
|
# 保存图片路径映射
|
|
paths_path = os.path.join(output_dir, 'image_paths.pkl')
|
|
with open(paths_path, 'wb') as f:
|
|
pickle.dump(self.image_paths, f)
|
|
print(f"图片路径映射已保存到: {paths_path}")
|
|
|
|
# 保存标签映射
|
|
labels_path = os.path.join(output_dir, 'labels.pkl')
|
|
labels_json_path = os.path.join(output_dir, 'labels.json')
|
|
with open(labels_path, 'wb') as f:
|
|
pickle.dump(self.labels, f)
|
|
with open(labels_json_path, "w", encoding="utf-8") as f:
|
|
json.dump(self.labels, f, ensure_ascii=False)
|
|
print(f"标签映射已保存到: {labels_path}")
|
|
|
|
# 保存类别信息
|
|
class_info = {
|
|
'class_names': self.class_names,
|
|
'class_to_idx': self.class_to_idx,
|
|
'idx_to_class': self.idx_to_class,
|
|
'embedding_dim': self.embedding_dim,
|
|
'total_images': len(self.image_paths),
|
|
'build_time': datetime.now().isoformat()
|
|
}
|
|
|
|
class_info_path = os.path.join(output_dir, 'class_info.json')
|
|
with open(class_info_path, 'w', encoding='utf-8') as f:
|
|
json.dump(class_info, f, ensure_ascii=False, indent=2)
|
|
print(f"类别信息已保存到: {class_info_path}")
|
|
|
|
# 保存特征向量（可选，用于调试）
|
|
embeddings_path = os.path.join(output_dir, 'embeddings.npy')
|
|
embeddings_json_path = os.path.join(output_dir, 'embeddings.json')
|
|
embeddings_array = np.array(self.embeddings)
|
|
np.save(embeddings_path, embeddings_array)
|
|
|
|
# 转成 Python 列表（JSON 不支持 numpy.array 直接序列化）
|
|
embeddings_list = embeddings_array.tolist()
|
|
# 保存成 JSON 文件
|
|
with open(embeddings_json_path, "w", encoding="utf-8") as f:
|
|
json.dump(embeddings_list, f, ensure_ascii=False)
|
|
|
|
print(f"特征向量已保存到: {embeddings_path}")
|
|
|
|
def build_complete_index(self, train_dir: str, output_dir: str,
|
|
batch_size: int = 32, index_type: str = 'flat'):
|
|
"""
|
|
完整的索引构建流程
|
|
|
|
Args:
|
|
train_dir: 训练数据目录
|
|
output_dir: 输出目录
|
|
batch_size: 批处理大小
|
|
index_type: 索引类型
|
|
"""
|
|
print("=" * 60)
|
|
print("开始构建FAISS索引")
|
|
print("=" * 60)
|
|
|
|
# 1. 扫描训练数据
|
|
image_paths, class_names, labels = self.scan_training_data(train_dir)
|
|
self.image_paths = image_paths
|
|
self.labels = labels
|
|
|
|
# 2. 提取特征向量
|
|
embeddings = self.extract_features_batch(image_paths, batch_size)
|
|
self.embeddings = embeddings
|
|
|
|
# 3. 构建FAISS索引
|
|
index = self.build_faiss_index(embeddings, index_type)
|
|
|
|
# 4. 保存索引和元数据
|
|
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("=" * 60)
|
|
|
|
return index
|
|
|
|
|
|
class FAISSSearcher:
|
|
"""FAISS相似度检索器"""
|
|
|
|
def __init__(self, index_dir: str, model_path: str):
|
|
"""
|
|
初始化FAISS检索器
|
|
|
|
Args:
|
|
index_dir: 索引文件目录
|
|
model_path: 模型路径
|
|
"""
|
|
self.index_dir = index_dir
|
|
self.model_path = model_path
|
|
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
|
|
|
# 加载索引和元数据
|
|
self.load_index_and_metadata()
|
|
|
|
# 加载模型
|
|
self.load_model()
|
|
|
|
def load_index_and_metadata(self):
|
|
"""加载FAISS索引和元数据"""
|
|
# 加载FAISS索引
|
|
index_path = os.path.join(self.index_dir, 'faiss_index.bin')
|
|
self.index = faiss.read_index(index_path)
|
|
print(f"FAISS索引已加载: {self.index.ntotal} 个向量")
|
|
|
|
# 加载图片路径
|
|
paths_path = os.path.join(self.index_dir, 'image_paths.pkl')
|
|
with open(paths_path, 'rb') as f:
|
|
self.image_paths = pickle.load(f)
|
|
|
|
# 加载标签
|
|
labels_path = os.path.join(self.index_dir, 'labels.pkl')
|
|
with open(labels_path, 'rb') as f:
|
|
self.labels = pickle.load(f)
|
|
|
|
# 加载类别信息
|
|
class_info_path = os.path.join(self.index_dir, 'class_info.json')
|
|
with open(class_info_path, 'r', encoding='utf-8') as f:
|
|
self.class_info = json.load(f)
|
|
|
|
self.class_names = self.class_info['class_names']
|
|
self.class_to_idx = self.class_info['class_to_idx']
|
|
self.idx_to_class = self.class_info['idx_to_class']
|
|
self.embedding_dim = self.class_info['embedding_dim']
|
|
|
|
print(f"元数据已加载: {len(self.image_paths)} 张图片, {len(self.class_names)} 个类别")
|
|
|
|
def load_model(self):
|
|
"""加载特征提取模型"""
|
|
self.model = create_resnet50_embedding(
|
|
embedding_dim=self.embedding_dim,
|
|
pretrained=True,
|
|
use_internal_preprocess=False
|
|
)
|
|
|
|
checkpoint = torch.load(self.model_path, map_location=self.device)
|
|
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:
|
|
self.model.load_state_dict(checkpoint)
|
|
else:
|
|
self.model.load_state_dict(checkpoint)
|
|
|
|
self.model.to(self.device)
|
|
self.model.eval()
|
|
print("特征提取模型已加载")
|
|
|
|
def search_similar(self, query_image_path: str, k: int = 5) -> List[Tuple[str, str, float]]:
|
|
"""
|
|
搜索相似图片
|
|
|
|
Args:
|
|
query_image_path: 查询图片路径
|
|
k: 返回最相似的k张图片
|
|
|
|
Returns:
|
|
List[Tuple[str, str, float]]: [(图片路径, 类别名称, 相似度分数), ...]
|
|
"""
|
|
# 提取查询图片的特征
|
|
query_image = Image.open(query_image_path).convert('RGB')
|
|
query_embedding = self.model.extract_embedding(query_image, normalize=True)
|
|
query_embedding = query_embedding.reshape(1, -1).astype(np.float32)
|
|
|
|
# 在FAISS索引中搜索
|
|
scores, indices = self.index.search(query_embedding, k)
|
|
|
|
# 整理结果
|
|
results = []
|
|
for i in range(k):
|
|
if i < len(indices[0]):
|
|
idx = indices[0][i]
|
|
score = scores[0][i]
|
|
image_path = self.image_paths[idx]
|
|
class_name = self.class_names[self.labels[idx]]
|
|
results.append((image_path, class_name, float(score)))
|
|
|
|
return results
|
|
|
|
def search_by_class(self, class_name: str, k: int = 10) -> List[Tuple[str, float]]:
|
|
"""
|
|
按类别搜索图片
|
|
|
|
Args:
|
|
class_name: 类别名称
|
|
k: 返回的图片数量
|
|
|
|
Returns:
|
|
List[Tuple[str, float]]: [(图片路径, 索引), ...]
|
|
"""
|
|
if class_name not in self.class_to_idx:
|
|
raise ValueError(f"未知类别: {class_name}")
|
|
|
|
class_idx = self.class_to_idx[class_name]
|
|
class_images = []
|
|
|
|
for i, label in enumerate(self.labels):
|
|
if label == class_idx:
|
|
class_images.append((self.image_paths[i], i))
|
|
if len(class_images) >= k:
|
|
break
|
|
|
|
return class_images
|
|
|
|
|
|
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/WholeIngredientRecognition/embedding_20251030_135024/best_embedding_model.pth"
|
|
# 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"
|
|
|
|
# OUTPUT_DIR = "ProcessedIngredientRecognition/faiss_index"
|
|
# OUTPUT_DIR = "WholeIngredientRecognition/faiss_index"
|
|
OUTPUT_DIR = "DishClassification/faiss_index"
|
|
BATCH_SIZE = 16
|
|
INDEX_TYPE = 'flat' # 'flat', 'ivf', 'hnsw'
|
|
EMBEDDING_DIM = 512
|
|
|
|
try:
|
|
# 构建索引
|
|
builder = FAISSIndexBuilder(MODEL_PATH, EMBEDDING_DIM)
|
|
index = builder.build_complete_index(
|
|
train_dir=TRAIN_DIR,
|
|
output_dir=OUTPUT_DIR,
|
|
batch_size=BATCH_SIZE,
|
|
index_type=INDEX_TYPE
|
|
)
|
|
|
|
print("\n" + "=" * 60)
|
|
print("测试相似度检索功能")
|
|
print("=" * 60)
|
|
|
|
# 测试检索功能
|
|
searcher = FAISSSearcher(OUTPUT_DIR, MODEL_PATH)
|
|
|
|
# 随机选择一张图片进行测试
|
|
test_image_path = builder.image_paths[0]
|
|
print(f"测试图片: {test_image_path}")
|
|
|
|
# 搜索相似图片
|
|
similar_images = searcher.search_similar(test_image_path, k=5)
|
|
|
|
print("\n最相似的5张图片:")
|
|
for i, (img_path, class_name, score) in enumerate(similar_images, 1):
|
|
print(f"{i}. {os.path.basename(img_path)} (类别: {class_name}, 相似度: {score:.4f})")
|
|
|
|
# 按类别搜索
|
|
print(f"\n'{builder.class_names[0]}'类别的前5张图片:")
|
|
class_images = searcher.search_by_class(builder.class_names[0], k=5)
|
|
for i, (img_path, idx) in enumerate(class_images, 1):
|
|
print(f"{i}. {os.path.basename(img_path)} (索引: {idx})")
|
|
|
|
except Exception as e:
|
|
print(f"错误: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |