base feature
This commit is contained in:
eust-w
parent
1de00fccda
commit
0cd7a4cb41
201
LICENSE
201
LICENSE
@ -1,201 +0,0 @@
|
||||
Apache License
|
||||
Version 2.0, January 2004
|
||||
http://www.apache.org/licenses/
|
||||
|
||||
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
|
||||
|
||||
1. Definitions.
|
||||
|
||||
"License" shall mean the terms and conditions for use, reproduction,
|
||||
and distribution as defined by Sections 1 through 9 of this document.
|
||||
|
||||
"Licensor" shall mean the copyright owner or entity authorized by
|
||||
the copyright owner that is granting the License.
|
||||
|
||||
"Legal Entity" shall mean the union of the acting entity and all
|
||||
other entities that control, are controlled by, or are under common
|
||||
control with that entity. For the purposes of this definition,
|
||||
"control" means (i) the power, direct or indirect, to cause the
|
||||
direction or management of such entity, whether by contract or
|
||||
otherwise, or (ii) ownership of fifty percent (50%) or more of the
|
||||
outstanding shares, or (iii) beneficial ownership of such entity.
|
||||
|
||||
"You" (or "Your") shall mean an individual or Legal Entity
|
||||
exercising permissions granted by this License.
|
||||
|
||||
"Source" form shall mean the preferred form for making modifications,
|
||||
including but not limited to software source code, documentation
|
||||
source, and configuration files.
|
||||
|
||||
"Object" form shall mean any form resulting from mechanical
|
||||
transformation or translation of a Source form, including but
|
||||
not limited to compiled object code, generated documentation,
|
||||
and conversions to other media types.
|
||||
|
||||
"Work" shall mean the work of authorship, whether in Source or
|
||||
Object form, made available under the License, as indicated by a
|
||||
copyright notice that is included in or attached to the work
|
||||
(an example is provided in the Appendix below).
|
||||
|
||||
"Derivative Works" shall mean any work, whether in Source or Object
|
||||
form, that is based on (or derived from) the Work and for which the
|
||||
editorial revisions, annotations, elaborations, or other modifications
|
||||
represent, as a whole, an original work of authorship. For the purposes
|
||||
of this License, Derivative Works shall not include works that remain
|
||||
separable from, or merely link (or bind by name) to the interfaces of,
|
||||
the Work and Derivative Works thereof.
|
||||
|
||||
"Contribution" shall mean any work of authorship, including
|
||||
the original version of the Work and any modifications or additions
|
||||
to that Work or Derivative Works thereof, that is intentionally
|
||||
submitted to Licensor for inclusion in the Work by the copyright owner
|
||||
or by an individual or Legal Entity authorized to submit on behalf of
|
||||
the copyright owner. For the purposes of this definition, "submitted"
|
||||
means any form of electronic, verbal, or written communication sent
|
||||
to the Licensor or its representatives, including but not limited to
|
||||
communication on electronic mailing lists, source code control systems,
|
||||
and issue tracking systems that are managed by, or on behalf of, the
|
||||
Licensor for the purpose of discussing and improving the Work, but
|
||||
excluding communication that is conspicuously marked or otherwise
|
||||
designated in writing by the copyright owner as "Not a Contribution."
|
||||
|
||||
"Contributor" shall mean Licensor and any individual or Legal Entity
|
||||
on behalf of whom a Contribution has been received by Licensor and
|
||||
subsequently incorporated within the Work.
|
||||
|
||||
2. Grant of Copyright License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
copyright license to reproduce, prepare Derivative Works of,
|
||||
publicly display, publicly perform, sublicense, and distribute the
|
||||
Work and such Derivative Works in Source or Object form.
|
||||
|
||||
3. Grant of Patent License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
(except as stated in this section) patent license to make, have made,
|
||||
use, offer to sell, sell, import, and otherwise transfer the Work,
|
||||
where such license applies only to those patent claims licensable
|
||||
by such Contributor that are necessarily infringed by their
|
||||
Contribution(s) alone or by combination of their Contribution(s)
|
||||
with the Work to which such Contribution(s) was submitted. If You
|
||||
institute patent litigation against any entity (including a
|
||||
cross-claim or counterclaim in a lawsuit) alleging that the Work
|
||||
or a Contribution incorporated within the Work constitutes direct
|
||||
or contributory patent infringement, then any patent licenses
|
||||
granted to You under this License for that Work shall terminate
|
||||
as of the date such litigation is filed.
|
||||
|
||||
4. Redistribution. You may reproduce and distribute copies of the
|
||||
Work or Derivative Works thereof in any medium, with or without
|
||||
modifications, and in Source or Object form, provided that You
|
||||
meet the following conditions:
|
||||
|
||||
(a) You must give any other recipients of the Work or
|
||||
Derivative Works a copy of this License; and
|
||||
|
||||
(b) You must cause any modified files to carry prominent notices
|
||||
stating that You changed the files; and
|
||||
|
||||
(c) You must retain, in the Source form of any Derivative Works
|
||||
that You distribute, all copyright, patent, trademark, and
|
||||
attribution notices from the Source form of the Work,
|
||||
excluding those notices that do not pertain to any part of
|
||||
the Derivative Works; and
|
||||
|
||||
(d) If the Work includes a "NOTICE" text file as part of its
|
||||
distribution, then any Derivative Works that You distribute must
|
||||
include a readable copy of the attribution notices contained
|
||||
within such NOTICE file, excluding those notices that do not
|
||||
pertain to any part of the Derivative Works, in at least one
|
||||
of the following places: within a NOTICE text file distributed
|
||||
as part of the Derivative Works; within the Source form or
|
||||
documentation, if provided along with the Derivative Works; or,
|
||||
within a display generated by the Derivative Works, if and
|
||||
wherever such third-party notices normally appear. The contents
|
||||
of the NOTICE file are for informational purposes only and
|
||||
do not modify the License. You may add Your own attribution
|
||||
notices within Derivative Works that You distribute, alongside
|
||||
or as an addendum to the NOTICE text from the Work, provided
|
||||
that such additional attribution notices cannot be construed
|
||||
as modifying the License.
|
||||
|
||||
You may add Your own copyright statement to Your modifications and
|
||||
may provide additional or different license terms and conditions
|
||||
for use, reproduction, or distribution of Your modifications, or
|
||||
for any such Derivative Works as a whole, provided Your use,
|
||||
reproduction, and distribution of the Work otherwise complies with
|
||||
the conditions stated in this License.
|
||||
|
||||
5. Submission of Contributions. Unless You explicitly state otherwise,
|
||||
any Contribution intentionally submitted for inclusion in the Work
|
||||
by You to the Licensor shall be under the terms and conditions of
|
||||
this License, without any additional terms or conditions.
|
||||
Notwithstanding the above, nothing herein shall supersede or modify
|
||||
the terms of any separate license agreement you may have executed
|
||||
with Licensor regarding such Contributions.
|
||||
|
||||
6. Trademarks. This License does not grant permission to use the trade
|
||||
names, trademarks, service marks, or product names of the Licensor,
|
||||
except as required for reasonable and customary use in describing the
|
||||
origin of the Work and reproducing the content of the NOTICE file.
|
||||
|
||||
7. Disclaimer of Warranty. Unless required by applicable law or
|
||||
agreed to in writing, Licensor provides the Work (and each
|
||||
Contributor provides its Contributions) on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
|
||||
implied, including, without limitation, any warranties or conditions
|
||||
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
|
||||
PARTICULAR PURPOSE. You are solely responsible for determining the
|
||||
appropriateness of using or redistributing the Work and assume any
|
||||
risks associated with Your exercise of permissions under this License.
|
||||
|
||||
8. Limitation of Liability. In no event and under no legal theory,
|
||||
whether in tort (including negligence), contract, or otherwise,
|
||||
unless required by applicable law (such as deliberate and grossly
|
||||
negligent acts) or agreed to in writing, shall any Contributor be
|
||||
liable to You for damages, including any direct, indirect, special,
|
||||
incidental, or consequential damages of any character arising as a
|
||||
result of this License or out of the use or inability to use the
|
||||
Work (including but not limited to damages for loss of goodwill,
|
||||
work stoppage, computer failure or malfunction, or any and all
|
||||
other commercial damages or losses), even if such Contributor
|
||||
has been advised of the possibility of such damages.
|
||||
|
||||
9. Accepting Warranty or Additional Liability. While redistributing
|
||||
the Work or Derivative Works thereof, You may choose to offer,
|
||||
and charge a fee for, acceptance of support, warranty, indemnity,
|
||||
or other liability obligations and/or rights consistent with this
|
||||
License. However, in accepting such obligations, You may act only
|
||||
on Your own behalf and on Your sole responsibility, not on behalf
|
||||
of any other Contributor, and only if You agree to indemnify,
|
||||
defend, and hold each Contributor harmless for any liability
|
||||
incurred by, or claims asserted against, such Contributor by reason
|
||||
of your accepting any such warranty or additional liability.
|
||||
|
||||
END OF TERMS AND CONDITIONS
|
||||
|
||||
APPENDIX: How to apply the Apache License to your work.
|
||||
|
||||
To apply the Apache License to your work, attach the following
|
||||
boilerplate notice, with the fields enclosed by brackets "[]"
|
||||
replaced with your own identifying information. (Don't include
|
||||
the brackets!) The text should be enclosed in the appropriate
|
||||
comment syntax for the file format. We also recommend that a
|
||||
file or class name and description of purpose be included on the
|
||||
same "printed page" as the copyright notice for easier
|
||||
identification within third-party archives.
|
||||
|
||||
Copyright [yyyy] [name of copyright owner]
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
170
README.md
Normal file
170
README.md
Normal file
@ -0,0 +1,170 @@
|
||||
# 多模态检索系统 (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 许可证。
|
||||
BIN
__pycache__/multimodal_retrieval_multigpu.cpython-310.pyc
Normal file
BIN
__pycache__/multimodal_retrieval_multigpu.cpython-310.pyc
Normal file
Binary file not shown.
632
multimodal_retrieval_multigpu.py
Normal file
632
multimodal_retrieval_multigpu.py
Normal file
@ -0,0 +1,632 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.nn.parallel import DataParallel, DistributedDataParallel
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
import faiss
|
||||
from transformers import AutoModel, AutoProcessor, AutoTokenizer
|
||||
from typing import List, Union, Tuple, Dict
|
||||
import os
|
||||
import json
|
||||
from pathlib import Path
|
||||
import logging
|
||||
import gc
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
import threading
|
||||
|
||||
# 设置日志
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class MultiGPUMultimodalRetrieval:
|
||||
"""多GPU优化的多模态检索系统，支持文搜图、文搜文、图搜图、图搜文"""
|
||||
|
||||
def __init__(self, model_name: str = "OpenSearch-AI/Ops-MM-embedding-v1-7B",
|
||||
use_all_gpus: bool = True, gpu_ids: List[int] = None, min_memory_gb=12):
|
||||
"""
|
||||
初始化多GPU多模态检索系统
|
||||
|
||||
Args:
|
||||
model_name: 模型名称
|
||||
use_all_gpus: 是否使用所有可用GPU
|
||||
gpu_ids: 指定使用的GPU ID列表
|
||||
min_memory_gb: 最小可用内存（GB）
|
||||
"""
|
||||
self.model_name = model_name
|
||||
|
||||
# 设置GPU设备
|
||||
self._setup_devices(use_all_gpus, gpu_ids, min_memory_gb)
|
||||
|
||||
# 清理GPU内存
|
||||
self._clear_all_gpu_memory()
|
||||
|
||||
logger.info(f"正在加载模型到多GPU: {self.device_ids}")
|
||||
|
||||
# 加载模型和处理器
|
||||
self.model = None
|
||||
self.tokenizer = None
|
||||
self.processor = None
|
||||
self._load_model_multigpu()
|
||||
|
||||
# 初始化索引
|
||||
self.text_index = None
|
||||
self.image_index = None
|
||||
self.text_data = []
|
||||
self.image_data = []
|
||||
|
||||
logger.info("多GPU模型加载完成")
|
||||
|
||||
def _setup_devices(self, use_all_gpus: bool, gpu_ids: List[int], min_memory_gb=12):
|
||||
"""设置GPU设备"""
|
||||
if not torch.cuda.is_available():
|
||||
raise RuntimeError("CUDA不可用，无法使用多GPU")
|
||||
|
||||
total_gpus = torch.cuda.device_count()
|
||||
logger.info(f"检测到 {total_gpus} 个GPU")
|
||||
|
||||
# 检查是否设置了CUDA_VISIBLE_DEVICES
|
||||
cuda_visible_devices = os.environ.get('CUDA_VISIBLE_DEVICES')
|
||||
if cuda_visible_devices is not None:
|
||||
# 如果设置了CUDA_VISIBLE_DEVICES，使用可见的GPU
|
||||
visible_gpu_count = len(cuda_visible_devices.split(','))
|
||||
self.device_ids = list(range(visible_gpu_count))
|
||||
logger.info(f"使用CUDA_VISIBLE_DEVICES指定的GPU: {cuda_visible_devices}")
|
||||
elif use_all_gpus:
|
||||
self.device_ids = self._select_best_gpus(min_memory_gb)
|
||||
elif gpu_ids:
|
||||
self.device_ids = gpu_ids
|
||||
else:
|
||||
self.device_ids = [0]
|
||||
|
||||
self.num_gpus = len(self.device_ids)
|
||||
self.primary_device = f"cuda:{self.device_ids[0]}"
|
||||
|
||||
logger.info(f"使用GPU: {self.device_ids}, 主设备: {self.primary_device}")
|
||||
|
||||
def _clear_all_gpu_memory(self):
|
||||
"""清理所有GPU内存"""
|
||||
for gpu_id in self.device_ids:
|
||||
torch.cuda.set_device(gpu_id)
|
||||
torch.cuda.empty_cache()
|
||||
torch.cuda.synchronize()
|
||||
gc.collect()
|
||||
logger.info("所有GPU内存已清理")
|
||||
|
||||
def _load_model_multigpu(self):
|
||||
"""加载模型到多GPU"""
|
||||
try:
|
||||
# 设置环境变量优化内存使用
|
||||
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'
|
||||
|
||||
# 清理GPU内存
|
||||
self._clear_gpu_memory()
|
||||
|
||||
# 首先尝试使用accelerate的自动设备映射
|
||||
if self.num_gpus > 1:
|
||||
# 设置最大内存限制（每个GPU 18GB，留出缓冲）
|
||||
max_memory = {i: "18GiB" for i in self.device_ids}
|
||||
|
||||
logger.info(f"正在加载模型到多GPU: {self.device_ids}")
|
||||
self.model = AutoModel.from_pretrained(
|
||||
self.model_name,
|
||||
trust_remote_code=True,
|
||||
torch_dtype=torch.float16,
|
||||
device_map="auto",
|
||||
max_memory=max_memory,
|
||||
low_cpu_mem_usage=True,
|
||||
offload_folder="./offload"
|
||||
)
|
||||
else:
|
||||
# 单GPU加载
|
||||
self.model = AutoModel.from_pretrained(
|
||||
self.model_name,
|
||||
trust_remote_code=True,
|
||||
torch_dtype=torch.float16,
|
||||
device_map=self.primary_device
|
||||
)
|
||||
|
||||
# 加载分词器和处理器到主设备
|
||||
try:
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(
|
||||
self.model_name,
|
||||
trust_remote_code=True
|
||||
)
|
||||
logger.info("Tokenizer加载成功")
|
||||
except Exception as e:
|
||||
logger.error(f"Tokenizer加载失败: {e}")
|
||||
return False
|
||||
|
||||
# 加载处理器用于图像处理
|
||||
try:
|
||||
self.processor = AutoProcessor.from_pretrained(
|
||||
self.model_name,
|
||||
trust_remote_code=True
|
||||
)
|
||||
logger.info("Processor加载成功")
|
||||
except Exception as e:
|
||||
logger.warning(f"Processor加载失败: {e}")
|
||||
# 如果AutoProcessor失败，尝试使用tokenizer作为fallback
|
||||
logger.info("尝试使用tokenizer作为processor的fallback")
|
||||
self.processor = self.tokenizer
|
||||
|
||||
logger.info(f"模型已成功加载到设备: {self.model.hf_device_map if hasattr(self.model, 'hf_device_map') else self.primary_device}")
|
||||
logger.info("多GPU模型加载完成")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"多GPU模型加载失败: {str(e)}")
|
||||
return False
|
||||
|
||||
def _clear_gpu_memory(self):
|
||||
"""清理GPU内存"""
|
||||
for gpu_id in self.device_ids:
|
||||
torch.cuda.set_device(gpu_id)
|
||||
torch.cuda.empty_cache()
|
||||
torch.cuda.synchronize()
|
||||
gc.collect()
|
||||
logger.info("GPU内存已清理")
|
||||
|
||||
def _get_gpu_memory_info(self):
|
||||
"""获取GPU内存使用情况"""
|
||||
try:
|
||||
import subprocess
|
||||
result = subprocess.run(['nvidia-smi', '--query-gpu=memory.used,memory.total', '--format=csv,nounits,noheader'],
|
||||
capture_output=True, text=True, check=True)
|
||||
lines = result.stdout.strip().split('\n')
|
||||
gpu_info = []
|
||||
for i, line in enumerate(lines):
|
||||
used, total = map(int, line.split(', '))
|
||||
free = total - used
|
||||
gpu_info.append({
|
||||
'gpu_id': i,
|
||||
'used': used,
|
||||
'total': total,
|
||||
'free': free,
|
||||
'usage_percent': (used / total) * 100
|
||||
})
|
||||
return gpu_info
|
||||
except Exception as e:
|
||||
logger.warning(f"无法获取GPU内存信息: {e}")
|
||||
return []
|
||||
|
||||
def _select_best_gpus(self, min_memory_gb=12):
|
||||
"""选择内存充足的GPU"""
|
||||
gpu_info = self._get_gpu_memory_info()
|
||||
if not gpu_info:
|
||||
return list(range(torch.cuda.device_count()))
|
||||
|
||||
# 按可用内存排序
|
||||
gpu_info.sort(key=lambda x: x['free'], reverse=True)
|
||||
|
||||
# 选择内存充足的GPU
|
||||
min_memory_mb = min_memory_gb * 1024
|
||||
suitable_gpus = []
|
||||
|
||||
for gpu in gpu_info:
|
||||
if gpu['free'] >= min_memory_mb:
|
||||
suitable_gpus.append(gpu['gpu_id'])
|
||||
logger.info(f"GPU {gpu['gpu_id']}: {gpu['free']}MB 可用 (合适)")
|
||||
else:
|
||||
logger.warning(f"GPU {gpu['gpu_id']}: {gpu['free']}MB 可用 (不足)")
|
||||
|
||||
if not suitable_gpus:
|
||||
# 如果没有GPU满足要求，选择可用内存最多的
|
||||
logger.warning(f"没有GPU有足够内存({min_memory_gb}GB)，选择可用内存最多的GPU")
|
||||
suitable_gpus = [gpu_info[0]['gpu_id']]
|
||||
|
||||
return suitable_gpus
|
||||
|
||||
def encode_text_batch(self, texts: List[str]) -> np.ndarray:
|
||||
"""
|
||||
批量编码文本为向量（多GPU优化）
|
||||
|
||||
Args:
|
||||
texts: 文本列表
|
||||
|
||||
Returns:
|
||||
文本向量
|
||||
"""
|
||||
if not texts:
|
||||
return np.array([])
|
||||
|
||||
with torch.no_grad():
|
||||
# 预处理输入
|
||||
inputs = self.tokenizer(
|
||||
text=texts,
|
||||
return_tensors="pt",
|
||||
padding=True,
|
||||
truncation=True,
|
||||
max_length=512
|
||||
)
|
||||
|
||||
# 将输入移动到主设备
|
||||
inputs = {k: v.to(self.primary_device) for k, v in inputs.items()}
|
||||
|
||||
# 前向传播
|
||||
outputs = self.model(**inputs)
|
||||
embeddings = outputs.last_hidden_state.mean(dim=1)
|
||||
|
||||
# 清理GPU内存
|
||||
del inputs, outputs
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
return embeddings.cpu().numpy().astype(np.float32)
|
||||
|
||||
def encode_image_batch(self, images: List[Union[str, Image.Image]]) -> np.ndarray:
|
||||
"""
|
||||
批量编码图像为向量
|
||||
|
||||
Args:
|
||||
images: 图像路径或PIL图像列表
|
||||
|
||||
Returns:
|
||||
图像向量
|
||||
"""
|
||||
if not images:
|
||||
return np.array([])
|
||||
|
||||
# 预处理图像
|
||||
processed_images = []
|
||||
for img in images:
|
||||
if isinstance(img, str):
|
||||
img = Image.open(img).convert('RGB')
|
||||
elif isinstance(img, Image.Image):
|
||||
img = img.convert('RGB')
|
||||
processed_images.append(img)
|
||||
|
||||
try:
|
||||
logger.info(f"处理 {len(processed_images)} 张图像")
|
||||
|
||||
# 使用多模态模型生成图像embedding
|
||||
# 为每张图像创建简单的文本描述作为输入
|
||||
conversations = []
|
||||
for i in range(len(processed_images)):
|
||||
# 使用简化的对话格式
|
||||
conversation = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "image", "image": processed_images[i]},
|
||||
{"type": "text", "text": "What is in this image?"}
|
||||
]
|
||||
}
|
||||
]
|
||||
conversations.append(conversation)
|
||||
|
||||
# 使用processor处理
|
||||
try:
|
||||
# 尝试使用apply_chat_template方法
|
||||
texts = []
|
||||
for conv in conversations:
|
||||
text = self.processor.apply_chat_template(conv, tokenize=False, add_generation_prompt=False)
|
||||
texts.append(text)
|
||||
|
||||
# 处理文本和图像
|
||||
inputs = self.processor(
|
||||
text=texts,
|
||||
images=processed_images,
|
||||
return_tensors="pt",
|
||||
padding=True
|
||||
)
|
||||
|
||||
# 移动到GPU
|
||||
inputs = {k: v.to(self.primary_device) for k, v in inputs.items()}
|
||||
|
||||
# 获取模型输出
|
||||
with torch.no_grad():
|
||||
outputs = self.model(**inputs)
|
||||
embeddings = outputs.last_hidden_state.mean(dim=1)
|
||||
|
||||
# 转换为numpy数组
|
||||
embeddings = embeddings.cpu().numpy().astype(np.float32)
|
||||
|
||||
except Exception as inner_e:
|
||||
logger.warning(f"多模态模型图像编码失败，使用文本模式: {inner_e}")
|
||||
return np.zeros((len(processed_images), 3584), dtype=np.float32)
|
||||
# 如果多模态失败，使用纯文本描述作为fallback
|
||||
image_descriptions = ["An image" for _ in processed_images]
|
||||
text_inputs = self.processor(
|
||||
text=image_descriptions,
|
||||
return_tensors="pt",
|
||||
padding=True,
|
||||
truncation=True,
|
||||
max_length=512
|
||||
)
|
||||
text_inputs = {k: v.to(self.primary_device) for k, v in text_inputs.items()}
|
||||
|
||||
with torch.no_grad():
|
||||
outputs = self.model(**text_inputs)
|
||||
embeddings = outputs.last_hidden_state.mean(dim=1)
|
||||
|
||||
embeddings = embeddings.cpu().numpy().astype(np.float32)
|
||||
|
||||
logger.info(f"生成图像embeddings: {embeddings.shape}")
|
||||
return embeddings
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"图像编码失败: {e}")
|
||||
# 返回与文本embedding维度一致的零向量作为fallback
|
||||
embedding_dim = 3584
|
||||
embeddings = np.zeros((len(processed_images), embedding_dim), dtype=np.float32)
|
||||
return embeddings
|
||||
|
||||
def build_text_index_parallel(self, texts: List[str], save_path: str = None):
|
||||
"""
|
||||
并行构建文本索引（多GPU优化）
|
||||
|
||||
Args:
|
||||
texts: 文本列表
|
||||
save_path: 索引保存路径
|
||||
"""
|
||||
logger.info(f"正在并行构建文本索引，共 {len(texts)} 条文本")
|
||||
|
||||
# 根据GPU数量调整批次大小
|
||||
batch_size = max(4, 16 // self.num_gpus)
|
||||
all_embeddings = []
|
||||
|
||||
# 分批处理
|
||||
for i in range(0, len(texts), batch_size):
|
||||
batch_texts = texts[i:i+batch_size]
|
||||
|
||||
try:
|
||||
embeddings = self.encode_text_batch(batch_texts)
|
||||
all_embeddings.append(embeddings)
|
||||
|
||||
# 显示进度
|
||||
if (i // batch_size + 1) % 10 == 0:
|
||||
logger.info(f"已处理 {i + len(batch_texts)}/{len(texts)} 条文本")
|
||||
|
||||
except torch.cuda.OutOfMemoryError:
|
||||
logger.warning(f"GPU内存不足，跳过批次 {i}-{i+len(batch_texts)}")
|
||||
self._clear_all_gpu_memory()
|
||||
continue
|
||||
except Exception as e:
|
||||
logger.error(f"处理文本批次时出错: {e}")
|
||||
continue
|
||||
|
||||
if not all_embeddings:
|
||||
raise ValueError("没有成功处理任何文本")
|
||||
|
||||
# 合并所有嵌入向量
|
||||
embeddings = np.vstack(all_embeddings)
|
||||
|
||||
# 构建FAISS索引
|
||||
dimension = embeddings.shape[1]
|
||||
self.text_index = faiss.IndexFlatIP(dimension)
|
||||
|
||||
# 归一化向量
|
||||
faiss.normalize_L2(embeddings)
|
||||
self.text_index.add(embeddings)
|
||||
|
||||
self.text_data = texts
|
||||
|
||||
if save_path:
|
||||
self._save_index(self.text_index, texts, save_path + "_text")
|
||||
|
||||
logger.info("文本索引构建完成")
|
||||
|
||||
def build_image_index_parallel(self, image_paths: List[str], save_path: str = None):
|
||||
"""
|
||||
并行构建图像索引（多GPU优化）
|
||||
|
||||
Args:
|
||||
image_paths: 图像路径列表
|
||||
save_path: 索引保存路径
|
||||
"""
|
||||
logger.info(f"正在并行构建图像索引，共 {len(image_paths)} 张图像")
|
||||
|
||||
# 图像处理使用更小的批次
|
||||
batch_size = max(2, 8 // self.num_gpus)
|
||||
all_embeddings = []
|
||||
|
||||
for i in range(0, len(image_paths), batch_size):
|
||||
batch_images = image_paths[i:i+batch_size]
|
||||
|
||||
try:
|
||||
embeddings = self.encode_image_batch(batch_images)
|
||||
all_embeddings.append(embeddings)
|
||||
|
||||
# 显示进度
|
||||
if (i // batch_size + 1) % 5 == 0:
|
||||
logger.info(f"已处理 {i + len(batch_images)}/{len(image_paths)} 张图像")
|
||||
|
||||
except torch.cuda.OutOfMemoryError:
|
||||
logger.warning(f"GPU内存不足，跳过图像批次 {i}-{i+len(batch_images)}")
|
||||
self._clear_all_gpu_memory()
|
||||
continue
|
||||
except Exception as e:
|
||||
logger.error(f"处理图像批次时出错: {e}")
|
||||
continue
|
||||
|
||||
if not all_embeddings:
|
||||
raise ValueError("没有成功处理任何图像")
|
||||
|
||||
embeddings = np.vstack(all_embeddings)
|
||||
|
||||
# 构建FAISS索引
|
||||
dimension = embeddings.shape[1]
|
||||
self.image_index = faiss.IndexFlatIP(dimension)
|
||||
|
||||
faiss.normalize_L2(embeddings)
|
||||
self.image_index.add(embeddings)
|
||||
|
||||
self.image_data = image_paths
|
||||
|
||||
if save_path:
|
||||
self._save_index(self.image_index, image_paths, save_path + "_image")
|
||||
|
||||
logger.info("图像索引构建完成")
|
||||
|
||||
def search_text_by_text(self, query: str, top_k: int = 5) -> List[Tuple[str, float]]:
|
||||
"""文搜文：使用文本查询搜索相似文本"""
|
||||
if self.text_index is None:
|
||||
raise ValueError("文本索引未构建，请先调用 build_text_index_parallel")
|
||||
|
||||
query_embedding = self.encode_text_batch([query]).astype(np.float32)
|
||||
faiss.normalize_L2(query_embedding)
|
||||
|
||||
scores, indices = self.text_index.search(query_embedding, top_k)
|
||||
|
||||
results = []
|
||||
for score, idx in zip(scores[0], indices[0]):
|
||||
if idx != -1:
|
||||
results.append((self.text_data[idx], float(score)))
|
||||
|
||||
return results
|
||||
|
||||
def search_images_by_text(self, query: str, top_k: int = 5) -> List[Tuple[str, float]]:
|
||||
"""文搜图：使用文本查询搜索相似图像"""
|
||||
if self.image_index is None:
|
||||
raise ValueError("图像索引未构建，请先调用 build_image_index_parallel")
|
||||
|
||||
query_embedding = self.encode_text_batch([query]).astype(np.float32)
|
||||
faiss.normalize_L2(query_embedding)
|
||||
|
||||
scores, indices = self.image_index.search(query_embedding, top_k)
|
||||
|
||||
results = []
|
||||
for score, idx in zip(scores[0], indices[0]):
|
||||
if idx != -1:
|
||||
results.append((self.image_data[idx], float(score)))
|
||||
|
||||
return results
|
||||
|
||||
def search_images_by_image(self, query_image: Union[str, Image.Image], top_k: int = 5) -> List[Tuple[str, float]]:
|
||||
"""图搜图：使用图像查询搜索相似图像"""
|
||||
if self.image_index is None:
|
||||
raise ValueError("图像索引未构建，请先调用 build_image_index_parallel")
|
||||
|
||||
query_embedding = self.encode_image_batch([query_image]).astype(np.float32)
|
||||
faiss.normalize_L2(query_embedding)
|
||||
|
||||
scores, indices = self.image_index.search(query_embedding, top_k)
|
||||
|
||||
results = []
|
||||
for score, idx in zip(scores[0], indices[0]):
|
||||
if idx != -1:
|
||||
results.append((self.image_data[idx], float(score)))
|
||||
|
||||
return results
|
||||
|
||||
def search_text_by_image(self, query_image: Union[str, Image.Image], top_k: int = 5) -> List[Tuple[str, float]]:
|
||||
"""图搜文：使用图像查询搜索相似文本"""
|
||||
if self.text_index is None:
|
||||
raise ValueError("文本索引未构建，请先调用 build_text_index_parallel")
|
||||
|
||||
query_embedding = self.encode_image_batch([query_image]).astype(np.float32)
|
||||
faiss.normalize_L2(query_embedding)
|
||||
|
||||
scores, indices = self.text_index.search(query_embedding, top_k)
|
||||
|
||||
results = []
|
||||
for score, idx in zip(scores[0], indices[0]):
|
||||
if idx != -1:
|
||||
results.append((self.text_data[idx], float(score)))
|
||||
|
||||
return results
|
||||
|
||||
# Web应用兼容的方法名称
|
||||
def search_text_to_image(self, query: str, top_k: int = 5) -> List[Tuple[str, float]]:
|
||||
"""文搜图：Web应用兼容方法"""
|
||||
return self.search_images_by_text(query, top_k)
|
||||
|
||||
def search_image_to_image(self, query_image: Union[str, Image.Image], top_k: int = 5) -> List[Tuple[str, float]]:
|
||||
"""图搜图：Web应用兼容方法"""
|
||||
return self.search_images_by_image(query_image, top_k)
|
||||
|
||||
def search_text_to_text(self, query: str, top_k: int = 5) -> List[Tuple[str, float]]:
|
||||
"""文搜文：Web应用兼容方法"""
|
||||
return self.search_text_by_text(query, top_k)
|
||||
|
||||
def search_image_to_text(self, query_image: Union[str, Image.Image], top_k: int = 5) -> List[Tuple[str, float]]:
|
||||
"""图搜文：Web应用兼容方法"""
|
||||
return self.search_text_by_image(query_image, top_k)
|
||||
|
||||
def _save_index(self, index, data, path_prefix):
|
||||
"""保存索引和数据"""
|
||||
faiss.write_index(index, f"{path_prefix}.index")
|
||||
with open(f"{path_prefix}.json", 'w', encoding='utf-8') as f:
|
||||
json.dump(data, f, ensure_ascii=False, indent=2)
|
||||
|
||||
def load_index(self, path_prefix, index_type="text"):
|
||||
"""加载已保存的索引"""
|
||||
index = faiss.read_index(f"{path_prefix}.index")
|
||||
with open(f"{path_prefix}.json", 'r', encoding='utf-8') as f:
|
||||
data = json.load(f)
|
||||
|
||||
if index_type == "text":
|
||||
self.text_index = index
|
||||
self.text_data = data
|
||||
else:
|
||||
self.image_index = index
|
||||
self.image_data = data
|
||||
|
||||
logger.info(f"已加载 {index_type} 索引")
|
||||
|
||||
def get_gpu_memory_info(self):
|
||||
"""获取所有GPU内存使用信息"""
|
||||
memory_info = {}
|
||||
for gpu_id in self.device_ids:
|
||||
torch.cuda.set_device(gpu_id)
|
||||
allocated = torch.cuda.memory_allocated(gpu_id) / 1024**3
|
||||
cached = torch.cuda.memory_reserved(gpu_id) / 1024**3
|
||||
total = torch.cuda.get_device_properties(gpu_id).total_memory / 1024**3
|
||||
free = total - cached
|
||||
|
||||
memory_info[f"GPU_{gpu_id}"] = {
|
||||
"total": f"{total:.1f}GB",
|
||||
"allocated": f"{allocated:.1f}GB",
|
||||
"cached": f"{cached:.1f}GB",
|
||||
"free": f"{free:.1f}GB"
|
||||
}
|
||||
|
||||
return memory_info
|
||||
|
||||
def check_multigpu_info():
|
||||
"""检查多GPU环境信息"""
|
||||
print("=== 多GPU环境信息 ===")
|
||||
|
||||
if not torch.cuda.is_available():
|
||||
print("❌ CUDA不可用")
|
||||
return
|
||||
|
||||
gpu_count = torch.cuda.device_count()
|
||||
print(f"✅ 检测到 {gpu_count} 个GPU")
|
||||
print(f"CUDA版本: {torch.version.cuda}")
|
||||
print(f"PyTorch版本: {torch.__version__}")
|
||||
|
||||
for i in range(gpu_count):
|
||||
gpu_name = torch.cuda.get_device_name(i)
|
||||
gpu_memory = torch.cuda.get_device_properties(i).total_memory / 1024**3
|
||||
print(f"GPU {i}: {gpu_name} ({gpu_memory:.1f}GB)")
|
||||
|
||||
print("=====================")
|
||||
|
||||
if __name__ == "__main__":
|
||||
# 检查多GPU环境
|
||||
check_multigpu_info()
|
||||
|
||||
# 示例使用
|
||||
print("\n正在初始化多GPU多模态检索系统...")
|
||||
|
||||
try:
|
||||
retrieval_system = MultiGPUMultimodalRetrieval()
|
||||
print("✅ 多GPU系统初始化成功！")
|
||||
|
||||
# 显示GPU内存使用情况
|
||||
memory_info = retrieval_system.get_gpu_memory_info()
|
||||
print("\n📊 GPU内存使用情况:")
|
||||
for gpu, info in memory_info.items():
|
||||
print(f" {gpu}: {info['allocated']} / {info['total']} (已用/总计)")
|
||||
|
||||
print("\n🚀 多GPU多模态检索系统就绪！")
|
||||
print("支持的检索模式:")
|
||||
print("1. 文搜文: search_text_by_text()")
|
||||
print("2. 文搜图: search_images_by_text()")
|
||||
print("3. 图搜图: search_images_by_image()")
|
||||
print("4. 图搜文: search_text_by_image()")
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ 多GPU系统初始化失败: {e}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
309
ops_mm_embedding_v1.py
Normal file
309
ops_mm_embedding_v1.py
Normal file
@ -0,0 +1,309 @@
|
||||
import math
|
||||
from typing import List, Optional, TypeAlias, Union
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from PIL import Image
|
||||
from tqdm import tqdm
|
||||
from transformers import AutoModelForImageTextToText, AutoProcessor
|
||||
|
||||
ImageInput: TypeAlias = Union[Image.Image, List[Image.Image]]
|
||||
BatchImageInput: TypeAlias = Union[List[Image.Image], List[List[Image.Image]]]
|
||||
|
||||
|
||||
class OpsMMEmbeddingV1(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
model_name: str,
|
||||
device: str = "cuda",
|
||||
max_length: Optional[int] = None,
|
||||
attn_implementation: Optional[str] = None,
|
||||
):
|
||||
super().__init__()
|
||||
self.device = device
|
||||
self.max_length = max_length
|
||||
self.default_instruction = "You are a helpful assistant."
|
||||
self.base_model = AutoModelForImageTextToText.from_pretrained(
|
||||
model_name,
|
||||
torch_dtype=torch.bfloat16,
|
||||
low_cpu_mem_usage=True,
|
||||
attn_implementation=attn_implementation,
|
||||
).to(self.device)
|
||||
|
||||
self.processor = AutoProcessor.from_pretrained(model_name, min_pixels=256 * 28 * 28, max_pixels=1280 * 28 * 28)
|
||||
self.processor.tokenizer.padding_side = "left"
|
||||
self.eval()
|
||||
|
||||
def encode_input(self, input):
|
||||
hidden_states = self.base_model(**input, return_dict=True, output_hidden_states=True)
|
||||
hidden_states = hidden_states.hidden_states[-1]
|
||||
pooled_output = self._pooling(hidden_states)
|
||||
return pooled_output
|
||||
|
||||
def _pooling(self, last_hidden_state):
|
||||
batch_size = last_hidden_state.shape[0]
|
||||
reps = last_hidden_state[torch.arange(batch_size), -1, :]
|
||||
reps = torch.nn.functional.normalize(reps, p=2, dim=-1)
|
||||
return reps
|
||||
|
||||
def _validate_instructions(
|
||||
self,
|
||||
texts: Optional[List[str]],
|
||||
images: Optional[BatchImageInput],
|
||||
instruction: Optional[Union[str, List[str]]],
|
||||
) -> List[str]:
|
||||
"""Validate and format instructions to match batch size"""
|
||||
batch_size = max(len(x) if x is not None else 0 for x in [texts, images])
|
||||
|
||||
if instruction is None:
|
||||
return [self.default_instruction] * batch_size
|
||||
|
||||
if isinstance(instruction, str):
|
||||
return [instruction] * batch_size
|
||||
|
||||
if isinstance(instruction, list):
|
||||
if len(instruction) != batch_size:
|
||||
raise ValueError(f"Length of instruction list ({len(instruction)}) must match batch size ({batch_size}) when texts/images are provided")
|
||||
return instruction
|
||||
|
||||
raise TypeError("instruction must be str, List[str] or None")
|
||||
|
||||
def _process_images(self, images: ImageInput) -> List[Image.Image]:
|
||||
"""Convert single image or list of images to processed format"""
|
||||
if isinstance(images, Image.Image) or isinstance(images, str):
|
||||
return [fetch_image(images)]
|
||||
return [fetch_image(i) for i in images]
|
||||
|
||||
def embed(
|
||||
self,
|
||||
texts: Optional[List[str]] = None,
|
||||
images: Optional[BatchImageInput] = None,
|
||||
instruction: Optional[Union[str, List[str]]] = None,
|
||||
**kwargs,
|
||||
) -> torch.Tensor:
|
||||
"""Generate embeddings for text, images, or combined inputs.
|
||||
|
||||
Args:
|
||||
texts: List of text inputs (optional)
|
||||
images: Can be:
|
||||
- List[Image.Image]: Single image per input
|
||||
- List[List[Image.Image]]: Multiple images per input
|
||||
instruction: Instruction(s) for the model. Can be:
|
||||
- None: use default instruction
|
||||
- str: use same instruction for all inputs
|
||||
- List[str]: per-input instructions (must match batch size)
|
||||
"""
|
||||
if texts is None and images is None:
|
||||
raise ValueError("Either texts or images must be provided")
|
||||
|
||||
instructions = self._validate_instructions(texts, images, instruction)
|
||||
|
||||
# Determine batch size
|
||||
batch_size = len(texts) if texts is not None else len(images) # type: ignore
|
||||
|
||||
input_texts, input_images = [], []
|
||||
for i in range(batch_size):
|
||||
text = texts[i] if texts is not None else None
|
||||
image = images[i] if images is not None else None
|
||||
|
||||
input_str = ""
|
||||
processed_image = None
|
||||
if image is not None:
|
||||
processed_image = self._process_images(image)
|
||||
input_str += "<|vision_start|><|image_pad|><|vision_end|>" * len(processed_image)
|
||||
|
||||
if text is not None:
|
||||
input_str += text
|
||||
|
||||
msg = f"<|im_start|>system\n{instructions[i]}<|im_end|>\n<|im_start|>user\n{input_str}<|im_end|>\n<|im_start|>assistant\n<|endoftext|>"
|
||||
|
||||
input_texts.append(msg)
|
||||
input_images.append(processed_image)
|
||||
|
||||
# Only pass to processor if we actually have images
|
||||
processed_images = input_images if any(img is not None for img in input_images) else None
|
||||
|
||||
inputs = self.processor(
|
||||
text=input_texts,
|
||||
images=processed_images,
|
||||
padding=True,
|
||||
truncation=True,
|
||||
max_length=self.max_length,
|
||||
return_tensors="pt",
|
||||
)
|
||||
inputs = {k: v.to(self.device) for k, v in inputs.items()}
|
||||
|
||||
with torch.inference_mode():
|
||||
embeddings = self.encode_input(inputs)
|
||||
|
||||
return embeddings
|
||||
|
||||
def get_text_embeddings(
|
||||
self,
|
||||
texts: List[str],
|
||||
instruction: Optional[Union[str, List[str]]] = None,
|
||||
**kwargs,
|
||||
) -> torch.Tensor:
|
||||
"""Convenience method for text-only embeddings"""
|
||||
return self.get_fused_embeddings(texts=texts, instruction=instruction, **kwargs)
|
||||
|
||||
def get_image_embeddings(
|
||||
self,
|
||||
images: BatchImageInput,
|
||||
instruction: Optional[Union[str, List[str]]] = None,
|
||||
**kwargs,
|
||||
) -> torch.Tensor:
|
||||
"""Convenience method for image-only embeddings.
|
||||
|
||||
Args:
|
||||
images: Can be:
|
||||
- List[Image.Image]: Single image per input
|
||||
- List[List[Image.Image]]: Multiple images per input
|
||||
"""
|
||||
return self.get_fused_embeddings(images=images, instruction=instruction, **kwargs)
|
||||
|
||||
def get_fused_embeddings(
|
||||
self,
|
||||
texts: Optional[List[str]] = None,
|
||||
images: Optional[BatchImageInput] = None,
|
||||
instruction: Optional[Union[str, List[str]]] = None,
|
||||
batch_size: int = 8,
|
||||
show_progress: bool = True,
|
||||
**kwargs,
|
||||
) -> torch.Tensor:
|
||||
"""Batch processing for large collections of texts/images.
|
||||
|
||||
Args:
|
||||
texts: List of text inputs (optional)
|
||||
images: Can be:
|
||||
- List[Image.Image]: Single image per input
|
||||
- List[List[Image.Image]]: Multiple images per input
|
||||
instruction: Instruction(s) for the model
|
||||
batch_size: Number of items to process at once
|
||||
show_progress: Whether to display progress bar
|
||||
"""
|
||||
|
||||
if texts is None and images is None:
|
||||
raise ValueError("Either texts or images must be provided")
|
||||
|
||||
total_items = len(texts) if texts is not None else len(images) # type: ignore
|
||||
num_batches = math.ceil(total_items / batch_size)
|
||||
|
||||
all_embeddings = []
|
||||
progress = tqdm(total=num_batches, disable=not show_progress, desc="Processing")
|
||||
|
||||
for i in range(0, total_items, batch_size):
|
||||
batch_texts = texts[i : i + batch_size] if texts is not None else None
|
||||
batch_images = images[i : i + batch_size] if images is not None else None
|
||||
batch_emb = self.embed(texts=batch_texts, images=batch_images, instruction=instruction)
|
||||
|
||||
all_embeddings.append(batch_emb.cpu())
|
||||
progress.update(1)
|
||||
|
||||
progress.close()
|
||||
return torch.cat(all_embeddings, dim=0).to(self.device)
|
||||
|
||||
def forward(self, **inputs) -> torch.Tensor:
|
||||
"""Alias for encode_input"""
|
||||
return self.encode_input(inputs)
|
||||
|
||||
|
||||
### Modified from qwen_vl_utils.vision_process.py
|
||||
import base64
|
||||
import logging
|
||||
import math
|
||||
from io import BytesIO
|
||||
|
||||
import requests
|
||||
|
||||
IMAGE_FACTOR = 28
|
||||
MIN_PIXELS = 256 * 28 * 28
|
||||
MAX_PIXELS = 1280 * 28 * 28
|
||||
MAX_RATIO = 200
|
||||
|
||||
|
||||
def round_by_factor(number: int, factor: int) -> int:
|
||||
"""Returns the closest integer to 'number' that is divisible by 'factor'."""
|
||||
return round(number / factor) * factor
|
||||
|
||||
|
||||
def ceil_by_factor(number: int | float, factor: int) -> int:
|
||||
"""Returns the smallest integer greater than or equal to 'number' that is divisible by 'factor'."""
|
||||
return math.ceil(number / factor) * factor
|
||||
|
||||
|
||||
def floor_by_factor(number: int | float, factor: int) -> int:
|
||||
"""Returns the largest integer less than or equal to 'number' that is divisible by 'factor'."""
|
||||
return math.floor(number / factor) * factor
|
||||
|
||||
|
||||
def smart_resize(
|
||||
height: int,
|
||||
width: int,
|
||||
factor: int = IMAGE_FACTOR,
|
||||
min_pixels: int = MIN_PIXELS,
|
||||
max_pixels: int = MAX_PIXELS,
|
||||
) -> tuple[int, int]:
|
||||
"""
|
||||
Rescales the image so that the following conditions are met:
|
||||
1. Both dimensions (height and width) are divisible by 'factor'.
|
||||
2. The total number of pixels is within the range ['min_pixels', 'max_pixels'].
|
||||
3. The aspect ratio of the image is maintained as closely as possible.
|
||||
"""
|
||||
h_bar = max(factor, round_by_factor(height, factor))
|
||||
w_bar = max(factor, round_by_factor(width, factor))
|
||||
if h_bar * w_bar > max_pixels:
|
||||
beta = math.sqrt((height * width) / max_pixels)
|
||||
h_bar = floor_by_factor(height / beta, factor)
|
||||
w_bar = floor_by_factor(width / beta, factor)
|
||||
elif h_bar * w_bar < min_pixels:
|
||||
beta = math.sqrt(min_pixels / (height * width))
|
||||
h_bar = ceil_by_factor(height * beta, factor)
|
||||
w_bar = ceil_by_factor(width * beta, factor)
|
||||
|
||||
if max(h_bar, w_bar) / min(h_bar, w_bar) > MAX_RATIO:
|
||||
logging.warning(f"Absolute aspect ratio must be smaller than {MAX_RATIO}, got {max(h_bar, w_bar) / min(h_bar, w_bar)}")
|
||||
if h_bar > w_bar:
|
||||
h_bar = w_bar * MAX_RATIO
|
||||
else:
|
||||
w_bar = h_bar * MAX_RATIO
|
||||
return h_bar, w_bar
|
||||
|
||||
|
||||
def fetch_image(
|
||||
image: str | Image.Image,
|
||||
size_factor: int = IMAGE_FACTOR,
|
||||
min_pixels: int = MIN_PIXELS,
|
||||
max_pixels: int = MAX_PIXELS,
|
||||
) -> Image.Image:
|
||||
image_obj = None
|
||||
if isinstance(image, Image.Image):
|
||||
image_obj = image
|
||||
elif image.startswith("http://") or image.startswith("https://"):
|
||||
image_obj = Image.open(requests.get(image, stream=True).raw) # type: ignore
|
||||
elif image.startswith("file://"):
|
||||
image_obj = Image.open(image[7:])
|
||||
elif image.startswith("data:image"):
|
||||
if "base64," in image:
|
||||
_, base64_data = image.split("base64,", 1)
|
||||
data = base64.b64decode(base64_data)
|
||||
image_obj = Image.open(BytesIO(data))
|
||||
else:
|
||||
image_obj = Image.open(image)
|
||||
if image_obj is None:
|
||||
raise ValueError(f"Unrecognized image input, support local path, http url, base64 and PIL.Image, got {image}")
|
||||
image = image_obj.convert("RGB")
|
||||
width, height = image.size
|
||||
resized_height, resized_width = smart_resize(
|
||||
height,
|
||||
width,
|
||||
factor=size_factor,
|
||||
min_pixels=min_pixels,
|
||||
max_pixels=max_pixels,
|
||||
)
|
||||
image = image.resize((resized_width, resized_height))
|
||||
|
||||
return image
|
||||
|
||||
|
||||
###
|
||||
12
requirements.txt
Normal file
12
requirements.txt
Normal file
@ -0,0 +1,12 @@
|
||||
torch>=2.0.0
|
||||
torchvision>=0.15.0
|
||||
transformers>=4.30.0
|
||||
accelerate>=0.20.0
|
||||
faiss-cpu>=1.7.4
|
||||
numpy>=1.21.0
|
||||
Pillow>=9.0.0
|
||||
scikit-learn>=1.3.0
|
||||
tqdm>=4.65.0
|
||||
flask>=2.3.0
|
||||
werkzeug>=2.3.0
|
||||
psutil>=5.9.0
|
||||
BIN
sample_images/1755681385_4__.jpg
Normal file
BIN
sample_images/1755681385_4__.jpg
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 189 KiB |
BIN
sample_images/1755681385_5__.jpg
Normal file
BIN
sample_images/1755681385_5__.jpg
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 201 KiB |
BIN
sample_images/1755681385_6__.jpg
Normal file
BIN
sample_images/1755681385_6__.jpg
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 144 KiB |
BIN
sample_images/1755681385_7__.jpg
Normal file
BIN
sample_images/1755681385_7__.jpg
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 166 KiB |
971
templates/index.html
Normal file
971
templates/index.html
Normal file
@ -0,0 +1,971 @@
|
||||
<!DOCTYPE html>
|
||||
<html lang="zh-CN">
|
||||
<head>
|
||||
<meta charset="UTF-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||||
<title>多模态检索系统</title>
|
||||
<link href="https://cdn.jsdelivr.net/npm/bootstrap@5.1.3/dist/css/bootstrap.min.css" rel="stylesheet">
|
||||
<link href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.0.0/css/all.min.css" rel="stylesheet">
|
||||
<style>
|
||||
:root {
|
||||
--primary-color: #2563eb;
|
||||
--secondary-color: #64748b;
|
||||
--success-color: #059669;
|
||||
--warning-color: #d97706;
|
||||
--danger-color: #dc2626;
|
||||
--bg-gradient: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
|
||||
}
|
||||
|
||||
body {
|
||||
background: var(--bg-gradient);
|
||||
min-height: 100vh;
|
||||
font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
|
||||
}
|
||||
|
||||
.main-container {
|
||||
background: rgba(255, 255, 255, 0.95);
|
||||
backdrop-filter: blur(10px);
|
||||
border-radius: 20px;
|
||||
box-shadow: 0 20px 40px rgba(0, 0, 0, 0.1);
|
||||
margin: 20px auto;
|
||||
max-width: 1200px;
|
||||
}
|
||||
|
||||
.header {
|
||||
background: linear-gradient(135deg, var(--primary-color), #3b82f6);
|
||||
color: white;
|
||||
padding: 2rem;
|
||||
border-radius: 20px 20px 0 0;
|
||||
text-align: center;
|
||||
}
|
||||
|
||||
.mode-card {
|
||||
background: white;
|
||||
border-radius: 15px;
|
||||
padding: 1.5rem;
|
||||
margin-bottom: 1.5rem;
|
||||
box-shadow: 0 4px 15px rgba(0, 0, 0, 0.1);
|
||||
transition: all 0.3s ease;
|
||||
border: 2px solid transparent;
|
||||
}
|
||||
|
||||
.mode-card:hover {
|
||||
transform: translateY(-5px);
|
||||
box-shadow: 0 8px 25px rgba(0, 0, 0, 0.15);
|
||||
}
|
||||
|
||||
.mode-card.active {
|
||||
border-color: var(--primary-color);
|
||||
background: linear-gradient(135deg, #eff6ff, #dbeafe);
|
||||
}
|
||||
|
||||
.mode-icon {
|
||||
font-size: 2.5rem;
|
||||
margin-bottom: 1rem;
|
||||
display: block;
|
||||
}
|
||||
|
||||
.text-to-text { color: #059669; }
|
||||
.text-to-image { color: #dc2626; }
|
||||
.image-to-text { color: #d97706; }
|
||||
.image-to-image { color: #7c3aed; }
|
||||
|
||||
.search-input {
|
||||
border-radius: 12px;
|
||||
border: 2px solid #e5e7eb;
|
||||
padding: 12px 16px;
|
||||
font-size: 16px;
|
||||
transition: all 0.3s ease;
|
||||
}
|
||||
|
||||
.search-input:focus {
|
||||
border-color: var(--primary-color);
|
||||
box-shadow: 0 0 0 3px rgba(37, 99, 235, 0.1);
|
||||
}
|
||||
|
||||
.btn-primary {
|
||||
background: var(--primary-color);
|
||||
border: none;
|
||||
border-radius: 12px;
|
||||
padding: 12px 24px;
|
||||
font-weight: 600;
|
||||
transition: all 0.3s ease;
|
||||
}
|
||||
|
||||
.btn-primary:hover {
|
||||
background: #1d4ed8;
|
||||
transform: translateY(-2px);
|
||||
}
|
||||
|
||||
.file-upload-area {
|
||||
border: 3px dashed #d1d5db;
|
||||
border-radius: 12px;
|
||||
padding: 3rem;
|
||||
text-align: center;
|
||||
transition: all 0.3s ease;
|
||||
cursor: pointer;
|
||||
}
|
||||
|
||||
.file-upload-area:hover {
|
||||
border-color: var(--primary-color);
|
||||
background: rgba(37, 99, 235, 0.05);
|
||||
}
|
||||
|
||||
.file-upload-area.dragover {
|
||||
border-color: var(--primary-color);
|
||||
background: rgba(37, 99, 235, 0.1);
|
||||
}
|
||||
|
||||
.result-card {
|
||||
background: white;
|
||||
border-radius: 12px;
|
||||
padding: 1.5rem;
|
||||
margin-bottom: 1rem;
|
||||
box-shadow: 0 2px 10px rgba(0, 0, 0, 0.1);
|
||||
border-left: 4px solid var(--primary-color);
|
||||
}
|
||||
|
||||
.result-image {
|
||||
max-width: 200px;
|
||||
max-height: 150px;
|
||||
border-radius: 8px;
|
||||
object-fit: cover;
|
||||
}
|
||||
|
||||
.score-badge {
|
||||
background: var(--success-color);
|
||||
color: white;
|
||||
padding: 4px 12px;
|
||||
border-radius: 20px;
|
||||
font-size: 0.85rem;
|
||||
font-weight: 600;
|
||||
}
|
||||
|
||||
.loading-spinner {
|
||||
display: none;
|
||||
text-align: center;
|
||||
padding: 2rem;
|
||||
}
|
||||
|
||||
.status-indicator {
|
||||
position: fixed;
|
||||
top: 20px;
|
||||
right: 20px;
|
||||
z-index: 1000;
|
||||
}
|
||||
|
||||
.fade-in {
|
||||
animation: fadeIn 0.5s ease-in;
|
||||
}
|
||||
|
||||
@keyframes fadeIn {
|
||||
from { opacity: 0; transform: translateY(20px); }
|
||||
to { opacity: 1; transform: translateY(0); }
|
||||
}
|
||||
|
||||
.query-image {
|
||||
max-width: 300px;
|
||||
max-height: 200px;
|
||||
border-radius: 12px;
|
||||
object-fit: cover;
|
||||
box-shadow: 0 4px 15px rgba(0, 0, 0, 0.2);
|
||||
}
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<!-- 状态指示器 -->
|
||||
<div class="status-indicator">
|
||||
<div id="statusBadge" class="badge bg-secondary">
|
||||
<i class="fas fa-circle-notch fa-spin"></i> 未初始化
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div class="container-fluid">
|
||||
<div class="main-container">
|
||||
<!-- 头部 -->
|
||||
<div class="header">
|
||||
<h1><i class="fas fa-search"></i> 多模态检索系统</h1>
|
||||
<p class="mb-0">支持文搜图、文搜文、图搜图、图搜文四种检索模式</p>
|
||||
</div>
|
||||
|
||||
<div class="p-4">
|
||||
<!-- 重新初始化按钮 -->
|
||||
<div class="text-center mb-4">
|
||||
<button id="reinitBtn" class="btn btn-warning">
|
||||
<i class="fas fa-redo"></i> 重新初始化系统
|
||||
</button>
|
||||
</div>
|
||||
|
||||
<!-- 检索模式选择 -->
|
||||
<div class="row mb-4" id="modeSelection">
|
||||
<div class="col-md-3">
|
||||
<div class="mode-card text-center" data-mode="text_to_text">
|
||||
<i class="fas fa-file-text mode-icon text-to-text"></i>
|
||||
<h5>文搜文</h5>
|
||||
<p class="text-muted">文本查找相似文本</p>
|
||||
</div>
|
||||
</div>
|
||||
<div class="col-md-3">
|
||||
<div class="mode-card text-center" data-mode="text_to_image">
|
||||
<i class="fas fa-image mode-icon text-to-image"></i>
|
||||
<h5>文搜图</h5>
|
||||
<p class="text-muted">文本查找相关图片</p>
|
||||
</div>
|
||||
</div>
|
||||
<div class="col-md-3">
|
||||
<div class="mode-card text-center" data-mode="image_to_text">
|
||||
<i class="fas fa-comment mode-icon image-to-text"></i>
|
||||
<h5>图搜文</h5>
|
||||
<p class="text-muted">图片查找相关文本</p>
|
||||
</div>
|
||||
</div>
|
||||
<div class="col-md-3">
|
||||
<div class="mode-card text-center" data-mode="image_to_image">
|
||||
<i class="fas fa-images mode-icon image-to-image"></i>
|
||||
<h5>图搜图</h5>
|
||||
<p class="text-muted">图片查找相似图片</p>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- 数据管理界面 -->
|
||||
<div class="row mb-4" id="dataManagement">
|
||||
<div class="col-12">
|
||||
<div class="card">
|
||||
<div class="card-header">
|
||||
<h5><i class="fas fa-database"></i> 数据管理</h5>
|
||||
<small class="text-muted">上传和管理检索数据库</small>
|
||||
</div>
|
||||
<div class="card-body">
|
||||
<div class="row">
|
||||
<!-- 批量上传图片 -->
|
||||
<div class="col-md-6">
|
||||
<div class="upload-section">
|
||||
<h6><i class="fas fa-images text-primary"></i> 批量上传图片</h6>
|
||||
<div class="file-upload-area" id="batchImageUpload">
|
||||
<i class="fas fa-cloud-upload-alt fa-2x text-muted mb-2"></i>
|
||||
<p>拖拽多张图片到此处或点击选择</p>
|
||||
<input type="file" id="batchImageFiles" multiple accept="image/*" style="display: none;">
|
||||
<button class="btn btn-outline-primary btn-sm mt-2" onclick="document.getElementById('batchImageFiles').click()">
|
||||
<i class="fas fa-folder-open"></i> 选择图片
|
||||
</button>
|
||||
</div>
|
||||
<div id="imageUploadProgress" class="mt-2" style="display: none;">
|
||||
<div class="progress">
|
||||
<div class="progress-bar" role="progressbar" style="width: 0%"></div>
|
||||
</div>
|
||||
<small class="text-muted mt-1 d-block">上传进度: <span id="imageProgressText">0/0</span></small>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- 批量上传文本 -->
|
||||
<div class="col-md-6">
|
||||
<div class="upload-section">
|
||||
<h6><i class="fas fa-file-text text-success"></i> 批量上传文本</h6>
|
||||
<div class="mb-3">
|
||||
<textarea id="batchTextInput" class="form-control" rows="8"
|
||||
placeholder="请输入文本数据，每行一条文本记录... 例如： 这是第一条文本记录 这是第二条文本记录 这是第三条文本记录"></textarea>
|
||||
</div>
|
||||
<div class="d-flex gap-2">
|
||||
<button id="uploadTextsBtn" class="btn btn-success">
|
||||
<i class="fas fa-upload"></i> 上传文本
|
||||
</button>
|
||||
<button class="btn btn-outline-secondary" onclick="document.getElementById('textFile').click()">
|
||||
<i class="fas fa-file-import"></i> 从文件导入
|
||||
</button>
|
||||
<input type="file" id="textFile" accept=".txt,.csv" style="display: none;">
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- 数据统计和管理 -->
|
||||
<div class="row mt-4">
|
||||
<div class="col-md-8">
|
||||
<div class="d-flex gap-3">
|
||||
<button id="buildIndexBtn" class="btn btn-warning" disabled>
|
||||
<i class="fas fa-cogs"></i> 构建索引
|
||||
</button>
|
||||
<button id="viewDataBtn" class="btn btn-info">
|
||||
<i class="fas fa-list"></i> 查看数据
|
||||
</button>
|
||||
<button id="clearDataBtn" class="btn btn-danger">
|
||||
<i class="fas fa-trash"></i> 清空数据
|
||||
</button>
|
||||
</div>
|
||||
</div>
|
||||
<div class="col-md-4">
|
||||
<div id="dataStats" class="text-end">
|
||||
<small class="text-muted">
|
||||
图片: <span id="imageCount">0</span> 张 |
|
||||
文本: <span id="textCount">0</span> 条
|
||||
</small>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- 搜索界面 -->
|
||||
<div id="searchInterface" style="display: none;">
|
||||
<!-- 文本搜索 -->
|
||||
<div id="textSearch" class="search-panel" style="display: none;">
|
||||
<div class="row">
|
||||
<div class="col-md-8">
|
||||
<input type="text" id="textQuery" class="form-control search-input"
|
||||
placeholder="请输入搜索文本...">
|
||||
</div>
|
||||
<div class="col-md-2">
|
||||
<select id="textTopK" class="form-select search-input">
|
||||
<option value="3">Top 3</option>
|
||||
<option value="5" selected>Top 5</option>
|
||||
<option value="10">Top 10</option>
|
||||
</select>
|
||||
</div>
|
||||
<div class="col-md-2">
|
||||
<button id="textSearchBtn" class="btn btn-primary w-100">
|
||||
<i class="fas fa-search"></i> 搜索
|
||||
</button>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- 图片搜索 -->
|
||||
<div id="imageSearch" class="search-panel" style="display: none;">
|
||||
<div class="row">
|
||||
<div class="col-md-8">
|
||||
<div class="file-upload-area" id="fileUploadArea">
|
||||
<i class="fas fa-cloud-upload-alt fa-3x text-muted mb-3"></i>
|
||||
<h5>拖拽图片到此处或点击选择</h5>
|
||||
<p class="text-muted">支持 PNG, JPG, JPEG, GIF, BMP, WebP 格式</p>
|
||||
<input type="file" id="imageFile" accept="image/*" style="display: none;">
|
||||
</div>
|
||||
</div>
|
||||
<div class="col-md-2">
|
||||
<select id="imageTopK" class="form-select search-input">
|
||||
<option value="3">Top 3</option>
|
||||
<option value="5" selected>Top 5</option>
|
||||
<option value="10">Top 10</option>
|
||||
</select>
|
||||
</div>
|
||||
<div class="col-md-2">
|
||||
<button id="imageSearchBtn" class="btn btn-primary w-100" disabled>
|
||||
<i class="fas fa-search"></i> 搜索
|
||||
</button>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- 加载动画 -->
|
||||
<div class="loading-spinner" id="loadingSpinner">
|
||||
<div class="spinner-border text-primary" role="status">
|
||||
<span class="visually-hidden">Loading...</span>
|
||||
</div>
|
||||
<p class="mt-2">正在搜索中...</p>
|
||||
</div>
|
||||
|
||||
<!-- 搜索结果 -->
|
||||
<div id="searchResults" class="mt-4"></div>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<script src="https://cdn.jsdelivr.net/npm/bootstrap@5.1.3/dist/js/bootstrap.bundle.min.js"></script>
|
||||
<script>
|
||||
let currentMode = null;
|
||||
let systemInitialized = false;
|
||||
|
||||
// 重新初始化系统
|
||||
document.getElementById('reinitBtn').addEventListener('click', async function() {
|
||||
const btn = this;
|
||||
const originalText = btn.innerHTML;
|
||||
|
||||
btn.innerHTML = '<i class="fas fa-spinner fa-spin"></i> 重新初始化中...';
|
||||
btn.disabled = true;
|
||||
|
||||
try {
|
||||
const response = await fetch('/api/init', {
|
||||
method: 'POST',
|
||||
headers: {'Content-Type': 'application/json'}
|
||||
});
|
||||
|
||||
const data = await response.json();
|
||||
|
||||
if (data.success) {
|
||||
systemInitialized = true;
|
||||
document.getElementById('statusBadge').innerHTML =
|
||||
'<i class="fas fa-check-circle"></i> 已重新初始化';
|
||||
document.getElementById('statusBadge').className = 'badge bg-success';
|
||||
|
||||
showAlert('success', `系统重新初始化成功！GPU: ${data.gpu_count} 个`);
|
||||
} else {
|
||||
throw new Error(data.message);
|
||||
}
|
||||
} catch (error) {
|
||||
showAlert('danger', '重新初始化失败: ' + error.message);
|
||||
} finally {
|
||||
btn.innerHTML = originalText;
|
||||
btn.disabled = false;
|
||||
}
|
||||
});
|
||||
|
||||
// 模式选择
|
||||
document.querySelectorAll('.mode-card').forEach(card => {
|
||||
card.addEventListener('click', function() {
|
||||
|
||||
// 更新选中状态
|
||||
document.querySelectorAll('.mode-card').forEach(c => c.classList.remove('active'));
|
||||
this.classList.add('active');
|
||||
|
||||
currentMode = this.dataset.mode;
|
||||
setupSearchInterface(currentMode);
|
||||
});
|
||||
});
|
||||
|
||||
// 设置搜索界面
|
||||
function setupSearchInterface(mode) {
|
||||
document.getElementById('searchInterface').style.display = 'block';
|
||||
document.getElementById('textSearch').style.display = 'none';
|
||||
document.getElementById('imageSearch').style.display = 'none';
|
||||
document.getElementById('searchResults').innerHTML = '';
|
||||
|
||||
if (mode === 'text_to_text' || mode === 'text_to_image') {
|
||||
document.getElementById('textSearch').style.display = 'block';
|
||||
document.getElementById('textQuery').placeholder =
|
||||
mode === 'text_to_text' ? '请输入要搜索的文本...' : '请输入要搜索图片的描述...';
|
||||
} else {
|
||||
document.getElementById('imageSearch').style.display = 'block';
|
||||
}
|
||||
}
|
||||
|
||||
// 文本搜索
|
||||
document.getElementById('textSearchBtn').addEventListener('click', performTextSearch);
|
||||
document.getElementById('textQuery').addEventListener('keypress', function(e) {
|
||||
if (e.key === 'Enter') performTextSearch();
|
||||
});
|
||||
|
||||
async function performTextSearch() {
|
||||
const query = document.getElementById('textQuery').value.trim();
|
||||
const topK = parseInt(document.getElementById('textTopK').value);
|
||||
|
||||
if (!query) {
|
||||
showAlert('warning', '请输入搜索文本');
|
||||
return;
|
||||
}
|
||||
|
||||
showLoading(true);
|
||||
|
||||
try {
|
||||
const endpoint = currentMode === 'text_to_text' ? '/api/search/text_to_text' : '/api/search/text_to_image';
|
||||
const response = await fetch(endpoint, {
|
||||
method: 'POST',
|
||||
headers: {'Content-Type': 'application/json'},
|
||||
body: JSON.stringify({query, top_k: topK})
|
||||
});
|
||||
|
||||
const data = await response.json();
|
||||
|
||||
if (data.success) {
|
||||
displayResults(data, currentMode);
|
||||
} else {
|
||||
throw new Error(data.message);
|
||||
}
|
||||
} catch (error) {
|
||||
showAlert('danger', '搜索失败: ' + error.message);
|
||||
} finally {
|
||||
showLoading(false);
|
||||
}
|
||||
}
|
||||
|
||||
// 图片上传处理
|
||||
const fileUploadArea = document.getElementById('fileUploadArea');
|
||||
const imageFile = document.getElementById('imageFile');
|
||||
|
||||
fileUploadArea.addEventListener('click', () => imageFile.click());
|
||||
fileUploadArea.addEventListener('dragover', handleDragOver);
|
||||
fileUploadArea.addEventListener('drop', handleDrop);
|
||||
imageFile.addEventListener('change', handleFileSelect);
|
||||
|
||||
function handleDragOver(e) {
|
||||
e.preventDefault();
|
||||
fileUploadArea.classList.add('dragover');
|
||||
}
|
||||
|
||||
function handleDrop(e) {
|
||||
e.preventDefault();
|
||||
fileUploadArea.classList.remove('dragover');
|
||||
const files = e.dataTransfer.files;
|
||||
if (files.length > 0) {
|
||||
handleFile(files[0]);
|
||||
}
|
||||
}
|
||||
|
||||
function handleFileSelect(e) {
|
||||
const file = e.target.files[0];
|
||||
if (file) handleFile(file);
|
||||
}
|
||||
|
||||
function handleFile(file) {
|
||||
if (!file.type.startsWith('image/')) {
|
||||
showAlert('warning', '请选择图片文件');
|
||||
return;
|
||||
}
|
||||
|
||||
const reader = new FileReader();
|
||||
reader.onload = function(e) {
|
||||
fileUploadArea.innerHTML = `
|
||||
<img src="${e.target.result}" class="query-image mb-3">
|
||||
<p class="text-success"><i class="fas fa-check"></i> 图片已选择: ${file.name}</p>
|
||||
`;
|
||||
document.getElementById('imageSearchBtn').disabled = false;
|
||||
};
|
||||
reader.readAsDataURL(file);
|
||||
}
|
||||
|
||||
// 图片搜索
|
||||
document.getElementById('imageSearchBtn').addEventListener('click', async function() {
|
||||
const file = imageFile.files[0];
|
||||
const topK = parseInt(document.getElementById('imageTopK').value);
|
||||
|
||||
if (!file) {
|
||||
showAlert('warning', '请选择图片文件');
|
||||
return;
|
||||
}
|
||||
|
||||
showLoading(true);
|
||||
|
||||
try {
|
||||
const formData = new FormData();
|
||||
formData.append('image', file);
|
||||
formData.append('top_k', topK);
|
||||
|
||||
const endpoint = currentMode === 'image_to_text' ? '/api/search/image_to_text' : '/api/search/image_to_image';
|
||||
const response = await fetch(endpoint, {
|
||||
method: 'POST',
|
||||
body: formData
|
||||
});
|
||||
|
||||
const data = await response.json();
|
||||
|
||||
if (data.success) {
|
||||
displayResults(data, currentMode);
|
||||
} else {
|
||||
throw new Error(data.message);
|
||||
}
|
||||
} catch (error) {
|
||||
showAlert('danger', '搜索失败: ' + error.message);
|
||||
} finally {
|
||||
showLoading(false);
|
||||
}
|
||||
});
|
||||
|
||||
// 显示结果
|
||||
function displayResults(data, mode) {
|
||||
const resultsContainer = document.getElementById('searchResults');
|
||||
|
||||
let html = `
|
||||
<div class="fade-in">
|
||||
<div class="d-flex justify-content-between align-items-center mb-3">
|
||||
<h4><i class="fas fa-search-plus"></i> 搜索结果</h4>
|
||||
<div>
|
||||
<span class="badge bg-info">找到 ${data.result_count} 个结果</span>
|
||||
<span class="badge bg-secondary">耗时 ${data.search_time}s</span>
|
||||
</div>
|
||||
</div>
|
||||
`;
|
||||
|
||||
if (data.query_image) {
|
||||
html += `
|
||||
<div class="result-card">
|
||||
<h6><i class="fas fa-image"></i> 查询图片</h6>
|
||||
<img src="data:image/jpeg;base64,${data.query_image}" class="query-image">
|
||||
</div>
|
||||
`;
|
||||
}
|
||||
|
||||
if (data.query) {
|
||||
html += `
|
||||
<div class="result-card">
|
||||
<h6><i class="fas fa-quote-left"></i> 查询文本</h6>
|
||||
<p class="mb-0">"${data.query}"</p>
|
||||
</div>
|
||||
`;
|
||||
}
|
||||
|
||||
data.results.forEach((result, index) => {
|
||||
html += '<div class="result-card">';
|
||||
|
||||
if (mode === 'text_to_image' || mode === 'image_to_image') {
|
||||
html += `
|
||||
<div class="row">
|
||||
<div class="col-md-3">
|
||||
<img src="data:image/jpeg;base64,${result.image_base64}"
|
||||
class="result-image" alt="Result ${index + 1}">
|
||||
</div>
|
||||
<div class="col-md-9">
|
||||
<div class="d-flex justify-content-between align-items-start">
|
||||
<h6><i class="fas fa-image"></i> ${result.filename}</h6>
|
||||
<span class="score-badge">相似度: ${(result.score * 100).toFixed(1)}%</span>
|
||||
</div>
|
||||
<p class="text-muted mb-0">路径: ${result.image_path}</p>
|
||||
</div>
|
||||
</div>
|
||||
`;
|
||||
} else {
|
||||
html += `
|
||||
<div class="d-flex justify-content-between align-items-start">
|
||||
<div>
|
||||
<h6><i class="fas fa-file-text"></i> 结果 ${index + 1}</h6>
|
||||
<p class="mb-0">${result.text || result}</p>
|
||||
</div>
|
||||
<span class="score-badge">相似度: ${((result.score || 0.95) * 100).toFixed(1)}%</span>
|
||||
</div>
|
||||
`;
|
||||
}
|
||||
|
||||
html += '</div>';
|
||||
});
|
||||
|
||||
html += '</div>';
|
||||
resultsContainer.innerHTML = html;
|
||||
}
|
||||
|
||||
// 工具函数
|
||||
function showLoading(show) {
|
||||
document.getElementById('loadingSpinner').style.display = show ? 'block' : 'none';
|
||||
}
|
||||
|
||||
function showAlert(type, message) {
|
||||
const alertDiv = document.createElement('div');
|
||||
alertDiv.className = `alert alert-${type} alert-dismissible fade show`;
|
||||
alertDiv.innerHTML = `
|
||||
${message}
|
||||
<button type="button" class="btn-close" data-bs-dismiss="alert"></button>
|
||||
`;
|
||||
|
||||
document.querySelector('.main-container .p-4').insertBefore(alertDiv, document.querySelector('.main-container .p-4').firstChild);
|
||||
|
||||
setTimeout(() => {
|
||||
if (alertDiv.parentNode) {
|
||||
alertDiv.remove();
|
||||
}
|
||||
}, 5000);
|
||||
}
|
||||
|
||||
// 检查系统状态
|
||||
async function checkStatus() {
|
||||
try {
|
||||
const response = await fetch('/api/status');
|
||||
const data = await response.json();
|
||||
|
||||
if (data.initialized) {
|
||||
systemInitialized = true;
|
||||
document.getElementById('statusBadge').innerHTML =
|
||||
'<i class="fas fa-check-circle"></i> 已初始化';
|
||||
document.getElementById('statusBadge').className = 'badge bg-success';
|
||||
} else {
|
||||
document.getElementById('statusBadge').innerHTML =
|
||||
'<i class="fas fa-exclamation-triangle"></i> 未初始化';
|
||||
document.getElementById('statusBadge').className = 'badge bg-warning';
|
||||
}
|
||||
} catch (error) {
|
||||
console.log('Status check failed:', error);
|
||||
document.getElementById('statusBadge').innerHTML =
|
||||
'<i class="fas fa-times-circle"></i> 连接失败';
|
||||
document.getElementById('statusBadge').className = 'badge bg-danger';
|
||||
}
|
||||
}
|
||||
|
||||
// 页面加载时检查状态
|
||||
checkStatus();
|
||||
|
||||
// 设置数据管理功能事件绑定
|
||||
setupDataManagement();
|
||||
|
||||
function setupDataManagement() {
|
||||
// 批量图片上传事件
|
||||
const batchImageUpload = document.getElementById('batchImageUpload');
|
||||
const batchImageFiles = document.getElementById('batchImageFiles');
|
||||
|
||||
// 拖拽上传
|
||||
batchImageUpload.addEventListener('dragover', function(e) {
|
||||
e.preventDefault();
|
||||
this.classList.add('dragover');
|
||||
});
|
||||
|
||||
batchImageUpload.addEventListener('dragleave', function(e) {
|
||||
e.preventDefault();
|
||||
this.classList.remove('dragover');
|
||||
});
|
||||
|
||||
batchImageUpload.addEventListener('drop', function(e) {
|
||||
e.preventDefault();
|
||||
this.classList.remove('dragover');
|
||||
const files = Array.from(e.dataTransfer.files).filter(file => file.type.startsWith('image/'));
|
||||
if (files.length > 0) {
|
||||
uploadBatchImages(files);
|
||||
}
|
||||
});
|
||||
|
||||
batchImageFiles.addEventListener('change', function(e) {
|
||||
if (e.target.files.length > 0) {
|
||||
uploadBatchImages(Array.from(e.target.files));
|
||||
}
|
||||
});
|
||||
|
||||
// 批量文本上传
|
||||
document.getElementById('uploadTextsBtn').addEventListener('click', function() {
|
||||
const textData = document.getElementById('batchTextInput').value.trim();
|
||||
if (textData) {
|
||||
const texts = textData.split('\n').filter(line => line.trim());
|
||||
if (texts.length > 0) {
|
||||
uploadBatchTexts(texts);
|
||||
} else {
|
||||
showAlert('warning', '请输入有效的文本数据');
|
||||
}
|
||||
} else {
|
||||
showAlert('warning', '请输入文本数据');
|
||||
}
|
||||
});
|
||||
|
||||
// 从文件导入文本
|
||||
document.getElementById('textFile').addEventListener('change', function(e) {
|
||||
const file = e.target.files[0];
|
||||
if (file) {
|
||||
const reader = new FileReader();
|
||||
reader.onload = function(e) {
|
||||
document.getElementById('batchTextInput').value = e.target.result;
|
||||
};
|
||||
reader.readAsText(file);
|
||||
}
|
||||
});
|
||||
|
||||
// 构建索引
|
||||
document.getElementById('buildIndexBtn').addEventListener('click', buildIndex);
|
||||
|
||||
// 查看数据
|
||||
document.getElementById('viewDataBtn').addEventListener('click', viewData);
|
||||
|
||||
// 清空数据
|
||||
document.getElementById('clearDataBtn').addEventListener('click', clearData);
|
||||
|
||||
// 初始化时更新数据统计
|
||||
updateDataStats();
|
||||
}
|
||||
|
||||
// 批量上传图片
|
||||
async function uploadBatchImages(files) {
|
||||
const progressDiv = document.getElementById('imageUploadProgress');
|
||||
const progressBar = progressDiv.querySelector('.progress-bar');
|
||||
const progressText = document.getElementById('imageProgressText');
|
||||
|
||||
progressDiv.style.display = 'block';
|
||||
progressText.textContent = `0/${files.length}`;
|
||||
progressBar.style.width = '0%';
|
||||
|
||||
const formData = new FormData();
|
||||
files.forEach(file => {
|
||||
formData.append('images', file);
|
||||
});
|
||||
|
||||
try {
|
||||
const response = await fetch('/api/upload/images', {
|
||||
method: 'POST',
|
||||
body: formData
|
||||
});
|
||||
|
||||
const data = await response.json();
|
||||
|
||||
if (data.success) {
|
||||
progressBar.style.width = '100%';
|
||||
progressText.textContent = `${files.length}/${files.length}`;
|
||||
showAlert('success', `成功上传 ${data.uploaded_count} 张图片`);
|
||||
updateDataStats();
|
||||
document.getElementById('buildIndexBtn').disabled = false;
|
||||
} else {
|
||||
showAlert('danger', `上传失败: ${data.message}`);
|
||||
}
|
||||
} catch (error) {
|
||||
showAlert('danger', `上传错误: ${error.message}`);
|
||||
} finally {
|
||||
setTimeout(() => {
|
||||
progressDiv.style.display = 'none';
|
||||
}, 2000);
|
||||
}
|
||||
}
|
||||
|
||||
// 批量上传文本
|
||||
async function uploadBatchTexts(texts) {
|
||||
try {
|
||||
const response = await fetch('/api/upload/texts', {
|
||||
method: 'POST',
|
||||
headers: {
|
||||
'Content-Type': 'application/json',
|
||||
},
|
||||
body: JSON.stringify({ texts: texts })
|
||||
});
|
||||
|
||||
const data = await response.json();
|
||||
|
||||
if (data.success) {
|
||||
showAlert('success', `成功上传 ${data.uploaded_count} 条文本`);
|
||||
document.getElementById('batchTextInput').value = '';
|
||||
updateDataStats();
|
||||
document.getElementById('buildIndexBtn').disabled = false;
|
||||
} else {
|
||||
showAlert('danger', `上传失败: ${data.message}`);
|
||||
}
|
||||
} catch (error) {
|
||||
showAlert('danger', `上传错误: ${error.message}`);
|
||||
}
|
||||
}
|
||||
|
||||
// 构建索引
|
||||
async function buildIndex() {
|
||||
const btn = document.getElementById('buildIndexBtn');
|
||||
const originalText = btn.innerHTML;
|
||||
btn.innerHTML = '<i class="fas fa-spinner fa-spin"></i> 构建中...';
|
||||
btn.disabled = true;
|
||||
|
||||
try {
|
||||
const response = await fetch('/api/build_index', {
|
||||
method: 'POST'
|
||||
});
|
||||
|
||||
const data = await response.json();
|
||||
|
||||
if (data.success) {
|
||||
showAlert('success', '索引构建完成！现在可以进行搜索了');
|
||||
} else {
|
||||
showAlert('danger', `索引构建失败: ${data.message}`);
|
||||
}
|
||||
} catch (error) {
|
||||
showAlert('danger', `构建错误: ${error.message}`);
|
||||
} finally {
|
||||
btn.innerHTML = originalText;
|
||||
btn.disabled = false;
|
||||
}
|
||||
}
|
||||
|
||||
// 查看数据
|
||||
async function viewData() {
|
||||
try {
|
||||
const response = await fetch('/api/data/list');
|
||||
const data = await response.json();
|
||||
|
||||
if (data.success) {
|
||||
let content = '<div class="row">';
|
||||
|
||||
// 显示图片数据
|
||||
if (data.images && data.images.length > 0) {
|
||||
content += '<div class="col-md-6"><h6>图片数据 (' + data.images.length + ')</h6>';
|
||||
content += '<div class="list-group" style="max-height: 300px; overflow-y: auto;">';
|
||||
data.images.forEach(img => {
|
||||
content += `<div class="list-group-item d-flex justify-content-between align-items-center">
|
||||
<span>${img}</span>
|
||||
<img src="/uploads/${img}" class="img-thumbnail" style="width: 50px; height: 50px; object-fit: cover;">
|
||||
</div>`;
|
||||
});
|
||||
content += '</div></div>';
|
||||
}
|
||||
|
||||
// 显示文本数据
|
||||
if (data.texts && data.texts.length > 0) {
|
||||
content += '<div class="col-md-6"><h6>文本数据 (' + data.texts.length + ')</h6>';
|
||||
content += '<div class="list-group" style="max-height: 300px; overflow-y: auto;">';
|
||||
data.texts.forEach((text, index) => {
|
||||
const shortText = text.length > 50 ? text.substring(0, 50) + '...' : text;
|
||||
content += `<div class="list-group-item">
|
||||
<small class="text-muted">#${index + 1}</small><br>
|
||||
${shortText}
|
||||
</div>`;
|
||||
});
|
||||
content += '</div></div>';
|
||||
}
|
||||
|
||||
content += '</div>';
|
||||
|
||||
showModal('数据列表', content);
|
||||
} else {
|
||||
showAlert('danger', `获取数据失败: ${data.message}`);
|
||||
}
|
||||
} catch (error) {
|
||||
showAlert('danger', `获取数据错误: ${error.message}`);
|
||||
}
|
||||
}
|
||||
|
||||
// 清空数据
|
||||
async function clearData() {
|
||||
if (!confirm('确定要清空所有数据吗？此操作不可恢复！')) {
|
||||
return;
|
||||
}
|
||||
|
||||
try {
|
||||
const response = await fetch('/api/data/clear', {
|
||||
method: 'POST'
|
||||
});
|
||||
|
||||
const data = await response.json();
|
||||
|
||||
if (data.success) {
|
||||
showAlert('success', '数据已清空');
|
||||
updateDataStats();
|
||||
document.getElementById('buildIndexBtn').disabled = true;
|
||||
} else {
|
||||
showAlert('danger', `清空失败: ${data.message}`);
|
||||
}
|
||||
} catch (error) {
|
||||
showAlert('danger', `清空错误: ${error.message}`);
|
||||
}
|
||||
}
|
||||
|
||||
// 更新数据统计
|
||||
async function updateDataStats() {
|
||||
try {
|
||||
const response = await fetch('/api/data/stats');
|
||||
const data = await response.json();
|
||||
|
||||
if (data.success) {
|
||||
document.getElementById('imageCount').textContent = data.image_count || 0;
|
||||
document.getElementById('textCount').textContent = data.text_count || 0;
|
||||
}
|
||||
} catch (error) {
|
||||
console.log('获取数据统计失败:', error);
|
||||
}
|
||||
}
|
||||
|
||||
// 显示模态框
|
||||
function showModal(title, content) {
|
||||
const modalHtml = `
|
||||
<div class="modal fade" id="dataModal" tabindex="-1">
|
||||
<div class="modal-dialog modal-lg">
|
||||
<div class="modal-content">
|
||||
<div class="modal-header">
|
||||
<h5 class="modal-title">${title}</h5>
|
||||
<button type="button" class="btn-close" data-bs-dismiss="modal"></button>
|
||||
</div>
|
||||
<div class="modal-body">
|
||||
${content}
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
`;
|
||||
|
||||
// 移除已存在的模态框
|
||||
const existingModal = document.getElementById('dataModal');
|
||||
if (existingModal) {
|
||||
existingModal.remove();
|
||||
}
|
||||
|
||||
document.body.insertAdjacentHTML('beforeend', modalHtml);
|
||||
const modal = new bootstrap.Modal(document.getElementById('dataModal'));
|
||||
modal.show();
|
||||
}
|
||||
</script>
|
||||
</body>
|
||||
</html>
|
||||
104
test_all_retrieval_modes.py
Normal file
104
test_all_retrieval_modes.py
Normal file
@ -0,0 +1,104 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
测试所有四种多模态检索模式
|
||||
"""
|
||||
|
||||
from multimodal_retrieval_multigpu import MultiGPUMultimodalRetrieval
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
import os
|
||||
|
||||
def test_all_retrieval_modes():
|
||||
print('正在初始化多GPU多模态检索系统...')
|
||||
retrieval = MultiGPUMultimodalRetrieval()
|
||||
|
||||
# 准备测试数据
|
||||
test_texts = [
|
||||
"一只可爱的小猫",
|
||||
"美丽的风景照片",
|
||||
"现代建筑设计",
|
||||
"colorful flowers in garden"
|
||||
]
|
||||
|
||||
test_images = [
|
||||
'sample_images/1755677101_1__.jpg',
|
||||
'sample_images/1755677101_2__.jpg',
|
||||
'sample_images/1755677101_3__.jpg',
|
||||
'sample_images/1755677101_4__.jpg'
|
||||
]
|
||||
|
||||
# 验证测试图像存在
|
||||
existing_images = [img for img in test_images if os.path.exists(img)]
|
||||
if not existing_images:
|
||||
print("❌ 没有找到测试图像文件")
|
||||
return
|
||||
|
||||
print(f"找到 {len(existing_images)} 张测试图像")
|
||||
|
||||
try:
|
||||
# 1. 构建文本索引
|
||||
print('\n=== 构建文本索引 ===')
|
||||
retrieval.build_text_index_parallel(test_texts)
|
||||
print('✅ 文本索引构建完成')
|
||||
|
||||
# 2. 构建图像索引
|
||||
print('\n=== 构建图像索引 ===')
|
||||
retrieval.build_image_index_parallel(existing_images)
|
||||
print('✅ 图像索引构建完成')
|
||||
|
||||
# 3. 测试文本到文本检索
|
||||
print('\n=== 测试文本到文本检索 ===')
|
||||
query = "小动物"
|
||||
results = retrieval.search_text_by_text(query, top_k=3)
|
||||
print(f'查询: "{query}"')
|
||||
for i, (text, score) in enumerate(results):
|
||||
print(f' {i+1}. {text} (相似度: {score:.4f})')
|
||||
|
||||
# 4. 测试文本到图像检索
|
||||
print('\n=== 测试文本到图像检索 ===')
|
||||
query = "beautiful image"
|
||||
results = retrieval.search_images_by_text(query, top_k=3)
|
||||
print(f'查询: "{query}"')
|
||||
for i, (image_path, score) in enumerate(results):
|
||||
print(f' {i+1}. {image_path} (相似度: {score:.4f})')
|
||||
|
||||
# 5. 测试图像到文本检索
|
||||
print('\n=== 测试图像到文本检索 ===')
|
||||
query_image = existing_images[0]
|
||||
results = retrieval.search_text_by_image(query_image, top_k=3)
|
||||
print(f'查询图像: {query_image}')
|
||||
for i, (text, score) in enumerate(results):
|
||||
print(f' {i+1}. {text} (相似度: {score:.4f})')
|
||||
|
||||
# 6. 测试图像到图像检索
|
||||
print('\n=== 测试图像到图像检索 ===')
|
||||
query_image = existing_images[0]
|
||||
results = retrieval.search_images_by_image(query_image, top_k=3)
|
||||
print(f'查询图像: {query_image}')
|
||||
for i, (image_path, score) in enumerate(results):
|
||||
print(f' {i+1}. {image_path} (相似度: {score:.4f})')
|
||||
|
||||
print('\n✅ 所有四种检索模式测试完成！')
|
||||
|
||||
# 7. 测试Web应用兼容的方法名
|
||||
print('\n=== 测试Web应用兼容方法 ===')
|
||||
try:
|
||||
results = retrieval.search_text_to_image("test query", top_k=2)
|
||||
print('✅ search_text_to_image 方法正常')
|
||||
|
||||
results = retrieval.search_image_to_text(existing_images[0], top_k=2)
|
||||
print('✅ search_image_to_text 方法正常')
|
||||
|
||||
results = retrieval.search_image_to_image(existing_images[0], top_k=2)
|
||||
print('✅ search_image_to_image 方法正常')
|
||||
|
||||
except Exception as e:
|
||||
print(f'❌ Web应用兼容方法测试失败: {e}')
|
||||
|
||||
except Exception as e:
|
||||
print(f'❌ 测试过程中出现错误: {e}')
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_all_retrieval_modes()
|
||||
49
test_image_encoding.py
Normal file
49
test_image_encoding.py
Normal file
@ -0,0 +1,49 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
测试图像编码功能
|
||||
"""
|
||||
|
||||
from multimodal_retrieval_multigpu import MultiGPUMultimodalRetrieval
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
def test_image_encoding():
|
||||
print('正在初始化多GPU多模态检索系统...')
|
||||
retrieval = MultiGPUMultimodalRetrieval()
|
||||
|
||||
# 测试文本编码
|
||||
print('测试文本编码...')
|
||||
text_embeddings = retrieval.encode_text_batch(['这是一个测试文本'])
|
||||
print(f'文本embedding形状: {text_embeddings.shape}')
|
||||
print(f'文本embedding数据类型: {text_embeddings.dtype}')
|
||||
|
||||
# 测试图像编码
|
||||
print('测试图像编码...')
|
||||
test_images = ['sample_images/1755677101_1__.jpg']
|
||||
image_embeddings = retrieval.encode_image_batch(test_images)
|
||||
print(f'图像embedding形状: {image_embeddings.shape}')
|
||||
print(f'图像embedding数据类型: {image_embeddings.dtype}')
|
||||
|
||||
# 测试两次相同图像的embedding是否一致
|
||||
print('测试embedding一致性...')
|
||||
image_embeddings2 = retrieval.encode_image_batch(test_images)
|
||||
consistency = np.allclose(image_embeddings, image_embeddings2, rtol=1e-5)
|
||||
print(f'相同图像embedding一致性: {consistency}')
|
||||
|
||||
# 测试不同图像的embedding差异
|
||||
print('测试不同图像embedding差异...')
|
||||
test_images2 = ['sample_images/1755677101_2__.jpg']
|
||||
image_embeddings3 = retrieval.encode_image_batch(test_images2)
|
||||
similarity = np.dot(image_embeddings[0], image_embeddings3[0]) / (np.linalg.norm(image_embeddings[0]) * np.linalg.norm(image_embeddings3[0]))
|
||||
print(f'不同图像间相似度: {similarity:.4f}')
|
||||
|
||||
# 验证维度一致性
|
||||
if text_embeddings.shape[1] == image_embeddings.shape[1]:
|
||||
print('✅ 文本和图像embedding维度一致')
|
||||
else:
|
||||
print('❌ 文本和图像embedding维度不一致')
|
||||
|
||||
print('测试完成！')
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_image_encoding()
|
||||
BIN
uploads/query_1755675080_10__.jpg
Normal file
BIN
uploads/query_1755675080_10__.jpg
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 172 KiB |
BIN
uploads/query_1755681423_5__.jpg
Normal file
BIN
uploads/query_1755681423_5__.jpg
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 201 KiB |
617
web_app_multigpu.py
Normal file
617
web_app_multigpu.py
Normal file
@ -0,0 +1,617 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
多GPU多模态检索系统 - Web应用
|
||||
专为双GPU部署优化
|
||||
"""
|
||||
|
||||
import os
|
||||
import json
|
||||
import time
|
||||
from flask import Flask, render_template, request, jsonify, send_file, url_for
|
||||
from werkzeug.utils import secure_filename
|
||||
from PIL import Image
|
||||
import base64
|
||||
import io
|
||||
import logging
|
||||
import traceback
|
||||
import glob
|
||||
|
||||
# 设置环境变量优化GPU内存
|
||||
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'
|
||||
|
||||
# 配置日志
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
app = Flask(__name__)
|
||||
app.config['SECRET_KEY'] = 'multigpu_multimodal_retrieval_2024'
|
||||
app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024 # 16MB max file size
|
||||
|
||||
# 配置上传文件夹
|
||||
UPLOAD_FOLDER = 'uploads'
|
||||
SAMPLE_IMAGES_FOLDER = 'sample_images'
|
||||
TEXT_DATA_FOLDER = 'text_data'
|
||||
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif', 'bmp', 'webp'}
|
||||
|
||||
# 确保文件夹存在
|
||||
os.makedirs(UPLOAD_FOLDER, exist_ok=True)
|
||||
os.makedirs(SAMPLE_IMAGES_FOLDER, exist_ok=True)
|
||||
os.makedirs(TEXT_DATA_FOLDER, exist_ok=True)
|
||||
|
||||
# 全局检索系统实例
|
||||
retrieval_system = None
|
||||
|
||||
def allowed_file(filename):
|
||||
"""检查文件扩展名是否允许"""
|
||||
return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
|
||||
|
||||
def image_to_base64(image_path):
|
||||
"""将图片转换为base64编码"""
|
||||
try:
|
||||
with open(image_path, "rb") as img_file:
|
||||
return base64.b64encode(img_file.read()).decode('utf-8')
|
||||
except Exception as e:
|
||||
logger.error(f"图片转换失败: {e}")
|
||||
return None
|
||||
|
||||
@app.route('/')
|
||||
def index():
|
||||
"""主页"""
|
||||
return render_template('index.html')
|
||||
|
||||
@app.route('/api/status')
|
||||
def get_status():
|
||||
"""获取系统状态"""
|
||||
global retrieval_system
|
||||
|
||||
status = {
|
||||
'initialized': retrieval_system is not None,
|
||||
'gpu_count': 0,
|
||||
'model_loaded': False
|
||||
}
|
||||
|
||||
try:
|
||||
import torch
|
||||
if torch.cuda.is_available():
|
||||
status['gpu_count'] = torch.cuda.device_count()
|
||||
|
||||
if retrieval_system and retrieval_system.model:
|
||||
status['model_loaded'] = True
|
||||
status['device_ids'] = retrieval_system.device_ids
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"获取状态失败: {e}")
|
||||
|
||||
return jsonify(status)
|
||||
|
||||
@app.route('/api/init', methods=['POST'])
|
||||
def initialize_system():
|
||||
"""初始化多GPU检索系统"""
|
||||
global retrieval_system
|
||||
|
||||
try:
|
||||
logger.info("正在初始化多GPU检索系统...")
|
||||
|
||||
# 导入多GPU检索系统
|
||||
from multimodal_retrieval_multigpu import MultiGPUMultimodalRetrieval
|
||||
|
||||
# 初始化系统
|
||||
retrieval_system = MultiGPUMultimodalRetrieval()
|
||||
|
||||
if retrieval_system.model is None:
|
||||
raise Exception("模型加载失败")
|
||||
|
||||
logger.info("✅ 多GPU系统初始化成功")
|
||||
|
||||
return jsonify({
|
||||
'success': True,
|
||||
'message': '多GPU系统初始化成功',
|
||||
'device_ids': retrieval_system.device_ids,
|
||||
'gpu_count': len(retrieval_system.device_ids)
|
||||
})
|
||||
|
||||
except Exception as e:
|
||||
error_msg = f"系统初始化失败: {str(e)}"
|
||||
logger.error(error_msg)
|
||||
logger.error(traceback.format_exc())
|
||||
|
||||
return jsonify({
|
||||
'success': False,
|
||||
'message': error_msg
|
||||
}), 500
|
||||
|
||||
@app.route('/api/search/text_to_text', methods=['POST'])
|
||||
def search_text_to_text():
|
||||
"""文本搜索文本"""
|
||||
return handle_search('text_to_text')
|
||||
|
||||
@app.route('/api/search/text_to_image', methods=['POST'])
|
||||
def search_text_to_image():
|
||||
"""文本搜索图片"""
|
||||
return handle_search('text_to_image')
|
||||
|
||||
@app.route('/api/search/image_to_text', methods=['POST'])
|
||||
def search_image_to_text():
|
||||
"""图片搜索文本"""
|
||||
return handle_search('image_to_text')
|
||||
|
||||
@app.route('/api/search/image_to_image', methods=['POST'])
|
||||
def search_image_to_image():
|
||||
"""图片搜索图片"""
|
||||
return handle_search('image_to_image')
|
||||
|
||||
@app.route('/api/search', methods=['POST'])
|
||||
def search():
|
||||
"""通用搜索接口（兼容旧版本）"""
|
||||
mode = request.form.get('mode') or request.json.get('mode', 'text_to_text')
|
||||
return handle_search(mode)
|
||||
|
||||
def handle_search(mode):
|
||||
"""处理搜索请求的通用函数"""
|
||||
global retrieval_system
|
||||
|
||||
if not retrieval_system:
|
||||
return jsonify({
|
||||
'success': False,
|
||||
'message': '系统未初始化，请先点击初始化按钮'
|
||||
}), 400
|
||||
|
||||
try:
|
||||
top_k = int(request.form.get('top_k', 5))
|
||||
|
||||
if mode in ['text_to_text', 'text_to_image']:
|
||||
# 文本查询
|
||||
query = request.form.get('query') or request.json.get('query', '')
|
||||
if not query.strip():
|
||||
return jsonify({
|
||||
'success': False,
|
||||
'message': '请输入查询文本'
|
||||
}), 400
|
||||
|
||||
logger.info(f"执行{mode}搜索: {query}")
|
||||
|
||||
# 执行搜索
|
||||
if mode == 'text_to_text':
|
||||
results = retrieval_system.search_text_to_text(query, top_k=top_k)
|
||||
else: # text_to_image
|
||||
results = retrieval_system.search_text_to_image(query, top_k=top_k)
|
||||
|
||||
return jsonify({
|
||||
'success': True,
|
||||
'mode': mode,
|
||||
'query': query,
|
||||
'results': results,
|
||||
'result_count': len(results)
|
||||
})
|
||||
|
||||
elif mode in ['image_to_text', 'image_to_image']:
|
||||
# 图片查询
|
||||
if 'image' not in request.files:
|
||||
return jsonify({
|
||||
'success': False,
|
||||
'message': '请上传查询图片'
|
||||
}), 400
|
||||
|
||||
file = request.files['image']
|
||||
if file.filename == '' or not allowed_file(file.filename):
|
||||
return jsonify({
|
||||
'success': False,
|
||||
'message': '请上传有效的图片文件'
|
||||
}), 400
|
||||
|
||||
# 保存上传的图片
|
||||
filename = secure_filename(file.filename)
|
||||
timestamp = str(int(time.time()))
|
||||
filename = f"query_{timestamp}_{filename}"
|
||||
filepath = os.path.join(UPLOAD_FOLDER, filename)
|
||||
file.save(filepath)
|
||||
|
||||
logger.info(f"执行{mode}搜索，图片: {filename}")
|
||||
|
||||
# 执行搜索
|
||||
if mode == 'image_to_text':
|
||||
results = retrieval_system.search_image_to_text(filepath, top_k=top_k)
|
||||
else: # image_to_image
|
||||
results = retrieval_system.search_image_to_image(filepath, top_k=top_k)
|
||||
|
||||
# 转换查询图片为base64
|
||||
query_image_b64 = image_to_base64(filepath)
|
||||
|
||||
return jsonify({
|
||||
'success': True,
|
||||
'mode': mode,
|
||||
'query_image': query_image_b64,
|
||||
'results': results,
|
||||
'result_count': len(results)
|
||||
})
|
||||
|
||||
else:
|
||||
return jsonify({
|
||||
'success': False,
|
||||
'message': f'不支持的搜索模式: {mode}'
|
||||
}), 400
|
||||
|
||||
except Exception as e:
|
||||
error_msg = f"搜索失败: {str(e)}"
|
||||
logger.error(error_msg)
|
||||
logger.error(traceback.format_exc())
|
||||
|
||||
return jsonify({
|
||||
'success': False,
|
||||
'message': error_msg
|
||||
}), 500
|
||||
|
||||
@app.route('/api/upload/images', methods=['POST'])
|
||||
def upload_images():
|
||||
"""批量上传图片"""
|
||||
try:
|
||||
uploaded_files = []
|
||||
|
||||
if 'images' not in request.files:
|
||||
return jsonify({'success': False, 'message': '没有选择文件'}), 400
|
||||
|
||||
files = request.files.getlist('images')
|
||||
|
||||
for file in files:
|
||||
if file and file.filename != '' and allowed_file(file.filename):
|
||||
filename = secure_filename(file.filename)
|
||||
timestamp = str(int(time.time()))
|
||||
filename = f"{timestamp}_{filename}"
|
||||
filepath = os.path.join(SAMPLE_IMAGES_FOLDER, filename)
|
||||
file.save(filepath)
|
||||
uploaded_files.append(filename)
|
||||
|
||||
return jsonify({
|
||||
'success': True,
|
||||
'message': f'成功上传 {len(uploaded_files)} 个图片文件',
|
||||
'uploaded_count': len(uploaded_files),
|
||||
'files': uploaded_files
|
||||
})
|
||||
|
||||
except Exception as e:
|
||||
return jsonify({
|
||||
'success': False,
|
||||
'message': f'上传失败: {str(e)}'
|
||||
}), 500
|
||||
|
||||
@app.route('/api/upload/texts', methods=['POST'])
|
||||
def upload_texts():
|
||||
"""批量上传文本数据"""
|
||||
try:
|
||||
data = request.get_json()
|
||||
|
||||
if not data or 'texts' not in data:
|
||||
return jsonify({'success': False, 'message': '没有提供文本数据'}), 400
|
||||
|
||||
texts = data['texts']
|
||||
if not isinstance(texts, list):
|
||||
return jsonify({'success': False, 'message': '文本数据格式错误'}), 400
|
||||
|
||||
# 保存文本数据到文件
|
||||
timestamp = str(int(time.time()))
|
||||
filename = f"texts_{timestamp}.json"
|
||||
filepath = os.path.join(TEXT_DATA_FOLDER, filename)
|
||||
|
||||
with open(filepath, 'w', encoding='utf-8') as f:
|
||||
json.dump(texts, f, ensure_ascii=False, indent=2)
|
||||
|
||||
return jsonify({
|
||||
'success': True,
|
||||
'message': f'成功上传 {len(texts)} 条文本',
|
||||
'uploaded_count': len(texts)
|
||||
})
|
||||
|
||||
except Exception as e:
|
||||
return jsonify({
|
||||
'success': False,
|
||||
'message': f'上传失败: {str(e)}'
|
||||
}), 500
|
||||
|
||||
@app.route('/api/upload/file', methods=['POST'])
|
||||
def upload_single_file():
|
||||
"""上传单个文件"""
|
||||
if 'file' not in request.files:
|
||||
return jsonify({'success': False, 'message': '没有选择文件'}), 400
|
||||
|
||||
file = request.files['file']
|
||||
if file.filename == '':
|
||||
return jsonify({'success': False, 'message': '没有选择文件'}), 400
|
||||
|
||||
if file and allowed_file(file.filename):
|
||||
filename = secure_filename(file.filename)
|
||||
timestamp = str(int(time.time()))
|
||||
filename = f"{timestamp}_{filename}"
|
||||
filepath = os.path.join(SAMPLE_IMAGES_FOLDER, filename)
|
||||
file.save(filepath)
|
||||
|
||||
return jsonify({
|
||||
'success': True,
|
||||
'message': '文件上传成功',
|
||||
'filename': filename
|
||||
})
|
||||
|
||||
return jsonify({'success': False, 'message': '不支持的文件类型'}), 400
|
||||
|
||||
@app.route('/api/data/list', methods=['GET'])
|
||||
def list_data():
|
||||
"""列出已上传的数据"""
|
||||
try:
|
||||
# 列出图片文件
|
||||
images = []
|
||||
if os.path.exists(SAMPLE_IMAGES_FOLDER):
|
||||
for filename in os.listdir(SAMPLE_IMAGES_FOLDER):
|
||||
if allowed_file(filename):
|
||||
filepath = os.path.join(SAMPLE_IMAGES_FOLDER, filename)
|
||||
stat = os.stat(filepath)
|
||||
images.append({
|
||||
'filename': filename,
|
||||
'size': stat.st_size,
|
||||
'modified': stat.st_mtime
|
||||
})
|
||||
|
||||
# 列出文本文件
|
||||
texts = []
|
||||
if os.path.exists(TEXT_DATA_FOLDER):
|
||||
for filename in os.listdir(TEXT_DATA_FOLDER):
|
||||
if filename.endswith('.json'):
|
||||
filepath = os.path.join(TEXT_DATA_FOLDER, filename)
|
||||
stat = os.stat(filepath)
|
||||
texts.append({
|
||||
'filename': filename,
|
||||
'size': stat.st_size,
|
||||
'modified': stat.st_mtime
|
||||
})
|
||||
|
||||
return jsonify({
|
||||
'success': True,
|
||||
'data': {
|
||||
'images': images,
|
||||
'texts': texts
|
||||
}
|
||||
})
|
||||
|
||||
except Exception as e:
|
||||
return jsonify({
|
||||
'success': False,
|
||||
'message': f'获取数据列表失败: {str(e)}'
|
||||
}), 500
|
||||
|
||||
@app.route('/api/gpu_status')
|
||||
def gpu_status():
|
||||
"""获取GPU状态"""
|
||||
try:
|
||||
from smart_gpu_launcher import get_gpu_memory_info
|
||||
gpu_info = get_gpu_memory_info()
|
||||
|
||||
return jsonify({
|
||||
'success': True,
|
||||
'gpu_info': gpu_info
|
||||
})
|
||||
|
||||
except Exception as e:
|
||||
return jsonify({
|
||||
'success': False,
|
||||
'message': f"获取GPU状态失败: {str(e)}"
|
||||
}), 500
|
||||
|
||||
@app.route('/api/build_index', methods=['POST'])
|
||||
def build_index():
|
||||
"""构建检索索引"""
|
||||
global retrieval_system
|
||||
|
||||
if not retrieval_system:
|
||||
return jsonify({
|
||||
'success': False,
|
||||
'message': '系统未初始化'
|
||||
}), 400
|
||||
|
||||
try:
|
||||
# 获取所有图片和文本文件
|
||||
image_files = []
|
||||
text_data = []
|
||||
|
||||
# 扫描图片文件
|
||||
for ext in ALLOWED_EXTENSIONS:
|
||||
pattern = os.path.join(SAMPLE_IMAGES_FOLDER, f"*.{ext}")
|
||||
image_files.extend(glob.glob(pattern))
|
||||
|
||||
# 读取文本文件（支持.json和.txt格式）
|
||||
text_files = glob.glob(os.path.join(TEXT_DATA_FOLDER, "*.json"))
|
||||
text_files.extend(glob.glob(os.path.join(TEXT_DATA_FOLDER, "*.txt")))
|
||||
|
||||
for text_file in text_files:
|
||||
try:
|
||||
if text_file.endswith('.json'):
|
||||
# 读取JSON格式的文本数据
|
||||
with open(text_file, 'r', encoding='utf-8') as f:
|
||||
data = json.load(f)
|
||||
if isinstance(data, list):
|
||||
text_data.extend([str(item).strip() for item in data if str(item).strip()])
|
||||
else:
|
||||
text_data.append(str(data).strip())
|
||||
else:
|
||||
# 读取TXT格式的文本数据
|
||||
with open(text_file, 'r', encoding='utf-8') as f:
|
||||
lines = [line.strip() for line in f.readlines() if line.strip()]
|
||||
text_data.extend(lines)
|
||||
except Exception as e:
|
||||
logger.warning(f"读取文本文件失败 {text_file}: {e}")
|
||||
|
||||
# 检查是否有数据可以构建索引
|
||||
if not image_files and not text_data:
|
||||
return jsonify({
|
||||
'success': False,
|
||||
'message': '没有找到可用的图片或文本数据，请先上传数据'
|
||||
}), 400
|
||||
|
||||
# 构建索引
|
||||
if image_files:
|
||||
logger.info(f"构建图片索引，共 {len(image_files)} 张图片")
|
||||
retrieval_system.build_image_index_parallel(image_files)
|
||||
|
||||
if text_data:
|
||||
logger.info(f"构建文本索引，共 {len(text_data)} 条文本")
|
||||
retrieval_system.build_text_index_parallel(text_data)
|
||||
|
||||
return jsonify({
|
||||
'success': True,
|
||||
'message': f'索引构建完成！图片: {len(image_files)} 张，文本: {len(text_data)} 条',
|
||||
'image_count': len(image_files),
|
||||
'text_count': len(text_data)
|
||||
})
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"构建索引失败: {str(e)}")
|
||||
return jsonify({
|
||||
'success': False,
|
||||
'message': f'构建索引失败: {str(e)}'
|
||||
}), 500
|
||||
|
||||
@app.route('/api/data/stats', methods=['GET'])
|
||||
def get_data_stats():
|
||||
"""获取数据统计信息"""
|
||||
try:
|
||||
# 统计图片文件
|
||||
image_count = 0
|
||||
for ext in ALLOWED_EXTENSIONS:
|
||||
pattern = os.path.join(SAMPLE_IMAGES_FOLDER, f"*.{ext}")
|
||||
image_count += len(glob.glob(pattern))
|
||||
|
||||
# 统计文本数据
|
||||
text_count = 0
|
||||
text_files = glob.glob(os.path.join(TEXT_DATA_FOLDER, "*.txt"))
|
||||
for text_file in text_files:
|
||||
try:
|
||||
with open(text_file, 'r', encoding='utf-8') as f:
|
||||
lines = [line.strip() for line in f.readlines() if line.strip()]
|
||||
text_count += len(lines)
|
||||
except Exception:
|
||||
continue
|
||||
|
||||
return jsonify({
|
||||
'success': True,
|
||||
'image_count': image_count,
|
||||
'text_count': text_count
|
||||
})
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"获取数据统计失败: {str(e)}")
|
||||
return jsonify({
|
||||
'success': False,
|
||||
'message': f'获取统计失败: {str(e)}'
|
||||
}), 500
|
||||
|
||||
@app.route('/api/data/clear', methods=['POST'])
|
||||
def clear_data():
|
||||
"""清空所有数据"""
|
||||
try:
|
||||
# 清空图片文件
|
||||
for ext in ALLOWED_EXTENSIONS:
|
||||
pattern = os.path.join(SAMPLE_IMAGES_FOLDER, f"*.{ext}")
|
||||
for file_path in glob.glob(pattern):
|
||||
try:
|
||||
os.remove(file_path)
|
||||
except Exception as e:
|
||||
logger.warning(f"删除图片文件失败 {file_path}: {e}")
|
||||
|
||||
# 清空文本文件
|
||||
text_files = glob.glob(os.path.join(TEXT_DATA_FOLDER, "*.txt"))
|
||||
for text_file in text_files:
|
||||
try:
|
||||
os.remove(text_file)
|
||||
except Exception as e:
|
||||
logger.warning(f"删除文本文件失败 {text_file}: {e}")
|
||||
|
||||
# 重置索引
|
||||
global retrieval_system
|
||||
if retrieval_system:
|
||||
retrieval_system.text_index = None
|
||||
retrieval_system.image_index = None
|
||||
retrieval_system.text_data = []
|
||||
retrieval_system.image_data = []
|
||||
|
||||
return jsonify({
|
||||
'success': True,
|
||||
'message': '数据已清空'
|
||||
})
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"清空数据失败: {str(e)}")
|
||||
return jsonify({
|
||||
'success': False,
|
||||
'message': f'清空数据失败: {str(e)}'
|
||||
}), 500
|
||||
|
||||
@app.route('/uploads/<filename>')
|
||||
def uploaded_file(filename):
|
||||
"""提供上传文件的访问"""
|
||||
return send_file(os.path.join(SAMPLE_IMAGES_FOLDER, filename))
|
||||
|
||||
def print_startup_info():
|
||||
"""打印启动信息"""
|
||||
print("🚀 启动多GPU多模态检索Web应用")
|
||||
print("=" * 60)
|
||||
print("访问地址: http://localhost:5000")
|
||||
print("支持功能:")
|
||||
print(" 📝 文搜文 - 文本查找相似文本")
|
||||
print(" 🖼️ 文搜图 - 文本查找相关图片")
|
||||
print(" 📝 图搜文 - 图片查找相关文本")
|
||||
print(" 🖼️ 图搜图 - 图片查找相似图片")
|
||||
print(" 📤 批量上传 - 图片和文本数据管理")
|
||||
print("GPU配置:")
|
||||
|
||||
try:
|
||||
import torch
|
||||
if torch.cuda.is_available():
|
||||
gpu_count = torch.cuda.device_count()
|
||||
print(f" 🖥️ 检测到 {gpu_count} 个GPU")
|
||||
for i in range(gpu_count):
|
||||
name = torch.cuda.get_device_name(i)
|
||||
props = torch.cuda.get_device_properties(i)
|
||||
memory_gb = props.total_memory / 1024**3
|
||||
print(f" GPU {i}: {name} ({memory_gb:.1f}GB)")
|
||||
else:
|
||||
print(" ❌ CUDA不可用")
|
||||
except Exception as e:
|
||||
print(f" ❌ GPU检查失败: {e}")
|
||||
|
||||
print("=" * 60)
|
||||
|
||||
def auto_initialize():
|
||||
"""启动时自动初始化系统"""
|
||||
global retrieval_system
|
||||
|
||||
try:
|
||||
logger.info("🚀 启动时自动初始化多GPU检索系统...")
|
||||
|
||||
# 导入多GPU检索系统
|
||||
from multimodal_retrieval_multigpu import MultiGPUMultimodalRetrieval
|
||||
|
||||
# 初始化系统
|
||||
retrieval_system = MultiGPUMultimodalRetrieval()
|
||||
|
||||
if retrieval_system.model is None:
|
||||
raise Exception("模型加载失败")
|
||||
|
||||
logger.info("✅ 系统自动初始化成功")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"❌ 系统自动初始化失败: {str(e)}")
|
||||
logger.error(traceback.format_exc())
|
||||
return False
|
||||
|
||||
if __name__ == '__main__':
|
||||
print_startup_info()
|
||||
|
||||
# 启动时自动初始化
|
||||
auto_initialize()
|
||||
|
||||
# 启动Flask应用
|
||||
app.run(
|
||||
host='0.0.0.0',
|
||||
port=5000,
|
||||
debug=False,
|
||||
threaded=True
|
||||
)
|
||||
Loading…
x
Reference in New Issue
Block a user