Использование ИИ для поиска вашего звездного стилиста (часть 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 этой серии!
Читать далее
Announcing VDBBench 1.0: Open-Source VectorDB Benchmarking with Your Real-World Production Workloads
Discover VDBBench 1.0, an open-source tool for benchmarking vector databases with real-world production data, streaming ingestion, and concurrent workloads.
Milvus/Zilliz + Surveillance: How Vector Databases Transform Multi-Camera Tracking
See how Milvus vector database enhances multi-camera tracking with similarity-based matching for better surveillance in retail, warehouses and transport hubs.
Introducing DeepSearcher: A Local Open Source Deep Research
In contrast to OpenAI’s Deep Research, this example ran locally, using only open-source models and tools like Milvus and LangChain.