검색 향상을 위해 CLIP과 YOLO 사용-파이썬 튜토리얼-php.cn

이 기사에서는 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 모델 시작

이 예에서는 OpenAI Clip-vit-base-patch32와 함께 최신 Ultralytics Yolo10x 모델을 사용하겠습니다.

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

로그인 후 복사

Using YOLO with CLIP to improve Retrieval

그래서 우리는 YOLO 모델이 이미지 속 물체를 감지하는 데 꽤 잘 작동한다는 것을 알 수 있습니다. 모니터를 TV로 태그한 경우 몇 가지 실수가 있습니다. 하지만 괜찮습니다. YOLO가 할당하는 실제 클래스는 CLIP을 사용하여 추론을 수행하므로 그다지 중요하지 않습니다.

일부 도우미 클래스 정의

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 클래스는 더 큰 상위 이미지에서 자른 이미지의 일부를 나타냅니다. 이는 상위 이미지의 경로, 자르기 영역을 정의하는 경계 상자 및 클래스 라벨(예: "cat" 또는 "dog")로 초기화됩니다. 이 클래스에는 자른 이미지를 표시하고 이를 이미지 객체로 검색하는 메서드가 포함되어 있습니다. 표시 방법을 사용하면 제공된 축에서 또는 새 그림을 생성하여 잘린 부분을 시각화할 수 있으므로 다양한 사용 사례에 맞게 다용도로 사용할 수 있습니다. 또한 객체에 대한 쉽고 유익한 문자열 표현을 위해 __str__ 및 __repr__ 메소드가 구현되었습니다.

YOLO이미지 클래스

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

로그인 후 복사

Using YOLO with CLIP to improve Retrieval