理解你的需求后，可以通过配置文件来管理多个图像目录，这样用户可以直接编辑配置文件来添加或删除图像路径，而不需要通过前端上传文件。以下是一个完整的实现方案，包括配置文件的读取、特征提取、索引构建和查询处理。
1. 配置文件
首先，创建一个配置文件 config.ini 来管理图像目录路径。
[directories]
image_dirs = D:\02人能智项目\04以图找图项目\query_image;D:\02人能智项目\04以图找图项目\data;D:\other_images
2. 读取配置文件
使用 configparser 模块来读取配置文件中的图像目录路径。
import configparser
def read_config(config_path):
config = configparser.ConfigParser()
config.read(config_path)
image_dirs = config.get('directories', 'image_dirs').split(';')
return image_dirs
3. 特征提取
读取配置文件中的所有图像目录，提取图像特征并保存。
import os
import torch
import clip
from PIL import Image
import faiss
import numpy as np
# 加载预训练的CLIP模型
device = "cuda" if torch.cuda.is_available() else "cpu"
model, preprocess = clip.load("ViT-B/32", device=device)
def extract_image_features(image_dirs):
all_features = []
all_image_paths = []
for image_dir in image_dirs:
if not os.path.exists(image_dir):
continue
image_paths = [os.path.join(image_dir, f) for f in os.listdir(image_dir) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
for path in image_paths:
try:
image = preprocess(Image.open(path)).unsqueeze(0).to(device)
with torch.no_grad():
feature = model.encode_image(image).cpu().numpy()
all_features.append(feature)
all_image_paths.append(path)
except Exception as e:
print(f"Error processing {path}: {e}")
return np.vstack(all_features), all_image_paths
# 读取配置文件
config_path = "config.ini"
image_dirs = read_config(config_path)
# 提取图像特征并保存
features, image_paths = extract_image_features(image_dirs)
np.save("image_features.npy", features)
np.save("image_paths.npy", image_paths)
4. 索引构建
构建 Faiss 索引并保存。
# 加载特征和路径
features = np.load("image_features.npy")
image_paths = np.load("image_paths.npy")
# 构建Faiss索引
dimension = features.shape[1]
index = faiss.IndexFlatL2(dimension)
index.add(features)
# 保存索引
faiss.write_index(index, "image_index.index")
5. 查询处理
创建一个 Flask 应用来处理查询请求。
from flask import Flask, request, jsonify
import faiss
import numpy as np
app = Flask(__name__)
# 加载索引和路径
index = faiss.read_index("image_index.index")
image_paths = np.load("image_paths.npy")
# 提取文本特征
def extract_text_features(text):
text_input = clip.tokenize([text]).to(device)
with torch.no_grad():
feature = model.encode_text(text_input).cpu().numpy()
return feature
# 处理查询
@app.route("/search", methods=["POST"])
def search():
data = request.json
query = data["query"]
text_feature = extract_text_features(query)
_, indices = index.search(text_feature, 5)
results = [image_paths[i] for i in indices[0]]
return jsonify({"results": results})
if __name__ == "__main__":
app.run(debug=True)
6. 前端示例
前端部分保持不变，仍然使用简单的 HTML 和 JavaScript 来创建查询界面。
<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>以文搜图</title>
</head>
<body>
<h1>以文搜图</h1>
<input type="text" id="query" placeholder="请输入查询文本">
<button onclick="search()">搜索</button>
<div id="results"></div>
<script>
async function search() {
const query = document.getElementById("query").value;
const response = await fetch("/search", {
method: "POST",
headers: {
"Content-Type": "application/json"
},
body: JSON.stringify({ query })
});
const data = await response.json();
const resultsDiv = document.getElementById("results");
resultsDiv.innerHTML = "";
data.results.forEach(result => {
const img = document.createElement("img");
img.src = result;
resultsDiv.appendChild(img);
});
}
</script>
</body>
</html>
总结
通过配置文件来管理多个图像目录，用户可以直接编辑配置文件来添加或删除图像路径。后端会读取配置文件中的路径，提取图像特征并构建索引。这样可以避免通过前端上传大量文件，提高运行效率。希望这个方案能满足你的需求！
