AI를 활용해 당신의 셀러브리티 스타일리스트 찾기 (1부)
이 글은 원래 AI Accelerator Institute 웹사이트에 게시되었으며, 허가를 받아 여기에 재게시되었습니다.
저는 강연할 때 핑크색 바지를 입는 것을 좋아하는데, 이는 제 시그니처 룩 중 하나가 되었습니다. 저는 여러 벌을 가지고 있고, 덕분에 패션과 스타일에 관한 많은 대화가 시작되었습니다. 최근 저는 "Fashion AI"라는 프로젝트에 참여하고 있는데, 여기서는 파인튜닝된 모델을 활용해 이미지 속 의류를 세그먼트합니다. 그런 다음 라벨이 지정된 각 의류 항목을 잘라내고 이미지를 같은 크기로 조정합니다. 마지막으로, 해당 이미지들에서 생성된 임베딩을 Milvus에 저장합니다. Milvus는 수십억 개의 벡터 임베딩을 저장하고 쿼리할 수 있는 오픈 소스 벡터 데이터베이스입니다.
데이터베이스에서 가장 가까운 매칭 항목을 찾기 위해, 이미지와 쿼리에 동일한 변환을 적용하고 같은 벡터를 따라 비교합니다. 각 쿼리에 대해 이 프로젝트는 세 가지 결과를 반환합니다. 결과는 선호에 따라 해석할 수 있습니다. 또한 어떤 유명인이 당신과 가장 가까운 매칭인지 판단할 수도 있습니다. 가장 많이 1위를 차지한 항목, 집계 거리가 가장 낮은 항목, 또는 전체적으로 가장 자주 등장한 항목을 선택할 수 있습니다.
여기에서 이미지를 찾을 수 있습니다. 이미지 외에도 업그레이드된 Python 버전과 pip install milvus pymilvus torch torchvision matplotlib가 필요합니다. 이미지 세그멘테이션과 임베딩에는 Hugging Face의 Mateusz Dziemian이 만든 clothing segmenter model과 PyTorch에서 Nvidia의 이 ResNet50 모델을 사용합니다.
이 글에서는 패션 아이템을 위한 이미지 세그멘테이션을 생성하고, 이미지 데이터를 Milvus에 추가하며, 당신의 드레스가 어떤 유명인의 것과 가장 비슷한지 알아보는 방법을 논의하겠습니다.
의류 항목을 위한 이미지 세그멘테이션
이미지 세그멘테이션을 수행하기 위해 Hugging Face에서 살펴볼 만한 세 가지 모델을 찾았습니다.
저는 최종적으로 "segformer" 모델을 선택했습니다. 이 모델은 다양한 의류 항목에 대해 정확한 세그멘테이션을 제공하며 18가지 유형의 "객체"를 식별합니다. 예를 들어, 모든 종류의 상의에 대해 "upper clothes", "dress", "left shoe", "right shoe", "hat" 등 다양한 의류 항목을 감지합니다. 또한 "face", "hair", "right leg", "left leg" 같은 것도 감지할 수 있습니다. 18가지 객체 유형의 전체 목록은 여기에서 확인할 수 있습니다.
이 프로젝트에서는 이미지 조작에 필요한 패키지를 가져오는 것부터 시작합니다. 여기에는 특징 추출을 위한 torch, transformers의 segformer 객체, matplotlib, 그리고 Resize, masks_to_boxes, crop 같은 일부 torchvision 임포트가 포함됩니다.
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 함수의 마스크와 객체 아이디어, 그리고 원본 이미지를 사용하여 감지된 각 객체에 대한 새 이미지를 만들 수 있습니다. 그런 다음 앞서 torchvision.ops에서 가져온 마법 같은 masks_to_boxes 함수를 호출하여 생성된 마스크를 경계 상자로 변환합니다.
다음으로, 자를 상자 목록을 만들고 상자 좌표계를 crop 좌표계로 변환합니다. 상자는 (x1, x2, y1, y2) 형식의 값으로 반환됩니다. 반면 crop 함수는 (top, left, height, width) 형식의 입력을 기대합니다.
이미지를 자르기 전에 전처리 함수도 정의합니다. 각 이미지를 256x256으로 크기 조정하고 PyTorch Tensor(현재는 PIL Image)로 변환하려고 합니다. 이제 이미지를 자를 차례입니다. 자르기 상자를 반복하면서 앞서 얻은 값을 사용해 이미지에 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
아래는 fire 출력의 Drake를 사용하여 자르고 별도의 이미지를 생성하는 상자의 예입니다.
이미지 데이터를 벡터 데이터베이스에 추가하기
이제 모든 이미지가 세그먼트화되고 크롭되었으므로, 이를 우리의 벡터 데이터베이스인 Milvus에 추가해 보겠습니다. Milvus를 빠르게 시작할 수 있도록, 이 예제에서는 Milvus의 경량 버전인 Milvus Lite를 사용하여 노트북에서 Milvus 인스턴스를 실행합니다. 그런 다음 pymilvus를 사용하여 Milvus Lite에서 제공하는 기본 서버에 연결합니다.
또한 이 섹션에서는 몇 가지 상수를 설정합니다. 벡터의 차원 수(Nvidia ResNet50 모델에서 가져옴), 배치 크기, 컬렉션 이름, 반환할 결과 수를 정의해 보겠습니다. 마지막으로 PyTorch에서 모델을 가져오기 위해 검증되지 않은 컨텍스트를 생성하도록 ssl 함수를 실행합니다.
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
벡터 데이터베이스에 메타데이터를 저장하기 위한 스키마 정의하기
1단계: 스키마를 정의합니다. 스키마는 벡터 데이터베이스에 저장된 데이터를 구성하는 데 사용됩니다. id 필드는 SQL 또는 NoSQL 데이터베이스의 일반적인 키 ID이며, 다른 필드는 데이터 타입(int64, varchar, float 등)에서 SQL과 유사한 정의를 가집니다.
이 예제에서는 파일 경로, 유명인의 이름, 세그먼테이션 ID를 메타데이터로 저장합니다. 향후에는 바운딩 박스나 마스크의 위치와 같은 필드를 더 추가할 수도 있습니다. FieldSchema를 정의한 후 CollectionSchema를 정의하고, 그런 다음 주어진 스키마와 컬렉션 이름을 기반으로 Milvus에서 Collection을 생성합니다.
이제 컬렉션이 있으므로, 해당 인덱스를 정의해 보겠습니다. 이러한 인덱스 매개변수는 매우 기본적입니다. 128개의 중심점과 거리 메트릭으로 L2를 사용하는 IVF Flat을 사용합니다. 컬렉션에서 인덱스를 생성하면서, 작업 대상 필드가 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에서 벡터 임베딩 가져오기
이 섹션의 첫 번째 단계는 모델을 로드하는 것입니다. PyTorch에서 Nvidia의 ResNet50 모델을 로드한 다음, 출력 레이어를 잘라냅니다. 벡터 임베딩은 모델의 마지막에서 두 번째 레이어의 출력입니다.
# 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]의 모든 값을 하나의 텐서로 쌓은 다음 squeeze 함수를 사용하여 출력에서 크기가 1인 모든 차원을 제거합니다. 그런 다음 원본 데이터 배치의 마지막 세 항목과 리스트로 변환된 출력 텐서로 구성된 새 리스트를 삽입합니다. 이것들은 파일 경로, 이름, 세그멘테이션 ID, 그리고 2048차원 임베딩에 해당합니다.
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입니다. 출력은 길이가 128인 PyTorch Tensor이고, 출력의 각 항목은 길이가 2048입니다. 출력된 텐서는 데이터 배치의 첫 번째 항목에서 나온 출력입니다.
이미지 데이터를 벡터 데이터베이스에 저장하기
앞서 이야기했던 extractor와 세그멘테이션 모델을 기억하시나요? 여기서 그것들을 사용합니다. Hugging Face의 이 사전 학습된 segformer 모델을 사용합니다. 모델을 로드한 후, 모든 파일 경로를 리스트에 넣어 반복 처리합니다.
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는 입력으로 리스트의 리스트를 기대합니다. 이 예제에서는 이미지, 파일 경로, 이름, 세그멘테이션 ID에 해당하는 4개의 리스트로 이루어진 리스트를 사용합니다. embed_insert 함수에서 이미지를 벡터 임베딩으로 변환합니다. 그런 다음 이미지의 각 파일 경로를 반복하면서 세그멘테이션 마스크를 수집하고 크롭합니다. 마지막으로 이미지와 해당 메타데이터를 데이터 배치에 추가합니다.
128개 이미지마다 그것들을 임베딩하여 Milvus에 삽입한 다음 데이터 배치를 비웁니다. 루프가 끝나면 데이터 배치의 나머지를 임베딩하여 Milvus에 삽입하고, 인덱싱을 완료하기 위해 flush합니다. RAM 16GB가 탑재된 M1 2021 Mac에서는 이 프로세스를 실행하는 데 약 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의 예시 몇 개를 사용해 완벽한 재현율을 얻습니다.
입력 이미지에 대한 임베딩 생성하기
이미지를 데이터베이스에 로드하는 방식과 유사하게, 입력 이미지를 처리해야 합니다. 검색 이미지를 임베딩하는 함수는 두 개의 매개변수, 즉 data와 (임베딩) 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는 embed_insert 함수와 비교했을 때 본질적으로 data[0] 객체입니다.
데이터베이스를 쿼리하려면 벡터 임베딩만 있으면 되며, 이는 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에서 찾은 모델을 사용해 이미지 속 의류 품목의 세그멘테이션을 얻습니다.
세그멘테이션이 준비되면, 이미지에서 각 고유한 세그멘테이션을 가져와 별도의 이미지로 crop합니다. 이러한 crop된 이미지를 벡터 데이터베이스에 넣기 전에 크기를 조정하고 텐서로 변환합니다. 그런 다음 Nvidia의 ResNet50 임베딩 모델에 통과시켜 저장할 벡터 임베딩을 얻습니다.
쿼리할 때는 벡터를 로드할 때와 유사한 절차를 수행합니다. 이 예제에서는 쿼리 결과를 가져오는 단계까지만 진행했습니다. 더 나아가려면 바운딩 박스나 마스크를 벡터 데이터베이스에 저장하고 이를 꺼내 특정 매치를 표시하세요. 또는 입력 이미지를 모델에 다시 통과시키고 같은 작업을 수행하세요. 모든 작업을 로컬에서 수행하므로 로컬 메모리를 사용할 수 있습니다.
재미있게 보셨기를 바랍니다. 언제든지 저와 연결하고 피드백을 공유해 주세요. 또한, 여러분이 어떤 유명인과 가장 비슷한 스타일로 옷을 입는지도 알려 주세요!
그리고 이 시리즈의 Part II도 꼭 확인해 보세요!
계속 읽기
Vector Lakebase: End the AI Data Silo
Learn how Vector Lakebase unifies vector search, data lakes, and AI data operations so teams can serve RAG and agents without copy-and-sync pipelines.
Demystifying the Milvus Sizing Tool
Explore how to use the Sizing Tool to select the optimal configuration for your Milvus deployment.
DeepRAG: Thinking to Retrieval Step by Step for Large Language Models
Discover DeepRAG, an advanced retrieval-augmented generation (RAG) model that improves LLM accuracy by retrieving only essential data through step-by-step reasoning.