YOLO を CLIP とともに使用して検索を改善する

この記事では、YOLO のような物体検出モデルと CLIP のようなマルチモーダル埋め込みモデルを併用して画像検索を改善する方法を見ていきます。

アイデアは次のとおりです: CLIP 画像の取得は次のように機能します: CLIP モデルを使用して画像を埋め込み、ベクトル データベースなどのどこかに保存します。次に、推論中に、クエリ画像またはプロンプトを使用してそれを埋め込み、取得可能な保存された埋め込みから最も近い画像を見つけることができます。問題は、埋め込まれた画像にオブジェクトが多すぎる場合、または一部のオブジェクトが背景にあり、それでもシステムにそれらを取得させたい場合です。これは、CLIP が画像全体を埋め込むためです。文埋め込みモデルに対する単語埋め込みモデルのようなものと考えてください。画像内のオブジェクトに相当する単語を検索できるようにしたいと考えています。したがって、解決策は、オブジェクト検出モデルを使用して画像をさまざまなオブジェクトに分解することです。次に、これらの分解された画像を埋め込みますが、親画像にリンクします。これにより、作物を取得し、その作物の由来となった親を取得できるようになります。 どのように機能するかを見てみましょう。

依存関係をインストールしてインポートする

!pip install -q ultralytics torch matplotlib numpy pillow zipfile36 transformers

from ultralytics import YOLO
import matplotlib.pyplot as plt
from PIL import pillow
import os
from Zipfile import Zipfile, BadZipFile
import torch
from transformers import CLIPProcessor, CLIPModel, CLIPVisionModelWithProjection, CLIPTextModelWithProjection

COCO データセットをダウンロードして解凍します。

!wget http://images.cocodataset.org/zips/val2017.zip -O coco_val2017.zip

def extract_zip_file(extract_path):
    try:
        with ZipFile(extract_path+".zip") as zfile:
            zfile.extractall(extract_path)
        # remove zipfile
        zfileTOremove=f"{extract_path}"+".zip"
        if os.path.isfile(zfileTOremove):
            os.remove(zfileTOremove)
        else:
            print("Error: %s file not found" % zfileTOremove)
    except BadZipFile as e:
        print("Error:", e)

extract_val_path = "./coco_val2017"
extract_zip_file(extract_val_path)

次に、いくつかの画像を取得して、例のリストを作成します。

source = ['coco_val2017/val2017/000000000139.jpg', '/content/coco_val2017/val2017/000000000632.jpg', '/content/coco_val2017/val2017/000000000776.jpg', '/content/coco_val2017/val2017/000000001503.jpg', '/content/coco_val2017/val2017/000000001353.jpg', '/content/coco_val2017/val2017/000000003661.jpg']

YOLO モデルと CLIP モデルを開始する

この例では、最新の Ultralytics Yolo10x モデルと OpenAI Clip-vit-base-patch32 を使用します。

device = "cuda"

 # YOLO Model
model = YOLO('yolov10x.pt')

# Clip model
model_id = "openai/clip-vit-base-patch32"
image_model = CLIPVisionModelWithProjection.from_pretrained(model_id, device_map = device)
text_model = CLIPTextModelWithProjection.from_pretrained(model_id, device_map = device)
processor = CLIPProcessor.from_pretrained(model_id)

検出モデルの実行

results = model(source=source, device = "cuda")

このコード スニペットで結果を示してみましょう

# Visualize the results
fig, ax = plt.subplots(2, 3, figsize=(15, 10))

for i, r in enumerate(results):
    # Plot results image
    im_bgr = r.plot()  # BGR-order numpy array
    im_rgb = Image.fromarray(im_bgr[..., ::-1])  # RGB-order PIL image

    ax[i%2, i//2].imshow(im_rgb)
    ax[i%2, i//2].set_title(f"Image {i+1}")

つまり、YOLO モデルが画像内のオブジェクトの検出に非常にうまく機能していることがわかります。モニターを TV としてタグ付けする箇所でいくつかの間違いを犯します。でもそれでいいのです。 CLIP を使用して推論を行うため、YOLO が割り当てる実際のクラスはそれほど重要ではありません。

いくつかのヘルパー クラスの定義

class CroppedImage:

  def __init__(self, parent, box, cls):

    self.parent = parent
    self.box = box
    self.cls = cls

  def display(self, ax = None):
    im_rgb = Image.open(self.parent)
    cropped_image = im_rgb.crop(self.box)

    if ax is not None:
      ax.imshow(cropped_image)
      ax.set_title(self.cls)
    else:
      plt.figure(figsize=(10, 10))
      plt.imshow(cropped_image)
      plt.title(self.cls)
      plt.show()

  def get_cropped_image(self):
    im_rgb = Image.open(self.parent)
    cropped_image = im_rgb.crop(self.box)
    return cropped_image

  def __str__(self):
    return f"CroppedImage(parent={self.parent}, boxes={self.box}, cls={self.cls})"

  def __repr__(self):
    return self.__str__()

class YOLOImage:
  def __init__(self, image_path, cropped_images):
    self.image_path = str(image_path)
    self.cropped_images = cropped_images

  def get_image(self):
    return Image.open(self.image_path)

  def get_caption(self):
    cls  =[]
    for cropped_image in self.cropped_images:
      cls.append(cropped_image.cls)

    unique_cls = set(cls)
    count_cls = {cls: cls.count(cls) for cls in unique_cls}

    count_string = " ".join(f"{count} {cls}," for cls, count in count_cls.items())
    return "this image contains " + count_string

  def __str__(self):
    return self.__repr__()

  def __repr__(self):
    cls  =[]
    for cropped_image in self.cropped_images:
      cls.append(cropped_image.cls)

    return f"YOLOImage(image={self.image_path}, cropped_images={cls})"

class ImageEmbedding:
  def __init__(self, image_path, embedding, cropped_image = None):
    self.image_path = image_path
    self.cropped_image = cropped_image
    self.embedding = embedding

CroppedImage クラス

CroppedImage クラスは、大きな親画像から切り取られた画像の一部を表します。これは、親画像へのパス、切り抜き領域を定義する境界ボックス、およびクラス ラベル (「猫」や「犬」など) で初期化されます。このクラスには、トリミングされた画像を表示し、それを画像オブジェクトとして取得するメソッドが含まれています。この表示方法では、指定された軸上で、または新しい図を作成することによって、切り取られた部分を視覚化できるため、さまざまなユースケースに多用途に使用できます。さらに、オブジェクトの簡単で有益な文字列表現のために、__str__ メソッドと __repr__ メソッドが実装されています。

YOLOImage クラス

YOLOImage クラスは、YOLO オブジェクト検出モデルで処理された画像を処理するように設計されています。これは、元の画像へのパスと、画像内で検出されたオブジェクトを表す CroppedImage インスタンスのリストを受け取ります。このクラスは、画像全体を開いて表示し、画像内で検出されたオブジェクトを要約するキャプションを生成するメソッドを提供します。キャプション メソッドは、トリミングされた画像から固有のクラス ラベルを集計してカウントし、画像の内容の簡潔な説明を提供します。このクラスは、物体検出タスクの結果の管理と解釈に特に役立ちます。

ImageEmbedding クラス

ImageEmbedding クラスには、画像とそれに関連付けられた埋め込みがあり、画像の特徴を数値的に表現したものです。このクラスは、画像へのパス、埋め込みベクトル、および埋め込みが画像の特定の切り取られた部分に対応する場合はオプションで CroppedImage インスタンスを使用して初期化できます。 ImageEmbedding クラスは、計算された特徴とともに画像データを保存およびアクセスするための構造化された方法を提供するため、画像の類似性、分類、および取得に関連するタスクに不可欠です。この統合により、効率的な画像処理と機械学習のワークフローが促進されます。

各画像をトリミングし、YOLOImage オブジェクトのリストを作成します。

yolo_images: list[YOLOImage]= []

names= model.names

for i, r in enumerate(results):
  crops:list[CroppedImage] = []
  boxes = r.boxes
  classes = r.boxes.cls
  for j, box in enumerate(r.boxes):
    box = tuple(box.xyxy.flatten().cpu().numpy())
    cropped_image = CroppedImage(parent = r.path, box = box, cls = names[classes[j].int().item()])
    crops.append(cropped_image)
  yolo_images.append(YOLOImage(image_path=r.path, cropped_images=crops))

CLIPを使用して画像を埋め込む

image_embeddings = []

for image in yolo_images:
  input = processor.image_processor(images= image.get_image(), return_tensors = 'pt')
  input.to(device)
  embeddings = image_model(pixel_values = input.pixel_values).image_embeds
  embeddings = embeddings/embeddings.norm(p=2, dim = -1, keepdim = True) # Normalize the embeddings
  image_embedding = ImageEmbedding(image_path = image.image_path, embedding = embeddings)
  image_embeddings.append(image_embedding)

  for cropped_image in image.cropped_images:
    input = processor.image_processor(images= cropped_image.get_cropped_image(), return_tensors = 'pt')
    input.to(device)
    embeddings = image_model(pixel_values = input.pixel_values).image_embeds
    embeddings = embeddings/embeddings.norm(p=2, dim = -1, keepdim = True) # Normalize the embeddings

    image_embedding = ImageEmbedding(image_path = image.image_path, embedding = embeddings, cropped_image = cropped_image)
    image_embeddings.append(image_embedding)

   **image_embeddings_tensor = torch.stack([image_embedding.embedding for image_embedding in image_embeddings]).squeeze()**

必要に応じて、これらの画像埋め込みを取得し、ベクトル データベースに保存できるようになりました。ただし、この例では、内積手法を使用して類似性をチェックし、画像を取得するだけです。

検索

query = "image of a flowerpot"

text_embedding = processor.tokenizer(query, return_tensors="pt").to(device)
text_embedding = text_model(**text_embedding).text_embeds

similarities = (torch.matmul(text_embedding, image_embeddings_tensor.T)).flatten().detach().cpu().numpy()

# get the top 5 similar images
k = 5
top_k_indices = similarities.argsort()[-k:]

# Display the top 5 results
fig, ax = plt.subplots(2, 5, figsize=(20, 5))
for i, index in enumerate(top_k_indices):
  if image_embeddings[index].cropped_image is not None:
    image_embeddings[index].cropped_image.display(ax = ax[0][i])
  else:
  ax[0][i].imshow(Image.open(image_embeddings[index].image_path))
  ax[1][i].imshow(Image.open(image_embeddings[index].image_path))
  ax[0][i].axis('off')
  ax[1][i].axis('off')
  ax[1][i].set_title("Original Image")
plt.show()

背景に隠れてしまった小さな植物も検索できることがわかります。また、元の画像も埋め込んでいるため、結果として元の画像を取得する場合もあります。

これは非常に強力なテクニックになる可能性があります。また、独自の画像の検出と埋め込み用に両方のモデルを微調整して、パフォーマンスをさらに向上させることもできます。

欠点の 1 つは、検出されたすべてのオブジェクトに対して CLIP モデルを実行する必要があることです。これを軽減する 1 つの方法は、YOLO が生成するボックスの数を制限することです。

このリンクで Colab のコードをチェックアウトできます。

---

接続したいですか?

?私のウェブサイト

?私のツイッター

?私のLinkedIn

以上がYOLO を CLIP とともに使用して検索を改善するの詳細内容です。詳細については、PHP 中国語 Web サイトの他の関連記事を参照してください。