基于 Embedding 的本地图像搜索

前言

通过上一篇 LLM多模态嵌入 - 图片嵌入，我们学习和了解了嵌入技术在图片领域的使用，并通过 openai/clip 的简单示例感受到了当文本嵌入和图片嵌入结合之后的强大之处。可以看到通过嵌入将两个图片的内容是否类似 ，一张猫的图片和一段描述一只可爱的小猫 这类抽象的问题转换为数学问题。实现嵌入之后，比较向量的相似度即可。

既然通过嵌入可以理解图片的内容，我们自然而然会想到用这种技术去进行图片搜索。下面我们尝试借助 clip ，探索图像搜索功能，支持通过文字从批量图片中找到符合描述的图片。

图像搜索实现

有了上一篇 LLM多模态嵌入 - 图片嵌入的基础，对于图像搜索我们很容易想到如下方案

对批量的图片进行 embedding ，获得每一张图片的 embedding 结果，将图片的地址（或者任意可以标识图片唯一性）和 embedding 保存下来。
搜索时将输入内容文本或者图片进行 embedding 后和已经存储的结果进行比较，返回相似度 topN 的图片地址。

下面我们就按照以上思路进行实操，看看实际效果如何。

批量图片的嵌入和存储

这里简单起见，我们将嵌入结果通过 Python 内置的 Pickle 序列化格式进行转换后存储在 .pkl 文件中。

在 Python 中，.pkl 文件（或 .pickle）是 Pickle 序列化格式的文件扩展名，用于将 Python 对象（如变量、函数、类实例等）序列化为二进制数据并保存到磁盘，或从磁盘反序列化还原对象。Pickle 是 Python 标准库的一部分（pickle 模块），专为 Python 对象设计。

python 复制代码

def compute_image_embeddings(image_paths):
    embeddings = []
    final_image_paths = []
    for path in tqdm(image_paths, desc="提取图像特征"):
        try:
            img = Image.open(path).convert("RGB")
            inputs = processor(images=img, return_tensors="pt", padding=True, use_fast=True).to(device)
            image_features = model.get_image_features(inputs.pixel_values).squeeze(0)
            embeddings.append(image_features)
            final_image_paths.append(path)
        except UnidentifiedImageError:
            print(path, "unsupported")
            continue

    return torch.stack(embeddings), final_image_paths

需要注意的是，由于在 clip 中默认是按批次处理图片的，因此返回 embedding 结果默认形状是 [1,512],这里为了方便处理，我们通过 squeeze 去掉第一个维度，只保留了特征数据。

保存嵌入结果

python 复制代码

def compute_embeddings():
    paths = get_all_images(IMAGE_DIR)
    features, paths = compute_image_embeddings(paths)
    os.makedirs(os.path.dirname(PKL_CACHE_PATH), exist_ok=True)
    with open(PKL_CACHE_PATH, "wb") as f:
        pickle.dump((features, paths), f)
    return features, paths

这里我们在 pickle 文件中保存的是每张图片的嵌入结果和其路径的字典，这样后续使用的时候我们就可以直接复用这个嵌入后的结果，不用每次都耗费时间去做嵌入，只有图片列表发生变化的时候再去做嵌入。

文本搜索

有了图片的嵌入结果，我们就可以基于文本就行搜索了。

python 复制代码

def search(text, image_features, image_paths, top_k=2):
    text_inputs = processor(text=text, return_tensors="pt", padding=True).to(device)
    text_emb = model.get_text_features(**text_inputs)

    # ✅ 1. image_features 是 [N, 512]，text_emb 是 [1, 512]
    print(image_features.shape, text_emb.shape)
    sims = torch.nn.functional.cosine_similarity(image_features, text_emb, dim=1)  # shape: [N]
    
    # ✅ 2. 取 Top-K 索引（确保是 1D）
    top_idx = torch.topk(sims, top_k).indices  # shape: [top_k]

    # ✅ 5. 遍历索引，取图像路径和相似度
    return [(image_paths[i], sims[i].item()) for i in top_idx]

对文本进行嵌入后，计算当前文本的向量表达和所有图片向量的相似度，这里我们取 top 2 的值作为结果。

为了快速验证效果，我们以上一篇示例中用到的图片为例，进行效果验证，总共有 7 张图片

shell 复制代码

features: torch.Size([7, 512])
['imgs\\animal.jpeg', 'imgs\\cat.jpg', 'imgs\\centaur_warrior.jpeg', 'imgs\\child.jpeg', 'imgs\\dog.jpeg', 'imgs\\horse.jpeg', 'imgs\\woman_with_horse.jpeg']
text = ：cat
torch.Size([7, 512]) torch.Size([1, 512])
top_idx: torch.Size([2])
cat.jpg: 0.2941
animal.jpeg: 0.2254

可以看到输入的文本 cat 嵌入之后和图片嵌入集合进过相似度计算后返回了值最高的 2 个结果，我们将这两个结果画出来。

可以看到返回结果就是张 7 张图片中唯一一张猫的图片和另外一张和猫长得相似的图片，结果似乎还不错。我们再换一个词 woman 试试

shell 复制代码

text = ：woman 
centaur_warrior.jpeg: 0.2489
woman_with_horse.jpeg: 0.2352

这次也是返回来图片集合中唯二包含woman 元素的图片。

小结

从以上结果来看，这个思路是可行的，用 clip 分别嵌入图像和文本数据后，再次计算相似度也能获得较好的结果。但是，有一点我们需要注意的是，单独计算文本和图片的嵌入之后，再次计算二者相似度的时候，相似度结果就没有之前统一计算那么高了。