使用火山云搜索ESCloud服务构建图文检索应用（以文搜图/以图搜图）

云计算 2023-08-13 大树手机阅读

图文检索在生活中具有广泛的应用，常见的图片检索包括基于文本内容搜索和基于图片内容搜索。用户通过输入文字描述或上传图片就可以在海量的图片库中快速找到同款或者相似图片，这种搜索方式被广泛应用于电商、广告、设计以及搜索引擎等热门领域。

本文基于火山引擎云搜索服务 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 作为图片数据集，详细介绍请参考：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:
            1. 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 SentenceTransformer
from elasticsearch7 import Elasticsearch as CloudSearch
from PIL import Image
import requests
import pandas as pd
import glob
from os.path import join 

1. We use the original clip-ViT-B-32 for encoding images
img_model = SentenceTransformer('clip-ViT-B-32')
text_model = SentenceTransformer('clip-ViT-B-32-multilingual-v1')

1. Construct request for es
def 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

1. download images
def 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)

1. 从unsplash的csv文件解出图片url，然后下载图片，
1. 下载完了后用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)
        1. 获取CSV 中的url
        photo_url = row["photo_image_url"]
        photo_id  = row["photo_id"]
        image = load_image(photo_url)
        1. 找到photo_id 且 suggested true 对应的图片描述
        filter = kwywords.loc[(kwywords['photo_id'] == photo_id) & (kwywords['suggested_by_user'] == 't')]
        text = ' '.join(set(filter['keyword']))
        1. 封装写入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:
            1. 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):
    1. 文搜图
    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():
    1. ...
    resp = fe.extract_text(text)

    return render_template('index.html',
                       query_text=text,
                       scores=resp['hits']['hits'])
    1. ...

搜 sunset 打印结果

以图搜图：图片向量化，执行 knn 查询

def extract(img):
    1. 图搜图
    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():
    1. ...
    1. 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)

    1. Run search
    resp = fe.extract(img)

    return render_template('index.html',
                       query_path=uploaded_img_path,
                       scores=resp['hits']['hits'])
    1. ...

搜海豹图片打印结果