Milvus를 사용한 텍스트-이미지 검색

텍스트-이미지 검색은 사용자가 자연어 텍스트 설명을 사용해 이미지를 검색할 수 있는 고급 기술입니다. 사전 학습된 멀티모달 모델을 활용하여 텍스트와 이미지를 모두 공유 의미 공간의 임베딩으로 변환하여 유사성 기반 비교를 가능하게 합니다.

이 튜토리얼에서는 OpenAI의 CLIP(대조 언어-이미지 사전 학습) 모델과 Milvus를 사용하여 텍스트 기반 이미지 검색을 구현하는 방법을 살펴봅니다. CLIP으로 이미지 임베딩을 생성하고 Milvus에 저장한 후 효율적인 유사도 검색을 수행해 보겠습니다.

전제 조건

시작하기 전에 필요한 모든 패키지와 예제 데이터가 준비되어 있는지 확인하세요.

종속성 설치

• Milvus 데이터베이스와 상호 작용하기 위한pymilvus>=2.4.2
• CLIP 모델 작업을 위한clip
• 이미지 처리 및 시각화를 위한pillow

$ pip install --upgrade pymilvus pillow
$ pip install git+https://github.com/openai/CLIP.git

예제 데이터 다운로드

이미지넷 데이터 세트의 하위 집합(100개 클래스, 각 클래스당 10개 이미지)을 예제 이미지로 사용합니다. 다음 명령은 예제 데이터를 다운로드하여 로컬 폴더 ./images_folder 에 추출합니다:

$ wget https://github.com/towhee-io/examples/releases/download/data/reverse_image_search.zip
$ unzip -q reverse_image_search.zip -d images_folder

Milvus 설정

계속 진행하기 전에 Milvus 서버를 설정하고 URI(선택 사항, 토큰)를 사용하여 연결합니다:

• Milvus Lite(편의상 권장): ./milvus.db와 같은 로컬 파일로 URI를 설정합니다. 이렇게 하면 Milvus Lite를 자동으로 활용하여 모든 데이터를 단일 파일에 저장합니다.

• Docker 또는 Kubernetes(대규모 데이터의 경우): 대규모 데이터 세트를 처리하려면 Docker 또는 Kubernetes를 사용하여 더 성능이 뛰어난 Milvus 서버를 배포하세요. 이 경우 http://localhost:19530 같은 서버 URI를 사용하여 연결하세요.

• 질리즈 클라우드(관리형 서비스): Milvus의 완전 관리형 클라우드 서비스인 Zilliz Cloud를 사용하는 경우, 퍼블릭 엔드포인트를 URI로 설정하고 API 키를 토큰으로 설정합니다.

from pymilvus import MilvusClient

milvus_client = MilvusClient(uri="milvus.db")

시작하기

이제 필요한 종속성과 데이터가 준비되었으므로 이제 특징 추출기를 설정하고 Milvus로 작업을 시작할 차례입니다. 이 섹션에서는 텍스트-이미지 검색 시스템을 구축하는 주요 단계를 안내합니다. 마지막으로 텍스트 쿼리를 기반으로 이미지를 검색하고 시각화하는 방법을 시연해 보겠습니다.

특징 추출기 정의하기

사전 학습된 CLIP 모델을 사용하여 이미지와 텍스트 임베딩을 생성합니다. 이 섹션에서는 사전 학습된 CLIP의 ViT-B/32 변형을 로드하고 이미지와 텍스트 인코딩을 위한 헬퍼 함수를 정의합니다:

• encode_image(image_path): 이미지를 특징 벡터로 처리하고 인코딩합니다.
• encode_text(text): 텍스트 쿼리를 특징 벡터로 인코딩

두 함수 모두 정확한 코사인 유사도 계산에 필수적인 벡터를 단위 길이로 변환하여 일관된 비교를 보장하기 위해 출력 피처를 정규화합니다.

import clip
from PIL import Image


# Load CLIP model
model_name = "ViT-B/32"
model, preprocess = clip.load(model_name)
model.eval()


# Define a function to encode images
def encode_image(image_path):
    image = preprocess(Image.open(image_path)).unsqueeze(0)
    image_features = model.encode_image(image)
    image_features /= image_features.norm(
        dim=-1, keepdim=True
    )  # Normalize the image features
    return image_features.squeeze().tolist()


# Define a function to encode text
def encode_text(text):
    text_tokens = clip.tokenize(text)
    text_features = model.encode_text(text_tokens)
    text_features /= text_features.norm(
        dim=-1, keepdim=True
    )  # Normalize the text features
    return text_features.squeeze().tolist()

데이터 수집

시맨틱 이미지 검색을 활성화하려면 먼저 모든 이미지에 대한 임베딩을 생성하고 효율적인 색인 및 검색을 위해 이를 벡터 데이터베이스에 저장해야 합니다. 이 섹션에서는 Milvus로 이미지 데이터를 수집하는 단계별 가이드를 제공합니다.

1. Milvus 컬렉션 생성

이미지 임베딩을 저장하기 전에 Milvus 컬렉션을 생성해야 합니다. 다음 코드는 기본 COSINE 메트릭 유형으로 빠른 설정 모드에서 컬렉션을 생성하는 방법을 보여줍니다. 컬렉션에는 다음 필드가 포함됩니다:

• id: 자동 ID가 활성화된 기본 필드.

• vector: 부동 소수점 벡터 임베딩을 저장하는 필드.

사용자 정의 스키마가 필요한 경우 Milvus 설명서를 참조하여 자세한 지침을 확인하세요.

collection_name = "image_collection"

# Drop the collection if it already exists
if milvus_client.has_collection(collection_name):
    milvus_client.drop_collection(collection_name)

# Create a new collection in quickstart mode
milvus_client.create_collection(
    collection_name=collection_name,
    dimension=512,  # this should match the dimension of the image embedding
    auto_id=True,  # auto generate id and store in the id field
    enable_dynamic_field=True,  # enable dynamic field for scalar fields
)

2. Milvus에 데이터 삽입

이 단계에서는 미리 정의된 이미지 인코더를 사용하여 예제 데이터 디렉토리에 있는 모든 JPEG 이미지에 대한 임베딩을 생성합니다. 그런 다음 이러한 임베딩은 해당 파일 경로와 함께 Milvus 컬렉션에 삽입됩니다. 컬렉션의 각 항목은 다음과 같이 구성됩니다:

• 임베딩 벡터: 이미지의 숫자 표현입니다. vector 필드에 저장됩니다.
• 파일 경로: 참조용 이미지 파일의 위치입니다. filepath 필드에 동적 필드로 저장됩니다.

import os
from glob import glob


image_dir = "./images_folder/train"
raw_data = []

for image_path in glob(os.path.join(image_dir, "**/*.JPEG")):
    image_embedding = encode_image(image_path)
    image_dict = {"vector": image_embedding, "filepath": image_path}
    raw_data.append(image_dict)
insert_result = milvus_client.insert(collection_name=collection_name, data=raw_data)

print("Inserted", insert_result["insert_count"], "images into Milvus.")

Inserted 1000 images into Milvus.

검색 수행

이제 예제 텍스트 쿼리를 사용하여 검색을 실행해 보겠습니다. 이렇게 하면 주어진 텍스트 설명과의 의미적 유사성을 기준으로 가장 관련성이 높은 이미지가 검색됩니다.

query_text = "a white dog"
query_embedding = encode_text(query_text)

search_results = milvus_client.search(
    collection_name=collection_name,
    data=[query_embedding],
    limit=10,  # return top 10 results
    output_fields=["filepath"],  # return the filepath field
)

결과를 시각화합니다:

from IPython.display import display


width = 150 * 5
height = 150 * 2
concatenated_image = Image.new("RGB", (width, height))

result_images = []
for result in search_results:
    for hit in result:
        filename = hit["entity"]["filepath"]
        img = Image.open(filename)
        img = img.resize((150, 150))
        result_images.append(img)

for idx, img in enumerate(result_images):
    x = idx % 5
    y = idx // 5
    concatenated_image.paste(img, (x * 150, y * 150))
print(f"Query text: {query_text}")
print("\nSearch results:")
display(concatenated_image)

Query text: a white dog

Search results:

png