Поиск по тексту в изображениях с помощью Milvus

Поиск по тексту и изображению - это передовая технология, позволяющая пользователям искать изображения по текстовым описаниям на естественном языке. Она использует предварительно обученную мультимодальную модель для преобразования текста и изображений во вложения в общее семантическое пространство, что позволяет проводить сравнение на основе сходства.

В этом учебном пособии мы рассмотрим, как реализовать поиск изображений по тексту с помощью модели CLIP (Contrastive Language-Image Pretraining) от OpenAI и Milvus. Мы будем генерировать вкрапления изображений с помощью CLIP, хранить их в Milvus и выполнять эффективный поиск по сходству.

Предварительные условия

Перед началом работы убедитесь, что у вас готовы все необходимые пакеты и данные примеров.

Установите зависимости

• pymilvus>=2.4.2 для работы с базой данных Milvus
• clip для работы с моделью CLIP
• pillow для обработки и визуализации изображений

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

Загрузка данных примера

В качестве примера изображений мы будем использовать подмножество набора данных ImageNet (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 (рекомендуется для удобства): Установите URI на локальный файл, например ./milvus.db. Это автоматически задействует Milvus Lite для хранения всех данных в одном файле.

• Docker или Kubernetes (для больших массивов данных): Для работы с большими массивами данных разверните более производительный сервер Milvus с помощью Docker или Kubernetes. В этом случае для подключения используйте URI сервера, например http://localhost:19530.

• Zilliz Cloud (управляемый сервис): Если вы используете Zilliz Cloud, полностью управляемый облачный сервис Milvus, задайте публичную конечную точку в качестве URI и ключ API в качестве токена.

from pymilvus import MilvusClient

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

Начало работы

Теперь, когда у вас есть необходимые зависимости и данные, пришло время настроить экстракторы функций и начать работу с Milvus. В этом разделе мы рассмотрим основные этапы создания системы поиска по тексту и изображению. Наконец, мы продемонстрируем, как получать и визуализировать изображения на основе текстовых запросов.

Определите экстракторы признаков

Мы будем использовать предварительно обученную модель CLIP для создания вкраплений изображений и текста. В этом разделе мы загрузим предварительно обученный ViT-B/32 вариант CLIP и определим вспомогательные функции для кодирования изображений и текста:

• 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: : Первичное поле с включенным автоматическим идентификатором.

• 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