244 lines
7.5 KiB
Python
244 lines
7.5 KiB
Python
"""
|
|
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 = "DishClassification/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("DishClassification/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() |