검색 향상을 위해 CLIP과 YOLO 사용
- WBOY원래의
- 2024-08-05 21:58:521182검색
이 기사에서는 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}")
그래서 우리는 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()
배경에 숨겨져 있는 작은 식물도 검색할 수 있다는 것을 알 수 있습니다. 또한 원본 이미지를 삽입하기 때문에 결과적으로 원본 이미지를 가져오는 경우도 있습니다.
이것은 매우 강력한 기술이 될 수 있습니다. 또한 자신의 이미지에 대한 감지 및 삽입 모델을 모두 미세 조정하고 성능을 더욱 향상시킬 수 있습니다.
한 가지 단점은 감지된 모든 개체에 대해 CLIP 모델을 실행해야 한다는 것입니다. 이를 완화하는 한 가지 방법은 YOLO가 생산하는 상자 수를 제한하는 것입니다.
이 링크에서 Colab의 코드를 확인할 수 있습니다.
연결하시겠습니까?
?내 웹사이트
?내 트위터
?내 LinkedIn
위 내용은 검색 향상을 위해 CLIP과 YOLO 사용의 상세 내용입니다. 자세한 내용은 PHP 중국어 웹사이트의 기타 관련 기사를 참조하세요!