# 多模态检索系统 (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 许可证。