AI 与向量

使用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 实现图像搜索以及反向图像搜索功能。