👬
Jina CLIP v2とDuckDBを使って類似画像検索してみた
はじめに
日本語対応の「Jina CLIP v2」と組み込みデータベースの「DuckDB」を使って、画像の類似検索をしてみました。
Jina CLIP v2の使い方、DuckDBによる類似検索については以下の記事を書きましたので、合わせてご参照ください。
仕様
今回の類似画像検索は、以下の仕様で実装しました。
- 画像は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リポジトリにも格納しています。
画像をデータベースに登録する
まずは、画像から特徴量を抽出し、データベースに登録します。使用した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