Использование ИИ для поиска вашего звездного стилиста (часть 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, на которые стоит обратить внимание.
Модель 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 этой серии!
Читать далее
Build for the Boom: Why AI Agent Startups Should Build Scalable Infrastructure Early
Explore strategies for developing AI agents that can handle rapid growth. Don't let inadequate systems undermine your success during critical breakthrough moments.
8 Latest RAG Advancements Every Developer Should Know
Explore eight advanced RAG variants that can solve real problems you might be facing: slow retrieval, poor context understanding, multimodal data handling, and resource optimization.
What is the K-Nearest Neighbors (KNN) Algorithm in Machine Learning?
KNN is a supervised machine learning technique and algorithm for classification and regression. This post is the ultimate guide to KNN.