图文检索在生活中具有广泛的应用，常见的图片检索包括基于文本内容搜索和基于图片内容搜索。用户通过输入文字描述或上传图片就可以在海量的图片库中快速找到同款或者相似图片，这种搜索方式被广泛应用于电商、广告、设计以及搜索引擎等热门领域。
本文基于火山引擎云搜索服务 ESCloud 和图文特征提取模型 CLIP，快速搭建一套以图搜图，以文搜图的端到端解决方案。
原理介绍
图片搜索技术，以文本描述和图片作为检索对象，分别对 image 和 text 进行特征提取，并在模型中对文本和图片建立相关联系，然后在海量图片数据库进行特征向量检索，返回与检索对象最相关的记录集合。其中特征提取部分采用 CLIP 模型，向量检索使用火山引擎云搜索服务在海量图片特征中进行快速的搜索。
环境依赖准备
登录火山引擎云搜索服务，创建实例集群，集群版本选择 7.10。
Python Client 关键依赖准备
pip install -U sentence-transformers # 模型相关pip install -U elasticsearch7==7.10.1 # ES向量数据库相关pip install -U pandas #分析splash的csv
复制代码
数据集准备
我们选择 Unsplash 作为图片数据集，详细介绍请参考：https://unsplash.com/data。在此示例中，我们选择下载 Lite 数据集，其中包含约 25,000 张照片。下载完成后会获得一个压缩文件，其中包含描述图片的 CSV 文件。通过使用 Pandas 读取 CSV 文件，我们将获得图片的 URL 地址。
def read_imgset(): path = '${下载的数据集所在路径}' documents = ['photos', 'keywords', 'collections', 'conversions', 'colors'] datasets = {} for doc in documents: files = glob.glob(path + doc + ".tsv*") subsets = [] for filename in files: # pd 分析csv df = pd.read_csv(filename, sep='\t', header=0) subsets.append(df) datasets[doc] = pd.concat(subsets, axis=0, ignore_index=True) return datasets
复制代码
模型选型
本文选取clip-ViT-B-32作为 以图搜图、以文搜图的模型，这个模型是基于 OpenAI 2021 论文的模型训练出来的，模型 CLIP 能将图片和文字联系在一起，目标是得到一个能同时表达图片和文字的模型。
ESCloud Mapping 准备
PUT image_search{ "mappings": { "dynamic": "false", "properties": { "photo_id": { "type": "keyword" }, "photo_url": { "type": "keyword" }, "describe": { "type": "text" }, "photo_embedding": { "type": "knn_vector", "dimension": 512 } } }, "settings": { "index": { "refresh_interval": "60s", "number_of_shards": "3", "knn.space_type": "cosinesimil", "knn": "true", "number_of_replicas": "1" } }}
复制代码
ESCloud 数据库操作
连接
登录火山引擎云搜索服务，选择刚刚创建好的实例，选择复制公网访问地址（如关闭，可选择开启）：
# 连接云搜索实例cloudSearch = CloudSearch("https://{user}:{password}@{ES_URL}", verify_certs=False, ssl_show_warn=False)
复制代码
写入
from sentence_transformers import SentenceTransformerfrom elasticsearch7 import Elasticsearch as CloudSearchfrom PIL import Imageimport requestsimport pandas as pdimport globfrom os.path import join
# We use the original clip-ViT-B-32 for encoding imagesimg_model = SentenceTransformer('clip-ViT-B-32')text_model = SentenceTransformer('clip-ViT-B-32-multilingual-v1')
# Construct request for esdef encodedataset(photo_id, photo_url, describe, image): encoded_sents = { "photo_id": photo_id, "photo_url": photo_url, "describe": describe, "photo_embedding": img_model.encode(image), } return encoded_sents
# download imagesdef load_image(url_or_path): if url_or_path.startswith("http://") or url_or_path.startswith("https://"): return Image.open(requests.get(url_or_path, stream=True).raw) else: return Image.open(url_or_path)
# 从unsplash的csv文件解出图片url，然后下载图片，# 下载完了后用model 生成embedding，并构造成ES的请求进行写入def get_imgset_and_bulk(): datasets = read_imgset() datasets['photos'].head() kwywords = datasets['keywords'] docs = [] #遍历CSV, 根据photo_url 去download photo for idx, row in datasets['photos'].iterrows(): print("Process id: ", idx) # 获取CSV 中的url photo_url = row["photo_image_url"] photo_id = row["photo_id"] image = load_image(photo_url) # 找到photo_id 且 suggested true 对应的图片描述 filter = kwywords.loc[(kwywords['photo_id'] == photo_id) & (kwywords['suggested_by_user'] == 't')] text = ' '.join(set(filter['keyword'])) # 封装写入ES的请求 one_document = encodedataset(photo_id=photo_id, photo_url=photo_url, describe=text, image=image) docs.append({"index": {}}) docs.append(one_document) if idx % 20 == 0: # 20条一组进行写入 resp = cloudSearch.bulk(docs, index='image_search') print(resp) docs = [] return docs if __name__ == '__main__': docs = get_imgset_and_bulk() print(docs)
复制代码
查询
以文搜图：文本向量化，执行 knn 查询
def extract_text(text): # 文搜图 res = cloudSearch.search( body={ "size": 5, "query": {"knn": {"photo_embedding": {"vector": text_model.encode(text), "k": 5}}}, "_source": ["describe", "photo_url"], }, index="image_search2", ) return res fe = FeatureExtractor()@app.route('/', methods=['GET', 'POST'])def index(): # ... resp = fe.extract_text(text)
return render_template('index.html', query_text=text, scores=resp['hits']['hits']) # ...
复制代码
搜 sunset 打印结果
以图搜图：图片向量化，执行 knn 查询
def extract(img): # 图搜图 res = cloudSearch.search( body={ "size": 5, "query": {"knn": {"photo_embedding": {"vector": img_model.encode(img), "k": 5}}}, "_source": ["describe", "photo_url"], }, index="image_search2", ) return res fe = FeatureExtractor()@app.route('/', methods=['GET', 'POST'])def index(): # ... # Save query image img = Image.open(file.stream) # PIL image uploaded_img_path = "static/uploaded/" + datetime.now().isoformat().replace(":", ".") + "_" + file.filename img.save(uploaded_img_path)
# Run search resp = fe.extract(img)
return render_template('index.html', query_path=uploaded_img_path, scores=resp['hits']['hits']) # ...
复制代码
搜海豹图片 打印结果
火山引擎云搜索服务 ESCloud 兼容 Elasticsearch、Kibana 等软件及常用开源插件，提供结构化、非结构化文本的多条件检索、统计、报表，可以实现一键部署、弹性扩缩、简化运维，快速构建日志分析、信息检索分析等业务能力。
了解更多产品详情：https://www.volcengine.com/product/es
评论