【PGCCC】在 Postgres 上构建图像搜索引擎

我最近看到的最有趣的电子商务功能之一是能够搜索与我手机上的图片相似的产品。例如，我可以拍一双鞋或其他产品的照片，然后搜索产品目录以查找类似商品。使用这样的功能可以是一个相当简单的项目，只要有合适的工具。如果我们可以将问题定义为向量搜索问题，那么我们就可以使用 Postgres 来解决它！

在这篇博文中，我们将使用Postgres构建一个基本的图像搜索引擎。我们将使用预先训练的模型为图像和文本生成嵌入，然后将这些嵌入存储在 Postgres 中。pgvector扩展将使我们能够使用图像和原始文本作为查询对这些嵌入进行相似性搜索。

使用 CLIP 和 Postgres 进行图像搜索

2021 年，OpenAI 发表了一篇论文和CLIP（对比语言-图像预训练）的模型权重，该模型经过训练可以预测给定图像的最相关文本片段。通过一些巧妙的实现，此模型还可以用作搜索引擎的主干，该搜索引擎接受图像和文本作为输入查询。我们可以将图像转换为向量（嵌入），将图像的嵌入存储在 Postgres 中，使用扩展对这些向量进行相似性搜索，并使用它在 Postgres 之上构建图像搜索引擎。Hugging Face 上有许多CLIP 模型的开源变体，但我们将使用 OpenAI 的clip-vit-base-patch32 模式进行演示。

在之前的博客中，我们写过关于为语义文本搜索生成嵌入的内容。其中一些原则也适用于此。我们将为数据存储库生成嵌入，在本例中是图像目录。然后我们将这些嵌入存储在 Postgres 中。当我们查询数据时，我们需要使用相同的模型来为查询生成嵌入。不同之处在于，在这种情况下，我们的模型将为文本和图像生成嵌入。

在本示例中，我们将使用 OpenAI 在 Hugging Face 上提供的开源CLIP模型之一。请注意， CLIP 在生产中的使用存在限制。使用这些模型非常方便，因为它们的接口可在transformers Python 库中找到。

使用图像嵌入加载 Postgres

首先，我们需要获取原始图像。我们使用来自 Kaggle 的 Amazon Products 数据集。该数据集包含每个示例产品的图像 URL，因此我们将下载图像并将其存储在目录中。

在本例中，我们将把图像文件存储在本地，但在生产系统中，您可以将它们存储在 S3 等云存储服务中。

import pandas as pd

df = pd.read_csv("data/amazon_product.csv")

for i, row in df.iterrows():
    url = row["product_photo"]
    asin = row["asin"]
    response = requests.get(url)
    img = Image.open(BytesIO(response.content))
    if img.mode == 'RGBA':
        img = img.convert('RGB')

    img.save(f"./data/{asin}.jpg")

接下来，我们需要为获取的图像生成嵌入。我们将在 Postgres 中设置一个表来存储嵌入。

CREATE TABLE IF NOT EXISTS image_embeddings (
    image_path TEXT PRIMARY KEY,
    embeddings VECTOR(512)
);

我们将使用 CLIP 模型为每个图像生成嵌入，并将它们保存到 Postgres 表中。并创建一些辅助函数来加载图像、生成嵌入并将它们插入到 Postgres 中。

from pydantic import BaseModel
from transformers import (
    CLIPImageProcessor,
    CLIPModel,
)

MODEL = "openai/clip-vit-base-patch32"

image_processor = CLIPImageProcessor.from_pretrained(MODEL)
image_model = CLIPModel.from_pretrained(MODEL)

class ImageEmbedding(BaseModel):
    image_path: str
    embeddings: list[float]

def get_image_embeddings(
    image_paths: list[str], normalize=True
) -> list[ImageEmbedding]:
    # Process image and generate embeddings
    images = []
    for path in image_paths:
        images.append(Image.open(path))
    inputs = image_processor(images=images, return_tensors="pt")
    with torch.no_grad():
        outputs = image_model.get_image_features(**inputs)

    image_embeddings: list[ImageEmbedding] = []
    for image_p, embedding in zip(image_paths, outputs):
        if normalize:
            embeds = F.normalize(embedding, p=2, dim=-1)
        else:
            embeds = embedding
        image_embeddings.append(
            ImageEmbedding(
                image_path=image_p,
                embeddings=embeds.tolist(),
            )
        )
    return image_embeddings


def list_jpg_files(directory: str) -> list[str]:
    # List to hold the full paths of files
    full_paths = []
    # Loop through the directory
    for filename in os.listdir(directory):
        # Check if the file ends with .jpg
        if filename.endswith(".jpg"):
            # Construct full path and add it to the list
            full_paths.append(os.path.join(directory, filename))
    return full_paths


def pg_insert_embeddings(images: list[ImageEmbedding]):
    init_pg_vector = "CREATE EXTENSION IF NOT EXISTS vector;"
    init_table = """
        CREATE TABLE IF NOT EXISTS image_embeddings (image_path TEXT PRIMARY KEY, embeddings VECTOR(512));
    """
    insert_query = """
        INSERT INTO image_embeddings (image_path, embeddings)
        VALUES (%s, %s)
        ON CONFLICT (image_path)
        DO UPDATE SET embeddings = EXCLUDED.embeddings
        ;
    """
    with psycopg.connect(DATABASE_URL) as conn:
        with conn.cursor() as cur:
            cur.execute(init_pg_vector)
            cur.execute(init_table)

            for image in images:
                cur.execute(insert_query, (image.image_path, image.embeddings))

我们的辅助函数是这样的，让我们按顺序执行它们。

# get the paths to all our jpg images
images = list_jpg_files("./images")
# generate embeddings
image_embeddings = get_image_embeddings(images)
# insert them into Postgres
pg_insert_embeddings(image_embeddings)

快速验证嵌入是否已插入 Postgres。我们应该看到

psql postgres://postgres:postgres@localhost:5433/postgres
\x
select image_path, embeddings from image_embeddings limit 1;

image_path   | ./data/B086QB7WZ1.jpg
embeddings | [0.01544646,0.062326625,-0.03682831,0 ...

使用pgvector搜索相似图片

现在我们有了为文本生成嵌入的函数，我们可以在向量相似度搜索查询中使用这些嵌入。pgvector 支持多种距离运算符，但在本例中我们将使用余弦相似度。我们要搜索的嵌入存储在Postgres中，因此我们可以使用 SQL 进行余弦相似度搜索（1 - 余弦相似度）并找到嵌入与文本查询的嵌入最相似的图像。

def similarity_search(txt_embedding: list[float]) -> list[tuple[str, float]]:
    with psycopg.connect(DATABASE_URL) as conn:
        with conn.cursor() as cur:
            cur.execute(
                """
                        SELECT
                            image_path,
                            1 - (embeddings <=> %s::vector) AS similarity_score
                        FROM image_embeddings
                        ORDER BY similarity_score DESC
                        LIMIT 2;
                        """,
                (txt_embedding,),
            )
            rows = cur.fetchall()

            return [(row[0], row[1]) for row in rows]

与使用原始文本对数据进行向量搜索类似，我们将使用嵌入来搜索相似的图像。让我们抓取一张Cher的图像，我们可以从她的维基百科页面使用该图像。将其保存到./cher_wikipedia.jpg。

现在我们可以将单个图像传递到我们的get_image_embeddings()函数中，然后使用“similarity_search()”搜索相似的图像。

search_embeddings = get_image_embeddings(["./cher_wikipedia.jpg"])[0].embeddings
results = similarity_search(search_embeddings)

for image_path, score in results[:2]:
    print((image_path, score))

('B0DBQY1PKS.jpg', 0.5851975926639095)
('B0DBR4KDRF.jpg', 0.5125825695644287)

产品B0DBQY1PKS和B0DBR4KDRF（雪儿的“Forever”专辑）是与我们的雪儿形象最相似的两种产品。

使用原始文本查询图像

在搜索产品时，搜索相似的图片非常有用。但是，有时人们会希望根据给定的文本字符串来搜索图片。例如，Google 早就具备了搜索猫图片的功能。

from transformers import (
    CLIPTokenizerFast,
    CLIPTextModel,
    CLIPImageProcessor
)

MODEL = "openai/clip-vit-base-patch32"

processor = CLIPProcessor.from_pretrained(MODEL)
clip_model = CLIPModel.from_pretrained(MODEL)

def get_text_embeddings(text: str) -> list[float]:
    inputs = processor(text=[text], return_tensors="pt", padding=True)
    text_features = clip_model.get_text_features(**inputs)
    text_embedding = text_features[0].detach().numpy()
    embeds = text_embedding / np.linalg.norm(text_embedding)
    return embeds.tolist()

最后，我们可以使用这些函数生成嵌入，然后使用原始文本查询搜索我们的图像。我们将在产品目录中搜索“电话”的图像。

text_embeddings = get_text_embeddings("telephones")

results: list[tuple[str, float]] = similarity_search(search_embeddings)

for image_path, score in results[:2]:
    print((image_path, score))

('./data/B086QB7WZ1.jpg', 0.26320752344041964)
('./data/B00FRSYS12.jpg', 0.2626421138474824)

产品B086QB7WZ1和B00FRSYS12是与文本查询“电话”最相似的两幅图像。

Postgres 上的多模式搜索

我们已经从概念上展示了如何在 Postgres 上构建多模式搜索引擎。提醒一下，本博客中的代码可在Tembo Github 存储库中找到。我们使用 CLIP 模型为图像和文本生成嵌入，然后将这些嵌入存储在 Postgres 中。我们使用扩展pgvector对这些嵌入进行相似性搜索。这是一个强大的工具，可用于构建可以接受文本和图像查询的搜索引擎。关注 Tembo 博客，了解有关 Postgres 上矢量搜索用例的更多信息。

其他阅读材料

如果您对此主题感兴趣，请查看geoMusings博客上有关使用 pgvector 进行图像相似性分析的内容。另请阅读《视觉表征对比学习的简单框架》，ICML2020，Ting ChenSimon Kornblith、Mohammad Norouzi、Geoffrey E. Hinton。
#PG证书#PG考试#postgresql培训#postgresql考试#postgresql认证