使用OpenAI CLIP实现图像搜索

结合OpenAI CLIP模型与Supabase Vector实现图像搜索功能

OpenAI CLIP模型经过多种(图像,文本)对的训练。您可以使用CLIP模型实现：

文本到图像/图像到文本/图像到图像/文本到文本搜索
可以使用常规的SentenceTransformers训练代码在您自己的图像和文本数据上进行微调

SentenceTransformers提供的模型允许您将图像和文本嵌入到相同的向量空间中。这可用于查找相似图像以及实现图像搜索。

您可以在GitHub上找到完整的Python Poetry项目应用代码。

使用Poetry创建新Python项目

Poetry为Python提供打包和依赖管理。如果尚未安装，请通过pip安装poetry：

1
pip install poetry

然后初始化新项目：

1
poetry new image-search

设置Supabase项目

如果尚未安装，请先安装Supabase CLI，然后在新创建的poetry项目根目录中初始化Supabase：

1
supabase init

接下来启动本地Supabase堆栈：

1
supabase start

这将启动本地Supabase堆栈并打印大量环境详细信息，包括本地DB URL。请记下该URL供后续使用。

安装依赖项

我们需要为项目添加以下依赖项：

vecs：Supabase向量Python客户端
sentence-transformers：用于句子、文本和图像嵌入的框架（与OpenAI CLIP模型配合使用）
matplotlib：用于显示图像结果

1
poetry add vecs sentence-transformers matplotlib

导入必要的依赖项

在您的Python主脚本顶部，导入所需的依赖项并将上述DB URL存储在一个变量中：

1
2
3
4
5
6
7
from PIL import Imagefrom sentence_transformers import SentenceTransformerimport vecsfrom matplotlib import pyplot as pltfrom matplotlib import image as mpimgDB_CONNECTION = "postgresql://postgres:postgres@localhost:54322/postgres"

为您的图像创建嵌入向量

在项目根目录下创建一个名为images的新文件夹，并添加一些图片。您可以使用GitHub上示例项目中的图片，也可以在Unsplash上寻找免版权的图片。

接下来，创建一个seed方法，该方法将创建一个新的Supabase向量集合，为您的图像生成嵌入向量，并将这些嵌入向量更新插入到数据库中：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
def seed():    # 创建向量存储客户端    vx = vecs.create_client(DB_CONNECTION)    # 创建一个维度为512的向量集合    images = vx.get_or_create_collection(name="image_vectors", dimension=512)    # 加载CLIP模型    model = SentenceTransformer('clip-ViT-B-32')    # 编码图像：    img_emb1 = model.encode(Image.open('./images/one.jpg'))    img_emb2 = model.encode(Image.open('./images/two.jpg'))    img_emb3 = model.encode(Image.open('./images/three.jpg'))    img_emb4 = model.encode(Image.open('./images/four.jpg'))    # 向*images*集合添加记录    images.upsert(        records=[            (                "one.jpg",        # 向量标识符                img_emb1,         # 向量值，可以是列表或np.array                {"type": "jpg"}   # 关联的元数据            ), (                "two.jpg",                img_emb2,                {"type": "jpg"}            ), (                "three.jpg",                img_emb3,                {"type": "jpg"}            ), (                "four.jpg",                img_emb4,                {"type": "jpg"}            )        ]    )    print("已插入图像")    # 为集合创建索引以提高搜索性能    images.create_index()    print("已创建索引")

将此方法作为脚本添加到您的pyproject.toml文件中：

1
2
3
[tool.poetry.scripts]seed = "image_search.main:seed"search = "image_search.main:search"

通过poetry shell激活虚拟环境后，您现在可以通过poetry run seed运行种子脚本。您可以通过访问本地Supabase仪表板localhost:54323，选择vecs模式和image_vectors数据库来检查本地数据库中生成的嵌入向量。

通过文本查询执行图像搜索

借助 Supabase Vector，我们可以查询嵌入向量。既可以使用图像作为搜索输入，也可以从字符串输入生成嵌入向量并将其用作查询输入：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
def search():    # 创建向量存储客户端    vx = vecs.create_client(DB_CONNECTION)    images = vx.get_or_create_collection(name="image_vectors", dimension=512)    # 加载 CLIP 模型    model = SentenceTransformer('clip-ViT-B-32')    # 编码文本查询    query_string = "a bike in front of a red brick wall"    text_emb = model.encode(query_string)    # 查询集合并筛选元数据中 "type" = "jpg" 的记录    results = images.query(        data=text_emb,                      # 必需参数        limit=1,                            # 返回记录数量        filters={"type": {"$eq": "jpg"}},   # 元数据筛选器    )    result = results[0]    print(result)    plt.title(result)    image = mpimg.imread('./images/' + result)    plt.imshow(image)    plt.show()

通过将查询结果限制为一条，我们可以向用户展示最相关的图像。最后使用 matplotlib 向用户展示图像结果。

您可以立即运行 poetry run search 进行测试，系统将展示一张"红色砖墙前的自行车"图片。

总结

仅需几行 Python 代码，您就能使用 OpenAI 的 CLIP 模型和 Supabase Vector 实现图像搜索以及反向图像搜索功能。