Использование ИИ для поиска вашего звездного стилиста (часть I)

Статья изначально была опубликована на сайте AI Accelerator Institute и перепубликована здесь с разрешения.

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

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

Вы можете найти изображения здесь. Помимо изображений, вам понадобится обновленная версия Python и pip install milvus pymilvus torch torchvision matplotlib. Мы используем модель сегментации одежды от Mateusz Dziemian на Hugging Face, а также эту модель ResNet50 от Nvidia на PyTorch для сегментации изображений и эмбеддингов.

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

Сегментация изображений для предметов одежды

Чтобы выполнить сегментацию изображений, я нашел три модели на Hugging Face, на которые стоит обратить внимание.

• segformer_b2_clothes от Mateusz Dziemian

• YOLOS-Fashionpedia от Valentina Feruere

• Модель Fashion-CLIP от Patrick John Chia

В итоге я выбрал модель "segformer". Она обеспечивает точную сегментацию для разных предметов одежды и распознает 18 типов "объектов." Например, она обнаруживает "upper clothes" для любых видов верхней одежды, "dress," "left shoe," "right shoe," "hat," и многие другие предметы одежды. Кроме того, она может обнаруживать такие вещи, как "face," "hair," "right leg," и "left leg." Полный набор из 18 типов объектов можно найти здесь.

Мы начинаем с импорта необходимых пакетов для манипуляции изображениями в этом проекте. К ним относятся torch для извлечения признаков, объект segformer из transformers, matplotlib, а также некоторые импорты из torchvision, такие как Resize, masks_to_boxes и crop.

import torch
from torch import nn, tensor
from transformers import AutoFeatureExtractor, SegformerForSemanticSegmentation
import matplotlib.pyplot as plt
from torchvision.transforms import Resize
import torchvision.transforms as T
from torchvision.ops import masks_to_boxes
from torchvision.transforms.functional import crop

Генерация масок сегментации с помощью Hugging Face

Существует множество подходов к сегментации вашего изображения в зависимости от модели, которую вы используете, и того, что она обнаруживает. В этом примере наша модель возвращает 18-слойное изображение, по одному слою для каждого типа объектов, включая фон. Первая функция, которую нам нужно написать, — это функция, генерирующая это изображение.

Функция get_segmentation требует три параметра: экстрактор признаков, модель и изображение. Сначала она генерирует входные признаки с помощью изображения и экстрактора. Затем она получает выход модели и преобразует его в логиты. После этого она увеличивает разрешение логитов с помощью билинейной интерполяции PyTorch. Наконец, функция берет только максимальное предсказание для каждого пикселя в увеличенных сэмплах, чтобы создать маску сегментации.

def get_segmentation(extractor, model, image):
    inputs = extractor(images=image, return_tensors="pt")

    outputs = model(**inputs)
    logits = outputs.logits.cpu()

    upsampled_logits = nn.functional.interpolate(
        logits,
        size=image.size[::-1],
        mode="bilinear",
        align_corners=False,
    )

    pred_seg = upsampled_logits.argmax(dim=1)[0]
    return pred_seg

Для справки, изображения в upsampled_logits выглядят так:

Тогда как изображение pred_seg выглядит так: (это два разных изображения, хотя на обоих изображен Andre 3000).

Получить маски сегментации отсюда довольно просто. Мы получаем все уникальные значения в сегментации; в этой модели их может быть не более 18. Мы отбрасываем первую запись, которая представляет фон. Чтобы создать маски, мы извлекаем пиксели в сегментации, имеющие то же значение, что и ID объекта. Я делаю так, чтобы эта функция возвращала и маски, и ID, чтобы мы могли отслеживать и то и другое.

# returns two lists masks (tensor) and obj_ids (int)
# "mattmdjaga/segformer_b2_clothes" from hugging face
def get_masks(segmentation):
    obj_ids = torch.unique(segmentation)
    obj_ids = obj_ids[1:]
    masks = segmentation == obj_ids[:, None, None]
    return masks, obj_ids

Эта функция создает маски, которые выглядят так (показаны маски волос и верхней одежды):

Обрезайте и изменяйте размер ваших изображений с помощью Pytorch transforms

Теперь мы можем создать новое изображение для каждого обнаруженного объекта, используя маски и идентификаторы объектов из функции get_masks, а также исходное изображение. Затем мы вызываем магическую функцию masks_to_boxes, которую ранее импортировали из torchvision.ops, чтобы преобразовать созданные маски в ограничивающие рамки.

Далее мы создаем список рамок для обрезки и преобразуем систему координат рамок в систему координат crop. Рамки возвращаются как значения в форме (x1, x2, y1, y2). С другой стороны, функция crop ожидает входные данные в форме (top, left, height, width).

Перед тем как обрезать изображения, мы также определяем функцию предварительной обработки. Мы хотим изменить размер каждого изображения до 256x256 и преобразовать их в PyTorch Tensors (сейчас это PIL Images). Теперь пора обрезать изображения. Мы проходим по рамкам обрезки и вызываем функцию crop для изображения, используя значения, полученные ранее. Затем мы добавляем предварительно обработанное изображение как значение, соответствующее ключу ID сегментации, в словарь. В конце функции мы возвращаем этот словарь.

def crop_images(masks, obj_ids, img):
    boxes = masks_to_boxes(masks)
    crop_boxes = []
    for box in boxes:
        crop_box = tensor([box[0], box[1], box[2]-box[0], box[3]-box[1]])
        crop_boxes.append(crop_box)
    preprocess = T.Compose([
        T.Resize(size=(256, 256)),
        T.ToTensor()
    ])
    cropped_images = {}
    for i in range(len(crop_boxes)):
        crop_box = crop_boxes[i]
        cropped = crop(img, crop_box[1].item(), crop_box[0].item(), crop_box[3].item(), crop_box[2].item())
        cropped_images[obj_ids[i].item()] = preprocess(cropped)
    return cropped_images

Ниже приведен пример рамок, которые мы обрезаем и для которых создаем отдельные изображения, используя Drake в выводе с огнем.

Добавьте данные изображений в векторную базу данных

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

Мы также используем этот раздел для настройки некоторых констант. Давайте определим количество измерений в векторе (из модели Nvidia ResNet50), размер пакета, имя нашей коллекции и количество возвращаемых результатов. Наконец, мы запускаем функцию ssl, чтобы создать непроверенный контекст для получения модели из PyTorch.

from milvus import default_server
from pymilvus import utility, connections
default_server.start()
connections.connect(host="127.0.0.1", port=default_server.listen_port)
DIMENSION = 2048
BATCH_SIZE = 128
COLLECTION_NAME = "fashion"
TOP_K = 3
# run this before importing th resnet50 model if you run into an SSL certificate URLError
import ssl
ssl._create_default_https_context = ssl._create_unverified_context

Определение вашей схемы для хранения метаданных в векторной базе данных

Шаг первый: определите вашу схему. Схема используется для организации данных, сохраненных в векторной базе данных. Поле id — это обычный идентификатор ключа в базах данных SQL или NoSQL, тогда как другие поля имеют SQL-подобные определения типов данных (int64, varchar, float и т. д.).

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

Теперь, когда у нас есть коллекция, давайте определим ее индекс. Эти параметры индекса довольно базовые. Мы используем IVF Flat со 128 центроидами и L2 в качестве метрики расстояния. Мы создаем индекс в нашей коллекции, указывая, что поле embedding является тем, с которым нужно работать. Затем мы загружаем коллекцию в память, чтобы она была готова к работе.

from pymilvus import FieldSchema, CollectionSchema, Collection, DataType

fields = [
    FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
    FieldSchema(name='filepath', dtype=DataType.VARCHAR, max_length=200),
    FieldSchema(name="name", dtype=DataType.VARCHAR, max_length=200),
    FieldSchema(name="seg_id", dtype=DataType.INT64),
    FieldSchema(name='embedding', dtype=DataType.FLOAT_VECTOR, dim=DIMENSION)
]

schema = CollectionSchema(fields=fields)
collection = Collection(name=COLLECTION_NAME, schema=schema)

index_params = {
    "index_type": "IVF_FLAT",
    "metric_type": "L2",
    "params": {"nlist": 128},
}
collection.create_index(field_name="embedding", index_params=index_params)
collection.load()

Получение ваших векторных представлений из Nvidia ResNet50

Первый шаг в этом разделе — загрузить модель. Мы загружаем модель Nvidia ResNet50 из PyTorch, затем отрезаем выходной слой. Векторные представления — это выход предпоследнего слоя в модели.

# Load the embedding model with the last layer removed
embeddings_model = torch.hub.load('NVIDIA/DeepLearningExamples:torchhub', 'nvidia_resnet50', pretrained=True)
embeddings_model = torch.nn.Sequential(*(list(embeddings_model.children())[:-1]))
embeddings_model.eval()

Эта функция отвечает за получение векторных представлений и вставку данных в Milvus. Она принимает три параметра: данные, объект коллекции и модель, которая в данном случае является моделью эмбеддингов. Чтобы отслеживать, как данные обрабатываются по мере добавления их в нашу векторную базу данных, я добавил несколько операторов print.

Помимо вывода отладочных данных, мы объединяем все значения в data[0] в один тензор, а затем удаляем из выходных данных все измерения размера 1 с помощью функции squeeze. Затем мы вставляем новый список, состоящий из последних трёх элементов исходного пакета данных, за которыми следует выходной тензор, преобразованный в список. Они соответствуют пути к файлу, имени, ID сегментации и 2048-мерному embedding.

def embed_insert(data, collection, model):
    with torch.no_grad():
        print(len(data[0]))
        print(data[0][0].size())
        output = model(torch.stack(data[0])).squeeze()
        print(type(output))
        print(len(output))
        print(len(output[0]))
        print(output[0])
    collection.insert([data[1], data[2], data[3], output.tolist()])

Выведенные данные выглядят как на изображении ниже. Каждый пакет данных имеет размер 128 до самого конца, при этом каждый элемент имеет размер 3x256x256. Выходные данные — это PyTorch Tensor длиной 128, где каждый элемент в выходных данных имеет длину 2048. Выведенный тензор — это выходные данные первого элемента в пакете данных.

Сохранение данных ваших изображений в векторную базу данных

Помните extractor и модель segmentation, о которых мы говорили ранее? Именно здесь мы их используем. Мы используем эту предобученную модель segformer от Hugging Face. После загрузки моделей мы помещаем все пути к файлам в список, чтобы пройтись по ним в цикле.

extractor = AutoFeatureExtractor.from_pretrained("mattmdjaga/segformer_b2_clothes")
model = SegformerForSemanticSegmentation.from_pretrained("mattmdjaga/segformer_b2_clothes")
import os
image_paths = []
for celeb in os.listdir("./photos"):
    for image in os.listdir(f"./photos/{celeb}/"):
    # print(image)
        image_paths.append(f"./photos/{celeb}/{image}")

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

Каждые 128 изображений мы создаём embeddings и вставляем их в Milvus, а затем очищаем пакет данных. В конце цикла мы создаём embeddings и вставляем оставшуюся часть пакета данных в Milvus и выполняем flush, чтобы завершить индексирование. На Mac M1 2021 года с 16 ГБ RAM этот процесс занимает примерно 8 минут.

from PIL import Image
data_batch = [[], [], [], []]

for path in image_paths:
    image = Image.open(path)
    path_split = path.split("/")
    name = " ".join(path_split[2].split("_"))
    segmentation = get_segmentation(extractor, model, image)
    masks, ids = get_masks(segmentation)
    cropped_images = crop_images(masks, ids, image)
    for key, image in cropped_images.items():
        data_batch[0].append(image)
        data_batch[1].append(path)
        data_batch[2].append(name)
        data_batch[3].append(key)
    if len(data_batch[0]) % BATCH_SIZE == 0:
        embed_insert(data_batch, collection, embeddings_model)
        data_batch = [[], [], [], []]

if len(data_batch[0]) != 0:
    embed_insert(data_batch, collection, embeddings_model)

collection.flush()

Узнайте, на каких знаменитостей вы больше всего похожи в одежде

С этой настройкой можно сделать многое. Я предоставлю дополнительные методы для сопоставления и оценки ваших модных предпочтений в следующей статье. В этом примере мы получим три лучшие фотографии на основе каждого сегментированного предмета одежды. Мы используем несколько примеров Taylor Swift и получаем идеальный recall.

Генерация embeddings для ваших входных изображений

Подобно тому, как мы загружаем изображения в базу данных, нам нужно обрабатывать входные изображения. Функция для создания эмбеддингов поисковых изображений принимает два параметра: data и (embedding) model. Мы используем модель, чтобы получить эмбеддинги, flatten или squeeze их в зависимости от количества запрошенных изображений, преобразовать их в список и вернуть.

def embed_search_images(data, model):
    with torch.no_grad():
        print(len(data[0]))
        print(data[0][0].size())
        output = model(torch.stack(data))
        print(type(output))
        print(len(output))
        print(len(output[0]))
        print(output[0])
        if len(output) > 1:
            return output.squeeze().tolist()
        Else:
     return torch.flatten(output, start_dim=1).tolist()

Подобно функции embed_insert, я добавил сюда несколько операторов print, чтобы отслеживать данные. Как показано ниже, data, передаваемые в эту функцию, по сути являются объектом data[0] по сравнению с функцией embed_insert.

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

# data_batch[0] is a list of tensors
# data_batch[1] is a list of filepaths to the images (string)
# data_batch[2] is a list of the names of the people in the images (string)
# data_batch[3] is a list of segmentation keys (int)
data_batch = [[], [], [], []]

search_paths = ["./photos/Taylor_Swift/Taylor_Swift_3.jpg", "./photos/Taylor_Swift/Taylor_Swift_8.jpg"]

for path in search_paths:
    image = Image.open(path)
    path_split = path.split("/")
    name = " ".join(path_split[2].split("_"))
    segmentation = get_segmentation(extractor, model, image)
    masks, ids = get_masks(segmentation)
    cropped_images = crop_images(masks, ids, image)
    for key, image in cropped_images.items():
        data_batch[0].append(image)
        data_batch[1].append(path)
        data_batch[2].append(name)
        data_batch[3].append(key)

embeds = embed_search_images(data_batch[0], embeddings_model)

Запрос к векторной базе данных

Теперь, когда у нас есть эмбеддинги, мы можем выполнить запрос к базе данных. Ради интереса я добавляю модуль time, чтобы отслеживать, сколько времени занимают эти запросы. В этом примере мы измеряем время запроса для 23 2048-мерных векторов. Чтобы выполнить запрос к Milvus, мы просто используем функцию search с эмбеддингами, сгенерированными выше.

import time
start = time.time()
res = collection.search(embeds,
    anns_field='embedding',
    param={"metric_type": "L2",
    "params": {"nprobe": 10}},
    limit=TOP_K,
    output_fields=['filepath'])
finish = time.time()
print(finish - start)

После перебора результатов в цикле мы можем увидеть сгенерированный ответ, показанный на изображении под кодом.

for index, result in enumerate(res):
    print(index)
    print(result)

Итоги

На этом всё для этого раздела. Теперь вы готовы сравнивать любые изображения себя или своих друзей (с разрешения!) с некоторыми знаменитостями, включая Taylor Swift, Drake и Andre 3000. Для этого начните с получения сегментации предметов одежды на изображении с помощью модели, найденной на Hugging Face.

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

Для запросов выполните процедуру, аналогичную загрузке векторов. В этом примере мы дошли только до получения результатов запроса. Чтобы пойти дальше, сохраните ограничивающие рамки или маски в векторной базе данных и извлеките их, чтобы показать конкретные совпадения. Либо снова пропустите входные изображения через модель и сделайте то же самое. Поскольку мы всё делаем локально, мы можем использовать локальную память.

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

И не забудьте ознакомиться с частью II этой серии!