使用OpenAI CLIP实现图像搜索
结合OpenAI CLIP模型与Supabase Vector实现图像搜索功能
OpenAI CLIP模型经过多种(图像,文本)对的训练。您可以使用CLIP模型实现:
- 文本到图像/图像到文本/图像到图像/文本到文本搜索
- 可以使用常规的
SentenceTransformers
训练代码在您自己的图像和文本数据上进行微调
SentenceTransformers
提供的模型允许您将图像和文本嵌入到相同的向量空间中。这可用于查找相似图像以及实现图像搜索。
您可以在GitHub上找到完整的Python Poetry项目应用代码。
使用Poetry创建新Python项目
Poetry为Python提供打包和依赖管理。如果尚未安装,请通过pip安装poetry:
1pip install poetry
然后初始化新项目:
1poetry new image-search
设置Supabase项目
如果尚未安装,请先安装Supabase CLI,然后在新创建的poetry项目根目录中初始化Supabase:
1supabase init
接下来启动本地Supabase堆栈:
1supabase start
这将启动本地Supabase堆栈并打印大量环境详细信息,包括本地DB URL
。请记下该URL供后续使用。
安装依赖项
我们需要为项目添加以下依赖项:
vecs
:Supabase向量Python客户端sentence-transformers
:用于句子、文本和图像嵌入的框架(与OpenAI CLIP模型配合使用)matplotlib
:用于显示图像结果
1poetry add vecs sentence-transformers matplotlib
导入必要的依赖项
在您的Python主脚本顶部,导入所需的依赖项并将上述DB URL
存储在一个变量中:
1234567from 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向量集合,为您的图像生成嵌入向量,并将这些嵌入向量更新插入到数据库中:
12345678910111213141516171819202122232425262728293031323334353637383940414243def 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
文件中:
123[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,我们可以查询嵌入向量。既可以使用图像作为搜索输入,也可以从字符串输入生成嵌入向量并将其用作查询输入:
1234567891011121314151617181920212223def 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 实现图像搜索以及反向图像搜索功能。