增加FAISS向量数据库对所有训练图片进行向量化的过程。
This commit is contained in:
zhanghuan
@@ -0,0 +1,517 @@
|
||||
"""
|
||||
食物分类器 - 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:
|
||||
self.model.load_state_dict(checkpoint['model_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("模型加载成功!")
|
||||
|
||||
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:
|
||||
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')
|
||||
with open(labels_path, 'wb') as f:
|
||||
pickle.dump(self.labels, f)
|
||||
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_array = np.array(self.embeddings)
|
||||
np.save(embeddings_path, embeddings_array)
|
||||
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 '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_20250917_145342/best_embedding_model.pth"
|
||||
TRAIN_DIR = "../dataset/train"
|
||||
OUTPUT_DIR = "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()
|
||||
Reference in New Issue
Block a user