""" FAISS相似度检索演示程序 提供简单的命令行界面来测试相似度检索功能 功能： 1. 加载预构建的FAISS索引 2. 对指定图片进行相似度检索 3. 显示检索结果和统计信息 4. 支持批量检索和性能测试 """ import os import sys import time import argparse from typing import List, Tuple from PIL import Image import matplotlib.pyplot as plt # 导入构建的FAISS检索器 from faiss_vector_db.build_faiss_index import FAISSSearcher class SimilaritySearchDemo: """相似度检索演示类""" def __init__(self, index_dir: str, model_path: str): """ 初始化演示程序 Args: index_dir: FAISS索引目录 model_path: 模型路径 """ self.searcher = FAISSSearcher(index_dir, model_path) # 设置中文字体（用于matplotlib显示） try: # 尝试使用系统中文字体 plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'DejaVu Sans'] plt.rcParams['axes.unicode_minus'] = False except: print("警告: 无法设置中文字体，图表中的中文可能显示异常") def search_and_display(self, query_image_path: str, k: int = 5, show_images: bool = False, save_result: bool = False): """ 搜索相似图片并显示结果 Args: query_image_path: 查询图片路径 k: 返回最相似的k张图片 show_images: 是否显示图片 save_result: 是否保存结果图片 """ print(f"\n查询图片: {query_image_path}") print("-" * 50) # 检查文件是否存在 if not os.path.exists(query_image_path): print(f"错误: 文件不存在 - {query_image_path}") return # 开始计时 start_time = time.time() try: # 执行相似度搜索 results = self.searcher.search_similar(query_image_path, k) # 计算耗时 search_time = time.time() - start_time # 显示结果 print(f"搜索完成，耗时: {search_time:.3f}秒") print(f"找到 {len(results)} 个相似结果:\n") for i, (img_path, class_name, score) in enumerate(results, 1): rel_path = os.path.relpath(img_path) print(f"{i:2d}. {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"\n类别分布:") for class_name, count in class_counts.items(): print(f" {class_name}: {count} 张") # 显示图片 if show_images: self.display_results(query_image_path, results, save_result) except Exception as e: print(f"搜索失败: {e}") import traceback traceback.print_exc() def display_results(self, query_path: str, results: List[Tuple[str, str, float]], save_result: bool = False): """ 显示搜索结果图片 Args: query_path: 查询图片路径 results: 搜索结果 save_result: 是否保存结果图片 """ try: # 计算子图布局 n_results = len(results) n_cols = min(3, n_results + 1) # 包括查询图片 n_rows = (n_results + 1 + n_cols - 1) // n_cols fig, axes = plt.subplots(n_rows, n_cols, figsize=(4*n_cols, 4*n_rows)) if n_rows == 1: axes = axes.reshape(1, -1) # 显示查询图片 query_img = Image.open(query_path).convert('RGB') axes[0, 0].imshow(query_img) axes[0, 0].set_title(f"查询图片\n{os.path.basename(query_path)}", fontsize=10) axes[0, 0].axis('off') # 显示搜索结果 for i, (img_path, class_name, score) in enumerate(results): row = (i + 1) // n_cols col = (i + 1) % n_cols if row < n_rows and col < n_cols: try: img = Image.open(img_path).convert('RGB') axes[row, col].imshow(img) axes[row, col].set_title( f"#{i+1} {class_name}\n" f"{os.path.basename(img_path)}\n" f"相似度: {score:.3f}", fontsize=9 ) axes[row, col].axis('off') except Exception as e: axes[row, col].text(0.5, 0.5, f"加载失败\n{e}", ha='center', va='center') axes[row, col].axis('off') # 隐藏多余的子图 for i in range(n_results + 1, n_rows * n_cols): row = i // n_cols col = i % n_cols if row < n_rows and col < n_cols: axes[row, col].axis('off') plt.tight_layout() # 保存结果 if save_result: timestamp = time.strftime("%Y%m%d_%H%M%S") save_path = f"similarity_search_result_{timestamp}.png" plt.savefig(save_path, dpi=150, bbox_inches='tight') print(f"结果图片已保存到: {save_path}") plt.show() except Exception as e: print(f"显示图片失败: {e}") def batch_search_test(self, test_dir: str, k: int = 5, max_tests: int = 10): """ 批量搜索测试 Args: test_dir: 测试图片目录 k: 每次搜索返回的结果数 max_tests: 最大测试数量 """ print(f"\n批量搜索测试 - 目录: {test_dir}") print("=" * 60) # 获取测试图片 test_images = [] for root, dirs, files in os.walk(test_dir): for file in files: if file.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp')): test_images.append(os.path.join(root, file)) if len(test_images) >= max_tests: break if len(test_images) >= max_tests: break if not test_images: print("未找到测试图片") return print(f"找到 {len(test_images)} 张测试图片") # 执行批量测试 total_time = 0 success_count = 0 for i, img_path in enumerate(test_images, 1): print(f"\n[{i}/{len(test_images)}] 测试: {os.path.basename(img_path)}") try: start_time = time.time() results = self.searcher.search_similar(img_path, k) search_time = time.time() - start_time total_time += search_time success_count += 1 print(f" 搜索耗时: {search_time:.3f}秒") print(f" 前3个结果:") for j, (result_path, class_name, score) in enumerate(results[:3], 1): print(f" {j}. {class_name} (相似度: {score:.3f})") except Exception as e: print(f" 搜索失败: {e}") # 显示统计信息 print("\n" + "=" * 60) print("批量测试统计:") print(f" 总测试数: {len(test_images)}") print(f" 成功数: {success_count}") print(f" 失败数: {len(test_images) - success_count}") if success_count > 0: print(f" 平均搜索时间: {total_time/success_count:.3f}秒") print(f" 总耗时: {total_time:.3f}秒") def show_index_info(self): """显示索引信息""" print("\n" + "=" * 60) print("FAISS索引信息") print("=" * 60) print(f"索引目录: {self.searcher.index_dir}") print(f"总向量数: {self.searcher.index.ntotal}") print(f"特征维度: {self.searcher.embedding_dim}") print(f"类别数: {len(self.searcher.class_names)}") print(f"类别列表: {', '.join(self.searcher.class_names)}") # 显示每个类别的图片数量 print(f"\n各类别图片数量:") class_counts = {} for label in self.searcher.labels: class_name = self.searcher.class_names[label] class_counts[class_name] = class_counts.get(class_name, 0) + 1 for class_name in self.searcher.class_names: count = class_counts.get(class_name, 0) print(f" {class_name}: {count} 张") print(f"\n构建时间: {self.searcher.class_info.get('build_time', '未知')}") def main(): """主函数""" parser = argparse.ArgumentParser(description='FAISS相似度检索演示程序') parser.add_argument('--index_dir', type=str, default='faiss_vector_db/faiss_index', help='FAISS索引目录路径') parser.add_argument('--model_path', type=str, default='model/embedding_20250917_145342/best_embedding_model.pth', help='模型文件路径') parser.add_argument('--query', type=str, help='查询图片路径') parser.add_argument('--k', type=int, default=5, help='返回最相似的k张图片') parser.add_argument('--show_images', action='store_true', help='显示搜索结果图片') parser.add_argument('--save_result', action='store_true', help='保存结果图片') parser.add_argument('--batch_test', type=str, help='批量测试目录') parser.add_argument('--max_tests', type=int, default=10, help='批量测试最大数量') parser.add_argument('--info', action='store_true', help='显示索引信息') args = parser.parse_args() try: # 检查索引目录 if not os.path.exists(args.index_dir): print(f"错误: 索引目录不存在 - {args.index_dir}") print("请先运行 build_faiss_index.py 构建索引") return # 检查模型文件 if not os.path.exists(args.model_path): print(f"错误: 模型文件不存在 - {args.model_path}") return # 初始化演示程序 demo = SimilaritySearchDemo(args.index_dir, args.model_path) # 显示索引信息 if args.info: demo.show_index_info() # 单张图片搜索 if args.query: demo.search_and_display(args.query, args.k, args.show_images, args.save_result) # 批量测试 if args.batch_test: demo.batch_search_test(args.batch_test, args.k, args.max_tests) # 如果没有指定任何操作，显示帮助信息 if not any([args.query, args.batch_test, args.info]): print("使用示例:") print(f" python {sys.argv[0]} --info # 显示索引信息") print(f" python {sys.argv[0]} --query dataset/test/回锅肉/aug_000.jpg --show_images") print(f" python {sys.argv[0]} --batch_test dataset/test --max_tests 5") except Exception as e: print(f"程序执行失败: {e}") import traceback traceback.print_exc() if __name__ == "__main__": main()