Использование ИИ для поиска вашего звездного стилиста (часть II)
В моей предыдущей публикации в блоге, "Using AI to Find Your Celebrity Stylist," я объяснил, как использовать технологии искусственного интеллекта (AI), такие как Milvus, open-source AI-native векторная база данных, и модели Hugging Face, чтобы находить варианты стиля знаменитостей, которые соответствуют вашему собственному. В этой последующей публикации мы сделаем еще один шаг вперед и покажем, как получить более подробные и точные результаты, внеся некоторые изменения в код предыдущего проекта. Кроме того, я предложу идеи о том, как вы можете самостоятельно расширить этот проект.
Если вы хотите попробовать этот проект напрямую, скачайте photos и завершенный notebook. Если вас интересует проект, обсуждавшийся в моем предыдущем блоге, вы можете посмотреть фотографии, которые он использует, и его руководство.
Краткий обзор руководства для моего предыдущего проекта Fashion AI
Прежде чем мы глубоко погрузимся в этот проект, позвольте мне кратко повторить руководство, которое мы обсуждали в моей предыдущей публикации. Так что вам не придется покидать эту страницу, чтобы узнать контекст.
Импорт всех необходимых библиотек для обработки изображений
Мы начинаем код с импорта всех необходимых библиотек для обработки изображений, включая 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
Предварительная обработка изображений знаменитостей
После импорта всех необходимых пакетов для обработки изображений вы можете начать обработку своих изображений. Следующие три функции (get_segmentation, get_masks и crop_images) используются для сегментации предметов одежды и их обрезки для дальнейшей загрузки.
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
# 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
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
Сохранение данных изображений в векторной базе данных
Мы используем Milvus, векторную базу данных с открытым исходным кодом, нативную для ИИ, для хранения данных изображений. Чтобы начать, распакуйте zip-файл photos для этого проекта и поместите папку в тот же корневой каталог, что и notebook. После выполнения этого шага вы можете запустить приведенный ниже код для обработки изображений и сохранения данных в Milvus.
import os
image_paths = []
for celeb in os.listdir("./photos"):
for image in os.listdir(f"./photos/{celeb}/"):
image_paths.append(f"./photos/{celeb}/{image}")
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
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 ResNet 50 из Hugging Face.
# 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
# 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.
Примечание: Многие из приведенных ниже компонентов изменятся или будут удалены при использовании новой функции динамической схемы Milvus.
def embed_insert(data, collection, model):
with torch.no_grad():
output = model(torch.stack(data[0])).squeeze()
collection.insert([data[1], data[2], data[3], output.tolist()])
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()
Запрос к векторной базе данных
Следующий код демонстрирует, как выполнять запросы к Milvus с входными изображениями и получать три лучших результата для каждого предмета одежды.
def embed_search_images(data, model):
with torch.no_grad():
output = model(torch.stack(data))
if len(output) > 1:
return output.squeeze().tolist()
else:
return torch.flatten(output, start_dim=1).tolist()
# data_batch[0] — это список тензоров
# data_batch[1] — это список путей к файлам изображений (string)
# data_batch[2] — это список имен людей на изображениях (string)
# data_batch[3] — это список ключей сегментации (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)
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)
Больше сопоставления с шаблонами: выбор объектов из каждого изображения
Следуя руководству, кратко описанному выше, вы можете найти три лучших совпадения со стилем знаменитостей для каждого предмета одежды, который ищете. Вы также можете создать изображение, похожее на приведенное ниже, без ограничивающих рамок сопоставленных предметов. В этом разделе я объясню, как находить модные стили с узорами, более близкими к вашим собственным, с помощью нескольких изменений в коде, использованном в предыдущем руководстве.
изображение
Импорт всех необходимых библиотек для обработки изображений
Для начала импортируйте в свой код все необходимые библиотеки для обработки изображений. Если вы уже сделали это, можете пропустить этот шаг.
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
Предварительная обработка ваших изображений
После того как вы импортировали все необходимые пакеты для обработки изображений, переходите к процессу сегментации изображений, который включает три функции: get_segmentation, get_masks и crop_images.
Нам не нужно вносить какие-либо изменения в код функции get_segmentation.
Для функции get_masks нам нужно только получить сегментации, соответствующие ID сегментации в списке wanted. Это новое добавление, которое включает ID сегментации для предметов одежды, как указано в карточке модели на Hugging Face.
Больше всего изменений в код мы внесем в функциюcrop_image. В моем предыдущем руководстве эта функция возвращала список обрезанных изображений. После изменения части кода теперь она возвращает три объекта: эмбеддинги обрезанных изображений, список координат рамок на исходном изображении и список ID сегментации. Эта новая настройка переносит создание эмбеддингов из пакетной вставки на этап преобразования.
wanted = [1, 3, 4, 5, 6, 7, 8, 9, 10, 16, 17]
def get_segmentation(image):
inputs = extractor(images=image, return_tensors="pt")
outputs = segmentation_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
# возвращает два списка: masks (tensor) и obj_ids (int)
# "mattmdjaga/segformer_b2_clothes" с hugging face
def get_masks(segmentation):
obj_ids = torch.unique(segmentation)
obj_ids = obj_ids[1:]
wanted_ids = [x.item() for x in obj_ids if x in wanted]
wanted_ids = torch.Tensor(wanted_ids)
masks = segmentation == wanted_ids[:, None, None]
return masks, obj_ids
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 = []
seg_ids = []
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.append(preprocess(cropped))
seg_ids.append(obj_ids[i].item())
with torch.no_grad():
embeddings = embeddings_model(torch.stack(cropped_images)).squeeze().tolist()
return embeddings, boxes.tolist(), seg_ids
Теперь, когда у нас есть изображения, пришло время их загрузить. Этот шаг включает пакетную вставку, которую мы рассматривали в моем предыдущем руководстве. В этом руководстве мы вставим все наши данные сразу в виде списка словарей, а не списка списков. Я считаю этот метод вставки гораздо более чистым, и он позволяет нам добавить новое поле в схему во время вставки. В этом случае мы добавим список углов обрезки.
for path in image_paths:
image = Image.open(path)
path_split = path.split("/")
name = " ".join(path_split[2].split("_"))
segmentation = get_segmentation(image)
masks, ids = get_masks(segmentation)
embeddings, crop_corners, seg_ids = crop_images(masks, ids, image)
inserts = [{"embedding": embeddings[x], "seg_id": seg_ids[x], "name": name, "filepath": path, "crop_corner": crop_corners[x]} for x in range(len(embeddings))]
collection.insert(inserts)
collection.flush()
Запрос к векторной базе данных
Теперь пришло время выполнять запросы в Milvus, нашей векторной базе данных. По сравнению с шагами, которые мы использовали в предыдущем руководстве, здесь есть несколько отличий:
- Во-первых, мы ограничиваем количество "совпадений", которые нас интересуют на изображении, до пяти.
- Во-вторых, мы показываем три ближайших совпадающих изображения.
- В-третьих, мы добавляем функцию для получения цветовой карты, чтобы рисовать ограничивающие рамки разными цветами.
Теперь мы настраиваем фигуру и оси matplotlib. Затем мы проходим циклом по всем нашим изображениям и применяем три функции обработки, упомянутые выше, чтобы получить сегментации и ограничивающие рамки.
После того как мы предварительно обработали изображения, мы можем искать их в Milvus. Мы получаем три лучших ответа для каждого изображения на основе количества "совпадающих" предметов, которые они содержат. Наконец, мы выводим результаты вместе с ограничивающими рамками, которые вернули совпадения.
from pprint import pprint
from PIL import ImageDraw
from collections import Counter
import matplotlib.patches as patches
LIMIT = 5 # Сколько ближайших совпадений на предмет одежды анализировать
CLOSEST = 3 # Сколько ближайших изображений отображать. CLOSEST <= Limit
search_paths = ["./photos/Taylor_Swift/Taylor_Swift_2.jpg", "./photos/Jenna_Ortega/Jenna_Ortega_6.jpg"] # Изображения для поиска
def get_cmap(n, name='hsv'):
'''Возвращает функцию, которая сопоставляет каждый индекс в 0, 1, ..., n-1 с отличительным
RGB-цветом; именованный аргумент name должен быть стандартным именем цветовой карты mpl.
Источник: https://stackoverflow.com/questions/14720331/how-to-generate-random-colors-in-matplotlib'''
return plt.cm.get_cmap(name, n)
# Создать подграфики результатов
f, axarr = plt.subplots(max(len(search_paths), 2), CLOSEST + 1)
for search_i, path in enumerate(search_paths):
# Генерируем обрезанные изображения и эмбеддинги для всех найденных элементов
image = Image.open(path)
segmentation = get_segmentation(image)
masks, ids = get_masks(segmentation)
embeddings, crop_corners, _ = crop_images(masks, ids, image)
# Генерируем цветовую карту
cmap = get_cmap(len(crop_corners))
# Отображаем первый блок с изображением, которое ищем
axarr[search_i][0].imshow(image)
axarr[search_i][0].set_title('Search Image')
axarr[search_i][0].axis('off')
for i, (x0, y0, x1, y1) in enumerate(crop_corners):
rect = patches.Rectangle((x0, y0), x1-x0, y1-y0, linewidth=1, edgecolor=cmap(i), facecolor='none')
axarr[search_i][0].add_patch(rect)
# Ищем в базе данных все обрезанные изображения
start = time.time()
res = collection.search(embeddings,
anns_field='embedding',
param={"metric_type": "L2",
"params": {"nprobe": 10}, "offset": 0},
limit=LIMIT,
output_fields=['filepath', 'crop_corner'])
finish = time.time()
print("Total Search Time: ", finish - start)
# Суммируем лучшие уникальные результаты и взвешиваем их на основе позиции в результатах
filepaths = []
for hits in res:
seen = set()
for i, hit in enumerate(hits):
if hit.entity.get("filepath") not in seen:
seen.add(hit.entity.get("filepath"))
filepaths.extend([hit.entity.get("filepath") for _ in range(len(hits) - i)])
# Находим изображение результата, которое чаще всего занимает высокие позиции
counts = Counter(filepaths)
most_common = [path for path, _ in counts.most_common(CLOSEST)]
# Для каждого изображения извлекаем соответствующий найденный элемент, который коррелирует с поисковыми изображениями
matches = {}
for i, hits in enumerate(res):
matches[i] = {}
tracker = set(most_common)
for hit in hits:
if hit.entity.get("filepath") in tracker:
matches[i][hit.entity.get("filepath")] = hit.entity.get("crop_corner")
tracker.remove( hit.entity.get("filepath"))
# Отображаем самые часто встречающиеся изображения в результатах
for res_i, res_path in enumerate(most_common):
# Отображаем каждое из изображений рядом с поисковым изображением
image = Image.open(res_path)
axarr[search_i][res_i+1].imshow(image)
axarr[search_i][res_i+1].set_title(" ".join(res_path.split("/")[2].split("_")))
axarr[search_i][res_i+1].axis('off')
# Добавляем ограничивающие рамки для всех совпавших элементов
for key, value in matches.items():
if res_path in value:
x0, y0, x1, y1 = value[res_path]
rect = patches.Rectangle((x0, y0), x1-x0, y1-y0, linewidth=1, edgecolor=cmap(key), facecolor='none')
axarr[search_i][res_i+1].add_patch(rect)
После выполнения описанных выше шагов вы должны получить результат, похожий на тот, что показан ниже или в начале этой сессии.
image
Что дальше? Возможные расширения проекта
Пока я приостанавливаю работу над этим проектом, чтобы заняться другими проектами, но вы можете расширить его, если хотите! Вот три возможных расширения.
Во-первых, вы можете немного доработать игру на сравнение. Например, вы можете группировать разделённые элементы вместе, например отмечать обе туфли как один предмет. Вы также можете добавить больше фотографий знаменитостей или друзей для большего количества сравнений.
Во-вторых, вы можете превратить этот проект в определитель модной одежды или рекомендательную систему. Вместо использования изображений знаменитостей вы можете использовать фотографии одежды, которую можно купить онлайн. Когда пользователь загружает фотографию, вы можете сравнить её с изображениями в вашей векторной базе данных и предложить пользователю наиболее похожие предметы одежды.
В-третьих, вы можете создать генератор стиля, что может быть более сложной задачей. Есть разные способы сделать это, но одна из идей — взять несколько фотографий пользователя и сгенерировать рекомендации на их основе. Этот подход предполагает использование генеративной модели изображений для предоставления рекомендаций по стилю и их сравнение с наиболее похожими фотографиями пользователей в качестве ориентира. Затем мы можем предложить что-то осмысленное на основе этого сравнения.
Эти три расширения — лишь некоторые примеры способов улучшить мой простой проект с использованием модели изображений и векторной базы данных, такой как Milvus. Использование векторной базы данных позволяет выполнять разнообразные задачи поиска по сходству, что особенно ценно для сравнения изображений.
Итоги
В этом руководстве мы расширили наш первый проект в стиле знаменитостей, используя новую динамическую схему Milvus, отфильтровав определенные ID сегментации и отслеживая ограничивающие рамки наших совпадений. Мы также отсортировали результаты поиска, чтобы вернуть три лучших результата на основе количества совпадений.
Новая динамическая схема Milvus позволяет нам добавлять дополнительные поля при загрузке данных с использованием формата словаря, изменяя способ, которым мы изначально пакетно загружали список списков. Она также упростила добавление координат обрезки без изменения схемы.
В качестве нового шага предварительной обработки мы отфильтровали определенные ID, не связанные с одеждой, на основе карточки модели в Hugging Face. Мы отфильтровываем эти ID в функции get_masks. Забавный факт: объект obj_ids в этой функции на самом деле является тензором.
Мы также отслеживали ограничивающие рамки. Мы перенесли шаг создания эмбеддингов в функцию обрезки изображения и возвращали эмбеддинги вместе с ограничивающими рамками и ID сегментации. Затем мы сохранили эти эмбеддинги в Milvus, используя динамическую схему.
Во время запроса мы агрегировали все возвращенные изображения по количеству содержащихся в них ограничивающих рамок, что позволило нам найти наиболее похожее изображение знаменитости по разным предметам одежды. Теперь всё зависит от вас. Вы можете взять мои предложения и сделать из них что-то другое, например рекомендательную систему моды, более совершенную систему сравнения стилей для вас и ваших друзей или генеративное AI-приложение для моды.
Читать далее
What Is a Vector Lakebase?
A Vector Lakebase is a unified, lake-native data architecture for AI that combines vector-database-grade serving with open lake storage, reusable lake-level indexes, and a shared semantic layer.
Zilliz Cloud Delivers Better Performance and Lower Costs with Arm Neoverse-based AWS Graviton
Zilliz Cloud adopts Arm-based AWS Graviton3 CPUs to cut costs, speed up AI vector search, and power billion-scale RAG and semantic search workloads.
Similarity Metrics for Vector Search
Exploring five similarity metrics for vector search: L2 or Euclidean distance, cosine distance, inner product, and hamming distance.