增加FAISS向量数据库对所有训练图片进行向量化的过程。
This commit is contained in:
zhanghuan
@@ -0,0 +1,318 @@
|
||||
"""
|
||||
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()
|
||||
Reference in New Issue
Block a user