Appleが今週公開したAIMv2を使って画像検索システム(類似画像検索)を構築してみた
はじめに
まず、知ってる方いると思うが、anacondaとnotebook形式が大嫌いです。なので、これはnotebook形式に対応していないやり方です。
こんにちは、しゅんです、いつも見てくれてありがとうございます。今回はた、Appleが公開した大規模ビジョンエンコーダ「AIMv2」を使用して、画像検索システム(類似画像検索)を構築する方法を紹介します。今回はaimv2-large-patch14-224のモデルを使います。では、解説を進めます：
Apple公式は特徴量を取り出すサンプルコードまでしか出していない、海外の皆さんも何もてを加えていない状況です。このような類似度検索のプロジェクトもまだ見当たらないです。このプロジェクトを思い浮かんでGPTと一緒にチャレンジして、うまく成功したので、ここでシェアします。よかったら、ぜひ拡散してください。よろしくお願いします
- AIMv2の概要と特徴
- システムの設計と分割されたコードの意義→ 全部notebook形式が大嫌いだからです。それ以外は下で解説書いています
- 特徴量抽出のコード (
feature_extraction.py) の詳細解説 - 類似画像検索のコード (
image_search.py) の詳細解説 - 実行結果と応用例
AIMv2とは
概要
AIMv2 (Autoregressive Image Model v2) は、Appleが開発した大規模なビジョンエンコーダであり、以下のようなタスクにおいて高いパフォーマンスを発揮します：
- マルチモーダル理解 (例: テキストと画像の関連性を理解する)
- オープンボキャブラリの物体検出
- 参照表現の理解
AIMv2は、単一のカメラやテキストデータと画像データを組み合わせたマルチモーダルアプローチを活用することで、より高度なビジョンタスクに対応可能です。
公式の画像です
主な特徴
-
シンプルかつスケーラブルなアーキテクチャ:
- AIMv2は、従来のモデル（例: CLIP、DINOv2）を多数のベンチマークで上回ります。
-
幅広い解像度に対応:
- AIMv2は、224px、336px、448px、ネイティブ解像度で事前学習されたモデルを提供します。
-
事前学習済みチェックポイント:
- AIMv2では、Hugging Face Hubを通じてモデルが提供されており、すぐに試せる環境が整っています。
-
Apple Silliconにも対応:
MLXとはMLX is an array framework for machine learning on Apple silicon, brought to you by Apple machine learning research.
MLXは、アップルの機械学習研究によってもたらされた、アップルのシリコン上での機械学習のための配列フレームワークです。
つまりこれでAppleのMAC シリコン M chip で動作も可能ですが、今回は
pytrochGPUをメインにしていますちなみに、海外の記事ではpytorchとMLXの能力スピードテストがありました（ベンチマークかなぁ）。MLXが少しだけ勝っていることがわかります。
公式リポジトリや論文はこちらを参照してください：
システム設計の概要
AIMv2を利用した画像検索システムでは、以下の2つのステップを分けて設計します：
-
特徴量の抽出と保存 (
prepare_features.py)- まず、
cocodatasetのval2017画像セットをダウンロードしプロジェクト内に置きますwget -c http://images.cocodataset.org/zips/val2017.zip unzip - データセット内のすべての画像をAPPLEのAIMV2に読み込ませる、このモデルが特徴量を計算し、
.ptファイルに保存します。 - 処理に時間がかかるため、1回の計算で使い回せるようにします。自分の場合は約45分かかりました
- まず、
-
類似画像検索の実行 (
image_search.py)- 保存済みの特徴量を読み込み、クエリ画像とデータセット画像の類似性を計算して最も近い画像を検索します。
特徴量抽出のコード (prepare_features.py)
以下は、画像の特徴量を抽出し、保存するコードです。
import os
from tqdm import tqdm
from PIL import Image
import torch
from transformers import AutoImageProcessor, AutoModel
# GPUの設定
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# モデルとプロセッサのロード
processor = AutoImageProcessor.from_pretrained("apple/aimv2-large-patch14-224")
model = AutoModel.from_pretrained("apple/aimv2-large-patch14-224", trust_remote_code=True).to(device)
# 特徴量抽出関数
def extract_features(image_dir, save_path):
features_dict = {}
image_files = [f for f in os.listdir(image_dir) if f.endswith(".jpg")]
print(f"Extracting features from {len(image_files)} images in {image_dir}...")
for img_file in tqdm(image_files, desc="Extracting features", unit="image"):
img_path = os.path.join(image_dir, img_file)
try:
image = Image.open(img_path).convert("RGB")
inputs = processor(images=image, return_tensors="pt").to(device)
outputs = model(**inputs)
# print(outputs)
features = outputs.last_hidden_state.mean(dim=1).detach().cpu()
features_dict[img_file] = features
except Exception as e:
print(f"Error processing {img_file}: {e.__class__.__name__} - {e}")
torch.save(features_dict, save_path)
print(f"Features saved to {save_path}")
# メイン処理
if __name__ == "__main__":
image_dir = "/media/syun/ssd02/python_learning/apple/qiita_project_AIMv2/coco_image/val2017"
save_path = "/media/syun/ssd02/python_learning/apple/qiita_project_AIMv2/model/coco_features.pt"
extract_features(image_dir, save_path)
コード解説
-
デバイス設定:
- GPUまたはCPUを自動選択します。
-
モデルのロード:
- Hugging FaceのAIMv2モデルを使用し、画像を処理可能な状態にします。
-
特徴量の抽出:
- データセット内の画像を順に処理し、特徴量を取得。
- 平均プーリングを用いて特徴量を1次元化します。
-
特徴量の保存:
- 計算済みの特徴量を辞書形式で保存します（
torch.saveを使用）。
- 計算済みの特徴量を辞書形式で保存します（
-
保存など画像パスは自分が置いていた場所に置き換えてください
image_dir = "/media/syun/ssd02/python_learning/apple/qiita_project_AIMv2/coco_image/val2017" save_path = "/media/syun/ssd02/python_learning/apple/qiita_project_AIMv2/model/coco_features.pt"
類似画像検索のコード (image_search.py)
クエリ画像はgoogleで適当に調べて、test_search_imageのフォルダーに保存して入力しました。
以下は、クエリ画像とデータセット内の画像を比較し、類似画像を検索するコードです。
import os
import torch
from PIL import Image
import cv2
from sklearn.metrics.pairwise import cosine_similarity
from transformers import AutoImageProcessor, AutoModel
# GPUの設定
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# モデルとプロセッサのロード
processor = AutoImageProcessor.from_pretrained("apple/aimv2-large-patch14-224")
model = AutoModel.from_pretrained("apple/aimv2-large-patch14-224", trust_remote_code=True).to(device)
# クエリ画像の特徴量を取得
def get_query_features(query_image_path):
query_image = Image.open(query_image_path).convert("RGB")
query_inputs = processor(images=query_image, return_tensors="pt").to(device)
query_outputs = model(**query_inputs)
# print(query_outputs)
query_features = query_outputs.last_hidden_state.mean(dim=1).detach().cpu()
return query_features
# 類似画像を検索
def find_most_similar_image(query_features, features_dict):
best_match_file = None
best_similarity = -1
query_features = query_features.numpy()
for img_file, img_features in features_dict.items():
img_features = img_features.numpy()
similarity = cosine_similarity(query_features, img_features)[0, 0]
if similarity > best_similarity:
best_similarity = similarity
best_match_file = img_file
return best_match_file, best_similarity
# メイン処理
if __name__ == "__main__":
import argparse
# コマンドライン引数
parser = argparse.ArgumentParser(description="Search for the most similar image")
parser.add_argument("query_image", type=str, help="Path to the query image")
parser.add_argument("--features", type=str, default="coco_features.pt", help="Path to the saved features file")
args = parser.parse_args()
# クエリ画像と特徴量のロード
query_image_path = args.query_image
features_path = args.features
if not os.path.exists(features_path):
print(f"Features file {features_path} not found. Run prepare_features.py first.")
exit(1)
print(f"Loading features from {features_path}...")
features_dict = torch.load(features_path)
print(f"Extracting features from query image: {query_image_path}...")
query_features = get_query_features(query_image_path)
print("Finding the most similar image...")
best_match_file, similarity_score = find_most_similar_image(query_features, features_dict)
print(f"Most similar image: {best_match_file}")
print(f"Similarity score: {similarity_score:.4f}")
best_image = Image.open(f"/media/syun/ssd02/python_learning/apple/qiita_project_AIMv2/coco_image/val2017/{best_match_file}")
query_search_image = Image.open(query_image_path)
best_image.show()
query_search_image.show()
コード解説
-
クエリ画像の処理:
- 指定されたクエリ画像を読み込み、AIMv2モデルを使って特徴量を抽出。
-
類似画像の検索:
- クエリ画像の特徴量と保存済みのデータセット特徴量をコサイン類似度で比較。
- 最も類似度が高い画像を特定。
- 下の実行結果で
Similarity score: 0.9307に精度が書いています - 下の実行結果で
Most similar image: 000000454661.jpgはcocodatasetのval2017内での一番類似度高い結果のファイル名を出力していること
ぜひ動画を見てください。
- 下の実行結果で
-
保存など画像パスは自分が置いていた場所に置き換えてください
実行結果
クエリ画像を指定してスクリプトを実行すると、データセット内で最も似ている画像と類似度スコアが出力されます。
(.venv) syun@syun:/media/syun/ssd02/python_learning/apple/qiita_project_AIMv2$ python3 aimv2-large-patch14-224/image_search.py /media/syun/ssd02/python_learning/apple/qiita_project_AIMv2/test_search_image/gtr.jpg --features /media/syun/ssd02/python_learning/apple/qiita_project_AIMv2/model/coco_features.pt
/media/syun/ssd02/python_learning/apple/qiita_project_AIMv2/.venv/lib/python3.12/site-packages/torch/cuda/__init__.py:129: UserWarning: CUDA initialization: CUDA unknown error - this may be due to an incorrectly set up environment, e.g. changing env variable CUDA_VISIBLE_DEVICES after program start. Setting the available devices to be zero. (Triggered internally at ../c10/cuda/CUDAFunctions.cpp:108.)
return torch._C._cuda_getDeviceCount() > 0
Loading features from /media/syun/ssd02/python_learning/apple/qiita_project_AIMv2/model/coco_features.pt...
/media/syun/ssd02/python_learning/apple/qiita_project_AIMv2/aimv2-large-patch14-224/image_search.py:60: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
features_dict = torch.load(features_path)
Extracting features from query image: /media/syun/ssd02/python_learning/apple/qiita_project_AIMv2/test_search_image/gtr.jpg...
Finding the most similar image...
Most similar image: 000000454661.jpg
Similarity score: 0.9307
音声とか何もないし、説明字幕もない動画ですが、ぜひ見てください。
一番ビックリしたのは新しいGTRの画像を入力し、古いGTRの類似度も認識してくれてたことです。
皆さんの反応次第で今後音声と字幕もある可能性もあります。登録よろしくお願いします
AIMv2の応用例GPTにも聞いていたが、中々面白いです。チャレンジしてみようと思います。
AIMv2は画像検索以外にも、以下のようなタスクに応用可能です：
- テキストと画像の関連性推論
- オープンボキャブラリ物体検出
- マルチモーダルアプリケーション（例: 画像キャプション生成）
まとめ
この記事では、AppleのAIMv2モデルを使って画像検索システムを構築する方法を解説しました。AIMv2の高度な特徴を活用することで、画像検索だけでなく、さまざまなビジョンタスクへの応用が可能です。ぜひ試してみてください！
今回も最後まで見てくれてありがとうございます。
次回はマルチモーダルを試そうと思っているが、AppleのISSUESの対応を待っているとことです。
またよろしくお願いします。