多模态检索系统 (Multimodal Retrieval System)

基于 OpenSearch-AI/Ops-MM-embedding-v1-7B 模型的多模态检索系统，支持四种检索模式：文搜图、文搜文、图搜图、图搜文。

功能特性

• 文搜文 (Text-to-Text): 使用文本查询搜索相似文本
• 文搜图 (Text-to-Image): 使用文本查询搜索相关图像
• 图搜图 (Image-to-Image): 使用图像查询搜索相似图像
• 图搜文 (Image-to-Text): 使用图像查询搜索相关文本

环境要求

• Python 3.8+
• CUDA 支持的 GPU (推荐，也可使用CPU)
• 至少 8GB 内存

安装依赖

pip install -r requirements.txt

快速开始

1. 检查GPU环境

from multimodal_retrieval import check_gpu_info
check_gpu_info()

2. 初始化系统

from multimodal_retrieval import MultimodalRetrieval

# 初始化检索系统
retrieval_system = MultimodalRetrieval()

3. 构建索引

# 构建文本索引
texts = ["一只可爱的小猫", "美丽的山景", "现代化城市"]
retrieval_system.build_text_index(texts)

# 构建图像索引
image_paths = ["./images/cat.jpg", "./images/mountain.jpg", "./images/city.jpg"]
retrieval_system.build_image_index(image_paths)

4. 执行检索

# 文搜文
results = retrieval_system.search_text_by_text("猫咪", top_k=5)

# 文搜图
results = retrieval_system.search_images_by_text("动物", top_k=5)

# 图搜图
results = retrieval_system.search_images_by_image("./query_image.jpg", top_k=5)

# 图搜文
results = retrieval_system.search_text_by_image("./query_image.jpg", top_k=5)

运行演示

python demo.py

演示脚本会自动：

1. 检查GPU环境信息
2. 初始化多模态检索系统
3. 演示四种检索模式
4. 显示检索结果和相似度分数

文件结构

mmeb/
├── multimodal_retrieval.py  # 主要检索系统类
├── demo.py                  # 演示脚本
├── requirements.txt         # 依赖包列表
├── README.md               # 项目说明
└── images/                 # 图像数据目录 (需要自行创建)

API 参考

MultimodalRetrieval 类

初始化参数

• model_name: 模型名称，默认 "OpenSearch-AI/Ops-MM-embedding-v1-7B"
• device: 设备类型，默认自动选择 ("cuda" 或 "cpu")

主要方法

build_text_index(texts, save_path=None)

构建文本索引

• texts: 文本列表
• save_path: 索引保存路径 (可选)

build_image_index(image_paths, save_path=None)

构建图像索引

• image_paths: 图像路径列表
• save_path: 索引保存路径 (可选)

search_text_by_text(query, top_k=5)

文搜文检索

• query: 查询文本
• top_k: 返回结果数量
• 返回: [(文本, 相似度分数), ...]

search_images_by_text(query, top_k=5)

文搜图检索

• query: 查询文本
• top_k: 返回结果数量
• 返回: [(图像路径, 相似度分数), ...]

search_images_by_image(query_image, top_k=5)

图搜图检索

• query_image: 查询图像路径或PIL图像
• top_k: 返回结果数量
• 返回: [(图像路径, 相似度分数), ...]

search_text_by_image(query_image, top_k=5)

图搜文检索

• query_image: 查询图像路径或PIL图像
• top_k: 返回结果数量
• 返回: [(文本, 相似度分数), ...]

注意事项

1. 首次运行: 首次运行时会自动下载模型，需要网络连接
2. 内存需求: 7B参数模型需要较大内存，建议使用GPU
3. 图像格式: 支持常见图像格式 (jpg, png, bmp, gif等)
4. 批处理: 系统自动进行批处理以提高效率
5. 索引保存: 可以保存和加载索引以避免重复构建

性能优化建议

• 使用GPU加速推理
• 合理设置批处理大小
• 保存索引文件避免重复构建
• 对大量数据使用分批处理

故障排除

常见问题

1. CUDA内存不足: 减小批处理大小或使用CPU
2. 模型下载失败: 检查网络连接或使用镜像源
3. 图像加载错误: 检查图像文件路径和格式

日志信息

系统会输出详细的日志信息，包括：

• GPU环境检测结果
• 模型加载进度
• 索引构建状态
• 检索执行情况

许可证

本项目遵循 MIT 许可证。