👬
Jina CLIP v2とDuckDBを使って類似画像検索してみた

2024/12/31に公開
 はじめに日本語対応の「Jina CLIP v2」と組み込みデータベースの「DuckDB」を使って、画像の類似検索をしてみました。
Jina CLIP v2の使い方、DuckDBによる類似検索については以下の記事を書きましたので、合わせてご参照ください。
https://zenn.dev/yuyakato/articles/175d6d590da13a
https://zenn.dev/yuyakato/articles/af67b301d37e17
https://zenn.dev/yuyakato/articles/e83650be2140b6

 仕様今回の類似画像検索は、以下の仕様で実装しました。
画像はJina CLIP v2を使って特徴量化する
1024次元の特徴量として扱う
Jina CLIP v2はマトリョーシカ表現（Matryoshka Representations）を使っているので、512次元、256次元などのより低い次元数でも類似画像検索できそうですが、まだ試していません

データベースにはDuckDBを使用する
指定したディレクトリ以下のJPEG画像をデータベースに登録する
今回は単純に*.jpgファイルを検索しています

指定したパスのJPEG画像に類似した画像上位10件のファイルパス、コサイン類似度を出力する
コサイン類似度の計算はDuckDBにて行う

 環境今回は以下の環境で試しました。
GPU: NVIDIA GeForce RTX 4080 16GB
OS: Ubuntu 22.04.5 LTS
NVIDIAドライバ: 550.127.08
Docker: 24.0.7
Docker Compose: 2.18.0
NVIDIA Container Toolkit: 1.13.5
Python: 3.11.6
requirements.txtの内容は以下の通りです。
requirements.txtの内容requirements.txt
anykeystore==0.2
apex==0.9.10.dev0
certifi==2024.12.14
charset-normalizer==3.4.0
click==8.1.8
cryptacular==1.6.2
defusedxml==0.7.1
duckdb==1.1.3
einops==0.8.0
filelock==3.16.1
flash-attn==2.7.2.post1
fsspec==2024.12.0
greenlet==3.1.1
huggingface-hub==0.27.0
hupper==1.12.1
idna==3.10
Jinja2==3.1.5
MarkupSafe==3.0.2
mpmath==1.3.0
networkx==3.4.2
numpy==2.2.1
nvidia-cublas-cu12==12.4.5.8
nvidia-cuda-cupti-cu12==12.4.127
nvidia-cuda-nvrtc-cu12==12.4.127
nvidia-cuda-runtime-cu12==12.4.127
nvidia-cudnn-cu12==9.1.0.70
nvidia-cufft-cu12==11.2.1.3
nvidia-curand-cu12==10.3.5.147
nvidia-cusolver-cu12==11.6.1.9
nvidia-cusparse-cu12==12.3.1.170
nvidia-nccl-cu12==2.21.5
nvidia-nvjitlink-cu12==12.4.127
nvidia-nvtx-cu12==12.4.127
oauthlib==3.2.2
packaging==24.2
PasteDeploy==3.1.0
pbkdf2==1.3
pillow==11.0.0
plaster==1.1.2
plaster-pastedeploy==1.0.1
pyramid==2.0.2
pyramid-mailer==0.15.1
python3-openid==3.2.0
PyYAML==6.0.2
regex==2024.11.6
repoze.sendmail==4.4.1
requests==2.32.3
requests-oauthlib==2.0.0
ruff==0.8.4
safetensors==0.4.5
SQLAlchemy==2.0.36
sympy==1.13.1
timm==1.0.12
tokenizers==0.21.0
torch==2.5.1
torchvision==0.20.1
tqdm==4.67.1
transaction==5.0
transformers==4.47.1
translationstring==1.4
triton==3.1.0
typing_extensions==4.12.2
urllib3==2.3.0
velruse==1.1.1
venusian==3.1.1
WebOb==1.8.9
WTForms==3.2.1
wtforms-recaptcha==0.3.2
xformers==0.0.28.post3
zope.deprecation==5.0
zope.interface==7.2
zope.sqlalchemy==3.1

 ソースコードソースコード一式は、以下のGitHubリポジトリにも格納しています。
https://github.com/nayutaya/202412-duckdb-similarity-search/tree/main/image-similarity-search

 画像をデータベースに登録するまずは、画像から特徴量を抽出し、データベースに登録します。使用したPythonスクリプトは以下の通りです。
add_image.py
from pathlib import Path

import click
import duckdb
import numpy as np
import torch
from PIL import Image
from transformers import AutoModel, AutoProcessor


def load_jina_clip_v2(device: torch.device) -> tuple:
    """Jina CLIP v2を読み込む。"""
    model_name = "jinaai/jina-clip-v2"
    model = AutoModel.from_pretrained(model_name, trust_remote_code=True).to(device)
    processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=True)
    return model, processor


def extract_image_feature(device: torch.device, model, processor, image) -> np.ndarray:
    """画像からCLIP特徴量を取得する。"""
    inputs = processor(images=image, return_tensors="pt")
    for k, v in inputs.items():
        inputs[k] = v.to(device)
    with torch.no_grad():
        features = model.get_image_features(**inputs)
    features /= features.norm(dim=-1, keepdim=True)
    features = features.float().cpu().numpy()
    return features[0]


@click.command()
@click.option(
    "--db-file",
    required=True,
    type=click.Path(path_type=Path, file_okay=True, dir_okay=False, writable=True),
    help="path to DuckDB file.",
)
@click.option(
    "--image-dir",
    required=True,
    type=click.Path(path_type=Path, file_okay=False, dir_okay=True, exists=True),
    help="path to image directory.",
)
@click.option(
    "--device",
    required=True,
    type=torch.device,
    default=(
        torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
    ),
    help="device name to use for inference.",
)
def main(db_file: Path, image_dir: Path, device: torch.device) -> None:
    print("Loading model...")
    model, processor = load_jina_clip_v2(device)

    print("Connecting to DB...")
    con = duckdb.connect(db_file)
    con.execute("""
    CREATE TABLE IF NOT EXISTS images(
        file_path VARCHAR NOT NULL PRIMARY KEY,
        feature FLOAT4[1024] NOT NULL
    )
    """)

    print("Finding images...")
    image_files = list(image_dir.rglob("*.jpg"))
    print(f'Found {len(image_files)} images under "{image_dir}"')

    added_paths = set(
        row[0] for row in con.execute("SELECT file_path FROM images").fetchall()
    )

    for index, image_file in enumerate(image_files, 1):
        if str(image_file) in added_paths:
            print(f"[{index}/{len(image_files)}] Skip (already added): {image_file}")
            continue

        try:
            image = Image.open(image_file).convert("RGB")
            feature = extract_image_feature(device, model, processor, image)
            con.execute(
                "INSERT INTO images(file_path, feature) VALUES ($file_path, $feature)",
                {"file_path": str(image_file), "feature": feature},
            )
            print(f"[{index}/{len(image_files)}] Added: {image_file}")
        except Exception as e:
            print(f"[{index}/{len(image_files)}] Error: {image_file}: {e}")

    con.commit()
    con.close()
    print("Done")


if __name__ == "__main__":
    main()
実行例は以下の通りです。（警告メッセージを一部省略しています）
$ python add_image.py --db-file example.duckdb --image-dir example
Loading model...
Connecting to DB...
Finding images...
Found 9 images under "example"
[1/9] Added: example/pakutaso_42965.jpg
[2/9] Added: example/pakutaso_30012.jpg
[3/9] Added: example/pakutaso_36665.jpg
[4/9] Added: example/pakutaso_2769.jpg
[5/9] Added: example/pakutaso_32080.jpg
[6/9] Added: example/pakutaso_86239.jpg
[7/9] Added: example/pakutaso_87869.jpg
[8/9] Added: example/pakutaso_5507.jpg
[9/9] Added: example/pakutaso_77003.jpg
Done
なお、動作確認用の画像はぱくたそからお借りしました。

 類似画像をデータベースから検索する続いて、登録済みの画像に類似する画像をデータベースから検索します。使用したPythonスクリプトは以下の通りです。
search_image.py
from pathlib import Path

import click
import duckdb


@click.command()
@click.option(
    "--image-file",
    required=True,
    type=click.Path(path_type=Path, file_okay=True, dir_okay=False, exists=True),
    help="path to query image",
)
@click.option(
    "--db-file",
    required=True,
    type=click.Path(
        path_type=Path, file_okay=True, dir_okay=False, writable=True, exists=True
    ),
    help="path to DuckDB file.",
)
def main(db_file, image_file):
    con = duckdb.connect(db_file)

    count = con.execute(
        "SELECT COUNT(*) FROM images WHERE file_path = ?", (str(image_file),)
    ).fetchone()[0]
    if count == 0:
        print(f'Error: "{image_file}" was not found in DB')
        return

    images = con.execute(
        """
        SELECT
            a.file_path,
            array_cosine_similarity(a.feature, b.feature) AS similarity
        FROM images AS a, images AS b
        WHERE b.file_path = $query_file_path AND a.file_path <> $query_file_path
        ORDER BY similarity DESC
        LIMIT 10
        """,
        {"query_file_path": str(image_file)},
    ).fetchall()

    for file_path, similarity in images:
        print(f"{file_path}: {similarity:.04f}")

    con.close()


if __name__ == "__main__":
    main()
実行例は以下の通りです。
$ python search_image.py --db-file example.duckdb --image-file example/pakutaso_30012.jpg
example/pakutaso_32080.jpg: 0.8817
example/pakutaso_36665.jpg: 0.7127
example/pakutaso_5507.jpg: 0.6243
example/pakutaso_87869.jpg: 0.6210
example/pakutaso_2769.jpg: 0.5790
example/pakutaso_77003.jpg: 0.5611
example/pakutaso_42965.jpg: 0.5031
example/pakutaso_86239.jpg: 0.4977
ちなみにクエリ画像は以下の通りです。
コサイン類似度が最も高い画像は以下の通りです。
コサイン類似度が最も低い画像は以下の通りです。
ちゃんとそれっぽく動作していますね。

 おわりにJina CLIP v2とDuckDBを使うことで、とても簡単に画像の類似検索を行うことができました。

Jina CLIP v2は日本語にも対応しているCLIPモデルのため、クエリ文字列からの画像検索も試してみたい所です。
本記事が何らかの参考になれば幸いです。
Discussion

ログインするとコメントできます