从零到一，教你搭建「CLIP 以文搜图」搜索服务（二）：5 分钟实现原型

郭人通

华中科技大学计算机软件与理论博士

在上一篇文章中，我们了解了关于搜索技术、以文搜图、以及 CLIP 模型的基础知识。本篇我们将花上5 分钟时间，对这些基础知识进行一次动手实践，快速构建一个「以文搜图」的搜索服务原型。

Notebook 链接

这里我们选取 “搜萌宠” 这个小例子：面对成千上万的萌宠，帮助用户快速在海量的图片中找到心仪的那只猫咪或狗狗

话不多说，先看 5 分钟出活儿的成品效果：

让我们来看看做这样一个原型都需要些什么：

一个宠物的小型图片库。
一个能将宠物图片的语义特征编码成向量的数据处理流水线。
一个能将查询文本的语义特征编码成向量的数据处理流水线。
一个可以支撑向量近邻搜索的向量数据库。
一段能将上述所有内容串起来的 python 脚本程序。

接下来，我们会陆续完成这张图的关键组件，开始干活~

安装基础工具包

我们用到了以下工具：

Towhee 用于构建模型推理流水线的框架，对于新手非常友好。
Faiss 高效的向量近邻搜索库。
Gradio 轻量级的机器学习 Demo 构建工具。

创建一个 conda 环境

conda create -n lovely_pet_retrieval python=3.9 
conda activate lovely_pet_retrieval

安装依赖

pip install towhee towhee.models gradio 
conda install -c pytorch faiss-cpu

准备图片库的数据

我们选取 ImageNet 数据集的子集作为本文所使用的 “小型宠物图片库”。首先，下载数据集并解压：

curl -L -O https://github.com/towhee-io/examples/releases/download/data/pet_small.zip 
unzip -q -o pet_small.zip

数据集的组织如下：

img: 包含 2500 张猫狗宠物图片
info.csv：包含 2500 张图片的基础信息，如图像的编号（id）、图片文件名（file_name）、以及类别（label）。

打开 jupyter notebook，简单观察一下数据：

import pandas as pd 
df = pd.read_csv('info.csv') 
df.head()

到这里，我们已经完成了关于图片库的准备工作。

将图片的特征编码成向量

我们通过 Towhee 调用 CLIP 模型推理来生成图像的 Embedding 向量：

import towhee 
img_vectors = ( 
    towhee.read_csv('info.csv') 
      .image_decode['file_name', 'img']() 
      .image_text_embedding.clip['img', 'vec'](model_name='clip_vit_b32', modality='image') 
      .tensor_normalize['vec','vec']() # normalize vector 
      .select['file_name', 'vec']() 
)

这里对代码做一些简要的说明：

read_csv('info.csv') 读取了三列数据到data collection，对应的 schema 为 (id，file_name，label）。
image_decode['file_name', 'img']() 通过每行的 file_name 读取图片文件，解码并将图片数据放入 img 列。
image_text_embedding.clip['img', 'vec'](model_name='clip_vit_b32',modality='image') 用 clip_vit_b32 将 img 列的每个图像的语义特征编码成向量，向量放到 vec 列。
tensor_normalize['vec','vec']() 将 vec 列的向量数据做归一化处理。
select['file_name', 'vec']() 选中 file_name 和 vec 两列作为最终结果。

创建向量库的索引

我们使用 Faiss 对图像的 Embedding 向量构建索引：

img_vectors.to_faiss['file_name', 'vec'](findex='./index.bin')

img_vectors 包含两列数据，分别是 file_name ， vec 。Faiss 对其中的 vec 列构建索引，并将每行的 file_name 与 vec 相关联。在向量搜索的过程中， file_name 信息会随结果返回。这一步可能会花上一些时间。

查询文本的向量化

查询文本的向量化过程与图像语义的向量化类似：

req = ( 
    towhee.dc['text'](['a samoyed lying down']) 
      .image_text_embedding.clip['text', 'vec'](model_name='clip_vit_b32', modality='text') 
      .tensor_normalize['vec', 'vec']() 
      .select['text','vec']() 
)

这里对代码做一些简要的说明：

dc['text'](['a samoyed lying down']) 创建了一个data collection，包含一行一列，列名为 text ，内容为'a samoyed lying down'。
image_text_embedding.clip['text', 'vec'](model_name='clip_vit_b32',modality='text') 用 clip_vit_b32 将文本 'query here' 编码成向量，向量放到 vec 列。注意，这里我们使用同样的模型( model_name='clip_vit_b32' )，但选择了文本模态（ modality='text' ）。这样可以保证图片和文本的语义向量存在于相同的向量空间。
tensor_normalize['vec','vec']() 将 vec 列的向量数据做归一化处理。
select['vec']() 选中 text ， vec 列作为最终结果。

查询

我们首先定义一个根据查询结果读取图片的函数 read_images ，用于支持召回后对原始图片的访问。

import cv2 
from towhee.types import Image 
def read_images(anns_results): 
    imgs = [] 
    for i in anns_results: 
        path = i.key 
        imgs.append(Image(cv2.imread(path), 'BGR')) 
    return imgs

接下来是查询的流水线：

results = ( 
    req.faiss_search['vec', 'results'](findex='./index.bin') 
        .runas_op['results', 'result_imgs'](func=read_images) 
        .select['text', 'result_imgs']() 
results.show()

faiss_search['vec', 'results'](findex='./index.bin', k = 5) 使用文本对应的Embedding 向量对图片的向量索引 index.bin 进行查询，找到与文本语义最接近的 5 张图片，并返回这 5 张图片所对应的文件名 results 。
runas_op['results', 'result_imgs'](func=read_images) 其中的read_images是我们定义的图片读取函数，我们使用 runas_op 将这个函数构造为 Towhee 推理流水线上的一个算子节点。这个算子根据输入的文件名读取图片。
select['text', 'result_imgs']() 选取 text 和 result_imgs 两列作为结果。

到这一步，我们以文搜图的完整流程就走完了，接下来，我们使用 Grado，将上面的代码包装成一个 demo。

使用 Gradio 打造 demo

首先，我们使用 Towhee 将查询过程组织成一个函数：

search_function = ( 
    towhee.dummy_input() 
        .image_text_embedding.clip(model_name='clip_vit_b32', modality='text') 
        .tensor_normalize() 
        .faiss_search(findex='./index.bin') 
        .runas_op(func=lambda results: [x.key for x in results]) 
        .as_function() 
)

然后，创建基于 Gradio 创建 demo 程序：

import gradio