Menggunakan YOLO dengan CLIP untuk menambah baik Retrieval

Dalam artikel ini kita akan melihat cara kita boleh menggunakan model pengesanan objek seperti YOLO bersama-sama dengan model benam berbilang mod seperti CLIP untuk menjadikan pengambilan imej lebih baik.

Inilah ideanya: Pengambilan imej CLIP berfungsi seperti berikut: Kami membenamkan imej yang kami ada menggunakan model CLIP dan menyimpannya di suatu tempat, seperti dalam pangkalan data vektor. Kemudian, semasa inferens, kita boleh menggunakan imej pertanyaan atau gesaan, membenamkannya dan mencari imej yang paling hampir daripada benam yang disimpan yang boleh diambil semula. Masalahnya ialah apabila imej terbenam mempunyai terlalu banyak objek atau beberapa objek berada di latar belakang, dan kami masih mahu sistem kami mendapatkannya semula. Ini kerana CLIP membenamkan imej secara keseluruhan. Fikirkannya seperti model benam perkataan kepada model benam ayat. Kami mahu dapat mencari perkataan yang setara dengan objek dalam imej. Jadi, penyelesaiannya adalah untuk menguraikan imej kepada objek yang berbeza menggunakan model pengesanan objek. Kemudian, benamkan imej terurai ini tetapi pautkannya ke imej induknya. Ini akan membolehkan kami mengambil semula tanaman dan mendapatkan induk dari mana tanaman itu berasal. Mari lihat bagaimana ia berfungsi.

Pasang Dependencies dan importnya

!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

Muat turun Set Data COCO dan nyahzip

!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)

Kami kemudiannya boleh mengambil beberapa imej dan membuat senarai contoh.

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']

Mulakan model YOLO dan Model CLIP

Dalam contoh ini, kami akan menggunakan model Ultralytics Yolo10x terkini bersama-sama dengan 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)

Menjalankan model pengesanan

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

Mari tunjukkan kepada kami hasil dengan coretan kod ini

# 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}")

Jadi kita dapat melihat bahawa model YOLO berfungsi dengan baik dalam mengesan objek dalam imej. Ia melakukan beberapa kesilapan apabila ia telah menandakan monitor sebagai TV. Tetapi itu tidak mengapa. Kelas sebenar yang YOLO tetapkan tidak begitu penting kerana kami akan menggunakan CLIP untuk membuat inferens.

Menentukan beberapa Kelas pembantu

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

Kelas CroppedImage

Kelas CroppedImage mewakili sebahagian daripada imej yang dipangkas daripada imej induk yang lebih besar. Ia dimulakan dengan laluan ke imej induk, kotak sempadan yang mentakrifkan kawasan tanaman dan label kelas (mis., "kucing" atau "anjing"). Kelas ini termasuk kaedah untuk memaparkan imej yang dipangkas dan untuk mendapatkannya semula sebagai objek imej. Kaedah paparan membolehkan untuk menggambarkan bahagian yang dipangkas sama ada pada paksi yang disediakan atau dengan mencipta angka baharu, menjadikannya serba boleh untuk kes penggunaan yang berbeza. Selain itu, kaedah __str__ dan __repr__ dilaksanakan untuk perwakilan rentetan objek yang mudah dan bermaklumat.

Kelas YOLOImage

Kelas YOLOImage direka untuk mengendalikan imej yang diproses dengan model pengesanan objek YOLO. Ia mengambil laluan ke imej asal dan senarai contoh CroppedImage yang mewakili objek yang dikesan dalam imej. Kelas menyediakan kaedah untuk membuka dan memaparkan imej penuh dan menjana kapsyen yang meringkaskan objek yang dikesan dalam imej. Kaedah kapsyen mengagregat dan mengira label kelas unik daripada imej yang dipangkas, memberikan penerangan ringkas tentang kandungan imej. Kelas ini amat berguna untuk mengurus dan mentafsir hasil daripada tugas pengesanan objek.

Kelas Pembenaman Imej

Kelas ImageEmbedding mempunyai imej dan pembenaman yang berkaitan, yang merupakan perwakilan berangka bagi ciri imej. Kelas ini boleh dimulakan dengan laluan ke imej, vektor benam, dan sebagai pilihan contoh CroppedImage jika benam sepadan dengan bahagian imej yang dipangkas tertentu. Kelas ImageEmbedding adalah penting untuk tugas yang melibatkan persamaan imej, pengelasan dan perolehan semula, kerana ia menyediakan cara berstruktur untuk menyimpan dan mengakses data imej bersama ciri yang dikiranya. Penyepaduan ini memudahkan pemprosesan imej dan aliran kerja pembelajaran mesin yang cekap.

Pangkas setiap imej dan buat senarai Objek 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))

Benamkan Imej menggunakan 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()**

Kami kini boleh mengambil benam imej ini dan menyimpan dalam pangkalan data vektor jika kami mahu. Tetapi dalam contoh ini kita hanya akan menggunakan teknik produk titik dalam untuk menyemak persamaan dan mendapatkan semula imej.

Pengambilan semula

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()

Anda dapat melihat bahawa kami dapat mendapatkan semula walaupun tumbuhan kecil yang tersembunyi di latar belakang. Juga kadangkala ia menarik imej asal sebagai hasilnya kerana kami juga membenamkannya .

Ini boleh menjadi teknik yang sangat berkuasa. Anda juga boleh memperhalusi kedua-dua model untuk pengesanan dan pembenaman untuk imej anda sendiri dan meningkatkan lagi prestasi.

Satu kelemahan ialah kita perlu menjalankan model CLIP pada semua objek yang dikesan. Satu cara untuk mengurangkan perkara ini ialah dengan mengehadkan bilangan kotak yang YOLO hasilkan.

Anda boleh menyemak kod pada Colab di pautan ini.

---

Mahu menyambung?

?Tapak Web Saya

?Twitter saya

? LinkedIn saya

Atas ialah kandungan terperinci Menggunakan YOLO dengan CLIP untuk menambah baik Retrieval. Untuk maklumat lanjut, sila ikut artikel berkaitan lain di laman web China PHP!