Hugging Faceデータセットとトランスフォーマーを使用した画像の類似性

'Hugging Faceデータセットとトランスフォーマーを使った画像の類似性'

この投稿では、🤗 Transformersを使用して画像の類似性システムを構築する方法を学びます。クエリ画像と候補画像の類似性を見つけることは、逆画像検索などの情報検索システムの重要なユースケースです。システムが答えようとしているのは、クエリ画像と候補画像セットが与えられた場合、どの画像がクエリ画像に最も類似しているかということです。

このシステムの構築には、このシステムの構築時に便利な並列処理をシームレスにサポートする🤗のdatasetsライブラリを活用します。

この投稿では、ViTベースのモデル（ nateraw/vit-base-beans ）と特定のデータセット（Beans）を使用していますが、ビジョンモダリティをサポートし、他の画像データセットを使用するために拡張することもできます。試してみることができるいくつかの注目すべきモデルには次のものがあります：

• Hugging Faceがフランスのデータ保護機関の強化サポートプログラムに選ばれました
• DuckDB Hugging Face Hubに保存されている50,000以上のデータセットを分析する
• 基礎モデルは人間のようにデータにラベルを付けることができますか？

• Swin Transformer
• ConvNeXT
• RegNet

また、投稿で紹介されているアプローチは、他のモダリティにも拡張できる可能性があります。

完全に動作する画像の類似性システムを学習するには、最初に2つの画像間の類似性をどのように定義するかを定義する必要があります。

このシステムを構築するためには、まず与えられた画像の密な表現（埋め込み）を計算し、その後、余弦類似性指標を使用して2つの画像の類似性を決定する一般的な方法があります。

この投稿では、画像をベクトル空間で表現するために「埋め込み」を使用します。これにより、画像の高次元ピクセル空間（たとえば224 x 224 x 3）を意味のある低次元空間（たとえば768）にうまく圧縮する方法が得られます。これによる主な利点は、後続のステップでの計算時間の削減です。

画像から埋め込みを計算するために、入力画像をベクトル空間で表現する方法について理解しているビジョンモデルを使用します。このタイプのモデルは画像エンコーダとも呼ばれます。

モデルをロードするために、AutoModelクラスを活用します。これにより、Hugging Face Hubから互換性のあるモデルチェックポイントをロードするためのインターフェースが提供されます。モデルと共に、データ前処理に関連するプロセッサもロードします。

from transformers import AutoFeatureExtractor, AutoModel


model_ckpt = "nateraw/vit-base-beans"
extractor = AutoFeatureExtractor.from_pretrained(model_ckpt)
model = AutoModel.from_pretrained(model_ckpt)

この場合、チェックポイントはbeansデータセットでビジョントランスフォーマーベースのモデルをファインチューニングして取得しました。

ここで生じるかもしれないいくつかの質問：

Q1 ：なぜAutoModelForImageClassificationを使用しなかったのですか？

これは、画像の密な表現ではなく、離散的なカテゴリを得るためです。これはAutoModelForImageClassificationが提供するものです。

Q2 ：なぜ特にこのチェックポイントを使用するのですか？

前述のように、特定のデータセットを使用してシステムを構築しています。したがって、一般的なモデル（たとえばImageNet-1kデータセットでトレーニングされたモデルなど）ではなく、使用されているデータセットでファインチューニングされたモデルを使用する方が良いです。そのようにすることで、基礎となるモデルは入力画像をより良く理解することができます。

なお、教師あり学習から得られたチェックポイントを使用する必要はありません。事前に自己教師あり学習でうまく学習されたモデルは、印象的な検索パフォーマンスを発揮することができます。

今や埋め込みを計算するためのモデルがあるので、クエリするためのいくつかの候補画像が必要です。

候補画像のデータセットのロード

しばらくすると、候補画像をハッシュにマッピングするハッシュテーブルを構築します。クエリ時には、これらのハッシュテーブルを使用します。ハッシュテーブルについては、該当するセクションで詳しく説明しますが、今のところ、候補画像のセットを持つために、beansデータセットのtrain分割を使用します。

from datasets import load_dataset


dataset = load_dataset("beans")

これはトレーニングデータセットの一つのサンプルの例です：

データセットには3つの特徴があります：

dataset["train"].features
>>> {'image_file_path': Value(dtype='string', id=None),
 'image': Image(decode=True, id=None),
 'labels': ClassLabel(names=['angular_leaf_spot', 'bean_rust', 'healthy'], id=None)}

画像の類似性システムをデモンストレーションするために、候補画像データセットから100のサンプルを使用して、全体的なランタイムを短く保ちます。

num_samples = 100
seed = 42
candidate_subset = dataset["train"].shuffle(seed=seed).select(range(num_samples))

類似画像の検索プロセス

以下に、類似画像の取得プロセスの概要を示します。

上記の図を少し詳しく見てみると、次のような要素があります：

1. 候補画像から埋め込みを抽出し、それらを行列に格納します。
2. クエリ画像から埋め込みを抽出します。
3. 埋め込み行列（ステップ1で計算された）を反復処理し、クエリ埋め込みと現在の候補埋め込みの類似スコアを計算します。通常、候補画像の識別子と類似スコアの対応を保持する辞書のようなマッピングを維持します。
4. マッピング構造を類似スコアに基づいてソートし、その下にある識別子を返します。これらの識別子を使用して候補サンプルを取得します。

効率的に埋め込みを計算するために、シンプルなユーティリティを作成し、データセットの候補画像にmap()します。

import torch 

def extract_embeddings(model: torch.nn.Module):
    """埋め込みを計算するためのユーティリティ。"""
    device = model.device

    def pp(batch):
        images = batch["image"]
        # `transformation_chain`は、入力画像をモデルに適用するための前処理変換の組み合わせです。
        # 詳細については、添付のColabノートブックをチェックしてください。
        image_batch_transformed = torch.stack(
            [transformation_chain(image) for image in images]
        )
        new_batch = {"pixel_values": image_batch_transformed.to(device)}
        with torch.no_grad():
            embeddings = model(**new_batch).last_hidden_state[:, 0].cpu()
        return {"embeddings": embeddings}

    return pp

そして、次のようにextract_embeddings()をmap()します。

device = "cuda" if torch.cuda.is_available() else "cpu"
extract_fn = extract_embeddings(model.to(device))
candidate_subset_emb = candidate_subset.map(extract_fn, batched=True, batch_size=batch_size)

次に、便宜のために、候補画像の識別子を含むリストを作成します。

candidate_ids = []

for id in tqdm(range(len(candidate_subset_emb))):
    label = candidate_subset_emb[id]["labels"]

    # ユニークな識別子を作成します。
    entry = str(id) + "_" + str(label)

    candidate_ids.append(entry)

クエリ画像とすべての候補画像の埋め込み行列を使用して類似スコアを計算するために、すべての候補画像の埋め込み行列を使用します。すでに候補画像の埋め込みを計算しています。次のセルでは、それらを行列にまとめるだけです。

all_candidate_embeddings = np.array(candidate_subset_emb["embeddings"])
all_candidate_embeddings = torch.from_numpy(all_candidate_embeddings)

コサイン類似度を使用して2つの埋め込みベクトル間の類似スコアを計算します。それを使用して、クエリサンプルに対して類似した候補サンプルを取得します。

def compute_scores(emb_one, emb_two):
    """2つのベクトル間のコサイン類似度を計算します。"""
    scores = torch.nn.functional.cosine_similarity(emb_one, emb_two)
    return scores.numpy().tolist()


def fetch_similar(image, top_k=5):
    """`image`をクエリとして、`top_k`個の類似画像を取得します。"""
    # 埋め込み計算のために入力クエリ画像を準備します。
    image_transformed = transformation_chain(image).unsqueeze(0)
    new_batch = {"pixel_values": image_transformed.to(device)}

    # 埋め込みを計算します。
    with torch.no_grad():
        query_embeddings = model(**new_batch).last_hidden_state[:, 0].cpu()

    # 一度にすべての候補画像との類似スコアを計算します。
    # 候補画像の識別子とクエリ画像の類似スコアのマッピングも作成します。
    sim_scores = compute_scores(all_candidate_embeddings, query_embeddings)
    similarity_mapping = dict(zip(candidate_ids, sim_scores))
 
    # マッピング辞書をソートし、`top_k`個の候補を返します。
    similarity_mapping_sorted = dict(
        sorted(similarity_mapping.items(), key=lambda x: x[1], reverse=True)
    )
    id_entries = list(similarity_mapping_sorted.keys())[:top_k]

    ids = list(map(lambda x: int(x.split("_")[0]), id_entries))
    labels = list(map(lambda x: int(x.split("_")[-1]), id_entries))
    return ids, labels

クエリを実行する

すべてのユーティリティを備えているため、類似検索を行うことができます。次に、beansデータセットのtest分割からクエリ画像を取得しましょう：

test_idx = np.random.choice(len(dataset["test"]))
test_sample = dataset["test"][test_idx]["image"]
test_label = dataset["test"][test_idx]["labels"]

sim_ids, sim_labels = fetch_similar(test_sample)
print(f"クエリのラベル：{test_label}")
print(f"上位5つの候補のラベル：{sim_labels}")

結果は次のようになります：

クエリのラベル：0
上位5つの候補のラベル：[0, 0, 0, 0, 0]

システムが正しい類似画像のセットを取得したようです。可視化すると次のようになります：

さらなる拡張と結論

現在、動作する画像の類似性システムがあります。ただし、実際の状況では、さらに多くの候補画像を扱うことになります。これを考慮に入れると、現在の手順にはいくつかの欠点があります：

• 埋め込みをそのまま保存すると、特に数百万の候補画像を扱う場合、メモリ要件が急速に増加する可能性があります。埋め込みは768次元であるため、大規模な場合でも比較的高い次元です。
• 高次元の埋め込みは、検索部分で関連する後続の計算に直接的な影響を与えます。

意味を損なうことなく埋め込みの次元を減らすことができれば、速度と検索品質の間で良いトレードオフを維持することができます。この投稿の付属のColabノートブックでは、ランダムプロジェクションと局所的に感じるハッシュを使用してこれを達成するためのユーティリティを実装してデモンストレーションしています。

🤗 Datasetsは、類似性システムの構築プロセスをさらに簡素化するFAISSとの直接的な統合を提供しています。すでに候補画像（beansデータセット）の埋め込みを抽出し、embeddingsという名前の特徴量内に保存しているとします。次のようにして、データセットのadd_faiss_index()を使用して密なインデックスを構築できます：

dataset_with_embeddings.add_faiss_index(column="embeddings")

インデックスが構築されたら、dataset_with_embeddingsを使用して、クエリの埋め込みを指定して最も近い例を取得することができます：get_nearest_examples()：

scores, retrieved_examples = dataset_with_embeddings.get_nearest_examples(
    "embeddings", qi_embedding, k=top_k
)

このメソッドはスコアと対応する候補の例を返します。詳細については、公式のドキュメントとこのノートブックをチェックしてください。

最後に、次のSpaceを試してみることで、ミニ画像類似性アプリケーションを構築できます：

この投稿では、画像の類似性システムのクイックスタートを実行しました。この投稿が興味深いと感じた場合は、ここで議論した概念を基にさらに進んで構築することを強くお勧めします。そのようにすることで、内部の仕組みにより快適になることができます。

さらに学びたいですか？以下はあなたに役立つ追加のリソースです：

• Faiss：効率的な類似性検索のためのライブラリ
• ScaNN：効率的なベクトル類似性検索
• モバイルアプリケーション内での画像検索の統合

We will continue to update VoAGI; if you have any questions or suggestions, please contact us!