增加FAISS向量数据库对所有训练图片进行向量化的过程。
This commit is contained in:
zhanghuan
@@ -0,0 +1,244 @@
|
||||
"""
|
||||
FAISS相似度检索系统快速演示脚本
|
||||
一键运行完整的索引构建和检索演示
|
||||
"""
|
||||
|
||||
import os
|
||||
import time
|
||||
|
||||
|
||||
def check_dependencies():
|
||||
"""检查依赖包"""
|
||||
required_packages = [
|
||||
'torch', 'torchvision', 'faiss', 'PIL', 'numpy', 'tqdm', 'matplotlib'
|
||||
]
|
||||
|
||||
missing_packages = []
|
||||
|
||||
for package in required_packages:
|
||||
try:
|
||||
if package == 'PIL':
|
||||
import PIL
|
||||
elif package == 'faiss':
|
||||
import faiss
|
||||
else:
|
||||
__import__(package)
|
||||
except ImportError:
|
||||
missing_packages.append(package)
|
||||
|
||||
if missing_packages:
|
||||
print("缺少以下依赖包:")
|
||||
for pkg in missing_packages:
|
||||
if pkg == 'faiss':
|
||||
print(f" {pkg} (安装命令: pip install faiss-cpu)")
|
||||
else:
|
||||
print(f" {pkg}")
|
||||
print("\n请先安装缺少的依赖包")
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def check_files():
|
||||
"""检查必要文件"""
|
||||
required_files = [
|
||||
'net/resnet_embedding.py',
|
||||
'settings/settings.py',
|
||||
'model/embedding_20250917_145342/best_embedding_model.pth',
|
||||
'dataset/train'
|
||||
]
|
||||
|
||||
missing_files = []
|
||||
|
||||
for file_path in required_files:
|
||||
if not os.path.exists(file_path):
|
||||
missing_files.append(file_path)
|
||||
|
||||
if missing_files:
|
||||
print("缺少以下必要文件:")
|
||||
for file_path in missing_files:
|
||||
print(f" {file_path}")
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def run_index_building():
|
||||
"""运行索引构建"""
|
||||
print("=" * 60)
|
||||
print("步骤 1: 构建FAISS索引")
|
||||
print("=" * 60)
|
||||
|
||||
try:
|
||||
from faiss_vector_db.build_faiss_index import main as build_main
|
||||
build_main()
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"索引构建失败: {e}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
return False
|
||||
|
||||
def run_search_demo():
|
||||
"""运行检索演示"""
|
||||
print("\n" + "=" * 60)
|
||||
print("步骤 2: 相似度检索演示")
|
||||
print("=" * 60)
|
||||
|
||||
try:
|
||||
from faiss_vector_db.build_faiss_index import FAISSSearcher
|
||||
|
||||
# 配置参数
|
||||
index_dir = "faiss_index"
|
||||
model_path = "../model/embedding_20250917_145342/best_embedding_model.pth"
|
||||
|
||||
# 检查索引是否存在
|
||||
if not os.path.exists(os.path.join(index_dir, 'faiss_index.bin')):
|
||||
print("错误: FAISS索引不存在，请先构建索引")
|
||||
return False
|
||||
|
||||
# 创建检索器
|
||||
searcher = FAISSSearcher(index_dir, model_path)
|
||||
|
||||
# 显示索引信息
|
||||
print(f"索引信息:")
|
||||
print(f" 总向量数: {searcher.index.ntotal}")
|
||||
print(f" 特征维度: {searcher.embedding_dim}")
|
||||
print(f" 类别数: {len(searcher.class_names)}")
|
||||
print(f" 类别: {', '.join(searcher.class_names)}")
|
||||
|
||||
# 找一些测试图片
|
||||
test_images = []
|
||||
test_dirs = ['dataset/test', 'dataset/val', 'dataset/train']
|
||||
|
||||
for test_dir in test_dirs:
|
||||
if os.path.exists(test_dir):
|
||||
for root, dirs, files in os.walk(test_dir):
|
||||
for file in files:
|
||||
if file.lower().endswith(('.jpg', '.jpeg', '.png')):
|
||||
test_images.append(os.path.join(root, file))
|
||||
if len(test_images) >= 3: # 只测试3张图片
|
||||
break
|
||||
if len(test_images) >= 3:
|
||||
break
|
||||
if len(test_images) >= 3:
|
||||
break
|
||||
|
||||
if not test_images:
|
||||
print("未找到测试图片")
|
||||
return False
|
||||
|
||||
# 执行相似度检索演示
|
||||
print(f"\n开始相似度检索演示 (测试 {len(test_images)} 张图片):")
|
||||
print("-" * 50)
|
||||
|
||||
total_time = 0
|
||||
|
||||
for i, test_image in enumerate(test_images, 1):
|
||||
print(f"\n[{i}/{len(test_images)}] 测试图片: {os.path.relpath(test_image)}")
|
||||
|
||||
start_time = time.time()
|
||||
results = searcher.search_similar(test_image, k=5)
|
||||
search_time = time.time() - start_time
|
||||
total_time += search_time
|
||||
|
||||
print(f"搜索耗时: {search_time:.3f}秒")
|
||||
print("最相似的5张图片:")
|
||||
|
||||
for j, (img_path, class_name, score) in enumerate(results, 1):
|
||||
rel_path = os.path.relpath(img_path)
|
||||
print(f" {j}. {os.path.basename(img_path):20s} "
|
||||
f"类别: {class_name:8s} "
|
||||
f"相似度: {score:.4f}")
|
||||
|
||||
# 统计类别分布
|
||||
class_counts = {}
|
||||
for _, class_name, _ in results:
|
||||
class_counts[class_name] = class_counts.get(class_name, 0) + 1
|
||||
|
||||
print(f" 类别分布: {dict(class_counts)}")
|
||||
|
||||
print(f"\n演示完成!")
|
||||
print(f"平均搜索时间: {total_time/len(test_images):.3f}秒")
|
||||
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
print(f"检索演示失败: {e}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
return False
|
||||
|
||||
def show_usage_examples():
|
||||
"""显示使用示例"""
|
||||
print("\n" + "=" * 60)
|
||||
print("使用示例")
|
||||
print("=" * 60)
|
||||
|
||||
print("\n1. 查看索引信息:")
|
||||
print(" python faiss_similarity_search.py --info")
|
||||
|
||||
print("\n2. 单张图片相似度检索:")
|
||||
print(" python faiss_similarity_search.py --query dataset/test/回锅肉/aug_000.jpg")
|
||||
|
||||
print("\n3. 显示检索结果图片:")
|
||||
print(" python faiss_similarity_search.py --query dataset/test/回锅肉/aug_000.jpg --show_images")
|
||||
|
||||
print("\n4. 批量测试:")
|
||||
print(" python faiss_similarity_search.py --batch_test dataset/test --max_tests 10")
|
||||
|
||||
print("\n5. 保存结果图片:")
|
||||
print(" python faiss_similarity_search.py --query dataset/test/回锅肉/aug_000.jpg --show_images --save_result")
|
||||
|
||||
def main():
|
||||
"""主函数"""
|
||||
print("FAISS相似度检索系统 - 快速演示")
|
||||
print("=" * 60)
|
||||
|
||||
# 检查依赖
|
||||
print("检查依赖包...")
|
||||
if not check_dependencies():
|
||||
return
|
||||
print("✓ 依赖包检查通过")
|
||||
|
||||
# 检查文件
|
||||
print("检查必要文件...")
|
||||
if not check_files():
|
||||
return
|
||||
print("✓ 文件检查通过")
|
||||
|
||||
# 检查是否已有索引
|
||||
index_exists = os.path.exists("faiss_index/faiss_index.bin")
|
||||
|
||||
if index_exists:
|
||||
print("✓ 发现已存在的FAISS索引")
|
||||
|
||||
# 询问是否重新构建
|
||||
while True:
|
||||
choice = input("\n是否重新构建索引? (y/n): ").lower().strip()
|
||||
if choice in ['y', 'yes', '是']:
|
||||
if not run_index_building():
|
||||
return
|
||||
break
|
||||
elif choice in ['n', 'no', '否']:
|
||||
print("跳过索引构建，直接进行检索演示")
|
||||
break
|
||||
else:
|
||||
print("请输入 y 或 n")
|
||||
else:
|
||||
print("未发现FAISS索引，开始构建...")
|
||||
if not run_index_building():
|
||||
return
|
||||
|
||||
# 运行检索演示
|
||||
if not run_search_demo():
|
||||
return
|
||||
|
||||
# 显示使用示例
|
||||
show_usage_examples()
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("演示完成! 🎉")
|
||||
print("您现在可以使用 faiss_similarity_search.py 进行更多测试")
|
||||
print("=" * 60)
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user