从头构建一个轻量的个人图片检索库

AIGC：仅部分代码

TL;DR

如果你有图片语义检索的需求，目前大多数相册都已支持此功能，而且开源流行的 immich 也可以满足私有化场景图片分类和语义检索的完善功能了。但是如果你想理解向量嵌入等一些 AI 概念，又需要动手来加深理解，同时还想创造一些有价值效果酷炫的小工具的话，一个表情包搜索引擎或许是再合适不过的选择了。

其实很早之前就收藏了一篇与此相关的 HN 热文——我不小心构建了一个表情包搜索引擎，但是当时只是粗略浏览一遍，没动手实践过。最新有些想法心血来潮决定上手试试能做到什么程度。

Step 0 - 开始之前

虽然之前或多或少听过以下一些名词，这里还是枚举简单回顾下：

  • 向量（Vector）：可以简单理解为一组数字表示，比如 $[x, y]$ 就是一个二维向量，而机器学习中一般会用更高维度的向量来表示多模态数据。
  • 向量嵌入（Vector Embeddings）：把文本或图像转换成数值向量，从而可高效进行相似性检索。
  • 向量数据库（Vector Database）：专门存储并检索这些向量的数据库，帮你迅速找出相似项。
  • CLIP（Contrastive Language-Image Pre-training）：OpenAI 的模型，可把图像和文本同时编码成向量。

而一个图片搜索引擎主要的流程其实很简单：

  1. 遍历读取所有图片输入模型，输出向量后存储到库或文件。
  2. 根据用户输入检索向量库，返回对应图片。

Step 1 - 嵌入向量数据

考虑到我需要处理的图片可能包含中英文，同时不想依赖云服务的 API，vibe search 了一下本来决定使用 jina-clip-v2 这个看起来效果比较好的多语言多模态嵌入模型，结果因为 PyTorch 兼容性等问题报错给我劝退，决定先回退到 CLIP 这个轻量流行的模型用。

1
2
3
4
from pathlib import Path
import torch
from PIL import Image
from transformers import CLIPModel, CLIPProcessor

虽然这几个库都大名鼎鼎，但是出于项目本身的学习目的，这里还是记录它们的作用：

  • transormers：负责模型的加载和运行，自动下载模型，并将输入向量化。
  • torch（aka PyTorch）：负责处理计算和模型的存储读取。
  • PIL（Pillow）：负责图像处理和编辑。
  • pathlib：负责目录遍历。

首先添加一些基本参数配置：

1
2
3
4
IMAGE_DIR = r"C:\Users\Rosin\Pictures"
SAVE_PATH = "vectors.pt"
MODEL_NAME = "openai/clip-vit-base-patch32"
SUPPORTED_EXTS = {".png", ".jpg", ".jpeg", ".webp", ".bmp"}
  • 图片路径前的 r 用于处理 Windows 路径转义
  • 模型先使用 openai 的 clip，该模型在中文输入下能力欠缺，但是后续可以简单替换

transformers 已经提供了标准的 CLIP 接口，直接拿来用就行：

1
2
3
4
model = CLIPModel.from_pretrained(MODEL_NAME)
processor = CLIPProcessor.from_pretrained(MODEL_NAME)
# 开启评估模式（推理用）
model.eval()

接下来可以开始循环遍历图片来创建向量：

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
image_paths = [
    # 递归遍历图片目录
    p for p in Path(IMAGE_DIR).rglob("*")
    if p.suffix.lower() in SUPPORTED_EXTS
]
print(f"📂 共找到 {len(image_paths)} 张图片\n")

paths: list[str] = []
vectors: list[torch.Tensor] = []

for i, img_path in enumerate(image_paths, 1):
    try:
        # 转换为 RGB 模式
        img = Image.open(img_path).convert("RGB")
    except Exception as e:
        print(f"  ⚠️ 跳过: {img_path.name} ({e})")
        continue

    # "pt" 表示返回 `torch.Tensor` 格式对象
    inputs = processor(images=img, return_tensors="pt")

    # 使用推理模式（不计算梯度）
    with torch.no_grad():
        # 获取图片特征向量
        vec = model.get_image_features(pixel_values=inputs["pixel_values"])
        if not isinstance(vec, torch.Tensor):
            vec = vec.pooler_output
    # 向量归一化，保留语义特征同时提升检索效率
    vec = vec / vec.norm(dim=-1, keepdim=True)

    paths.append(str(img_path))
    vectors.append(vec)
    print(f"  [{i}/{len(image_paths)}] {img_path.name}")

遍历完成后保存到向量文件 vectors.pt：

1
2
3
# 打包所有向量
matrix = torch.cat(vectors, dim=0)
torch.save({"paths": paths, "vectors": matrix}, SAVE_PATH)

完成后可以打印查看向量格式：

1
2
3
4
5
6
7
8
9
10
print(f"   图片数量: {len(paths)}")
print(f"   向量维度: {matrix.shape[1]}")
print(f"   矩阵形状: {matrix.shape}")

print("\n🔍 数据预览 (前 3 条):")
for i in range(min(3, len(paths))):
    print(f"  {i+1}. 路径: {Path(paths[i]).name}")
    # 打印前 5 个数值作为示例
    sample = matrix[i][:5].tolist()
    print(f"     向量预览: {[round(x, 4) for x in sample]}...")
1
2
3
4
5
6
7
8
   图片数量: 42
   向量维度: 512
   矩阵形状: torch.Size([42, 512])

🔍 数据预览 (前 3 条):
     向量预览: [0.0392, -0.0132, 0.0034, -0.0282, 0.0004]...
     向量预览: [0.0088, -0.0059, -0.0319, 0.0002, -0.0078]...
     向量预览: [0.0275, 0.0034, -0.0185, -0.0581, 0.0342]...

可以注意到每张图片都被转换为了一条 512 维度的向量数据嵌入，而后续如果需要更新这个向量库（添加图片），也可以单独转换插入即可。

Step 2 - 实现图片检索

有了向量库后，语义检索需要以下几步操作：

  1. 加载模型和向量库
  2. 将搜索文字输入模型（CLIP）得到输出向量
  3. 计算向量输出和向量库中的数据相似度并排序
  4. 展示结果

首先是向量库的加载，模型的加载和建库时一样：

1
2
3
4
data = torch.load(VECTORS_PATH, weights_only=False)
model = CLIPModel.from_pretrained(MODEL_NAME)
processor = CLIPProcessor.from_pretrained(MODEL_NAME)
model.eval()

搜索词通过模型编码为向量：

1
2
3
4
5
6
7
8
9
10
11
inputs = processor(text=query, return_tensors="pt", padding=True)

with torch.no_grad():
    text_vec = model.get_text_features(
        input_ids=inputs["input_ids"],
        attention_mask=inputs["attention_mask"],
    )
    if not isinstance(text_vec, torch.Tensor):
        text_vec = text_vec.pooler_output

text_vec = text_vec / text_vec.norm(dim=-1, keepdim=True)

然后计算相似度并排序：

1
2
3
similarities = (text_vec @ vectors.T).squeeze(0)
top_k = min(TOP_K, len(paths))
scores, indices = similarities.topk(top_k)
  • TOP_K 是需要打印比较的排序结果数量

测试目录图片比较杂，计算完成的结果直接使用 os.startfile(base_path) 打开看效果：

输入猫或 cat：

输入 Cillian Murphy 或者 whore：

Bad Case 的话就是输入基利安·墨菲的话并不能得到上面的结果，但是在替换为阿里的 chinese-clip-vit-base-patch16 之后就能中文检索了。Chinese-CLIP 的接口是 CLIP 兼容的，所以只需要简单替换包引用和模型名称就行：

1
2
3
4
# from transformers import CLIPModel, CLIPProcessor
from transformers import ChineseCLIPModel as CLIPModel, ChineseCLIPProcessor as CLIPProcessor

MODEL_NAME = "OFA-Sys/chinese-clip-vit-base-patch16"

Step 3 - Time to cook

Demo 跑通之后，正常的节奏应该是让 AI 整一个花里胡哨的前端页面，后端 fast api 一把梭。但是作为个人（特指本人）使用的场景，部署在 Nas 上虽然可以轻松完成建库定时任务和前端的部署，但每次都走 Tailscale 访问前端去搜 meme 有点麻烦。