171 lines
4.2 KiB
Markdown
171 lines
4.2 KiB
Markdown
# 多模态检索系统 (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 内存
|
||
|
||
## 安装依赖
|
||
|
||
```bash
|
||
pip install -r requirements.txt
|
||
```
|
||
|
||
## 快速开始
|
||
|
||
### 1. 检查GPU环境
|
||
|
||
```python
|
||
from multimodal_retrieval import check_gpu_info
|
||
check_gpu_info()
|
||
```
|
||
|
||
### 2. 初始化系统
|
||
|
||
```python
|
||
from multimodal_retrieval import MultimodalRetrieval
|
||
|
||
# 初始化检索系统
|
||
retrieval_system = MultimodalRetrieval()
|
||
```
|
||
|
||
### 3. 构建索引
|
||
|
||
```python
|
||
# 构建文本索引
|
||
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. 执行检索
|
||
|
||
```python
|
||
# 文搜文
|
||
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)
|
||
```
|
||
|
||
## 运行演示
|
||
|
||
```bash
|
||
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 许可证。
|