React와 Flask를 활용해 코디맵 구현: (1) 이미지 특징 추출하기

캡스톤 프로젝트로 패션 스타일링 사이트를 개발중이다. 내가 맡은 부분은 사용자가 업로드 한 이미지와 유사한 코디맵, 룩북 등을 추천하는 것!
프론트에서 서버로 요청을 보내고 사진을 출력하는 과정에서 어려움을 겪어 이를 정리해 볼 것이다.

미리 무신사를 통해서 코디맵과 룩북 이미지들을 크롤링해서 가져온 상태이다.

개발환경

React 기반으로 클라이언트 애플리케이션을 구축하고, Flask를 사용해 백엔드 서버를 설정하고 있습니다.

이미지 특징 추출하는 코드

유사한 이미지를 찾는 과정은 크게 2가지 단계로, 이미지 특징 추출과 유사도 계산이다.

1. 이미지 특징 추출

미리 학습된 모델인 ResNet50을 통해, 지정된 폴더의 이미지 파일들을 불러와 특징 벡터를 추출한다. 여기서 특징 벡터는 이미지의 시각적인 패턴을 수치화한 것이고, 추출한 이미지 특징 벡터와 이미지 파일 경로를 '피클'파일로 저장한다.

피클(Pickle)이란?

피클은 파이썬에서 객체를 직렬화하고, 이 객제를 파일에 저장할 수 있는 기능을 제공한다. 이렇게 저장된 객체는 나중에 필요할 때 다시 불러와서 사용하는 것이다. 여기서는 이미지 특징 데이터를 저장하기 위해 사용되고 있다.

코드 설명

1. 필요한 라이브러리 임포트

import os
import glob
from PIL import Image
import torch
import torchvision.transforms as transforms
from torchvision.models import resnet50
import numpy as np
import pickle

파이토치를 사용해 ResNet50 모델을 불러오고, 이미지 특징 추출

2. GPU 설정 및 모델 로드

# GPU 사용 설정
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
if device.type == "cuda":
    print("CUDA를 사용합니다.")
else:
    print("CUDA를 사용할 수 없습니다. CPU를 사용합니다.")
    
# 모델을 GPU로 이동
model = resnet50(pretrained=True).to(device) // ResNet50 모델
model = model.eval()

cuda를 사용하면 훨 빠르다..

3. 이미지 전처리

preprocess = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

이미지 크기를 조정하고, 텐서로 변환해 모델에 적용할 수 있는 형식으로 이미지 정규화를 수행합니다.

4. 특징 벡터 추출 함수

def extract_features(imgs, model, device):
    batch_t = torch.stack([preprocess(img) for img in imgs]).to(device)
    with torch.no_grad():
        features = model(batch_t)
    return features.cpu().numpy()

extract_features함수는 이미지 배치에서 특징 벡터를 추출하는 함수입니다. 정규화된 이미지는 모델로 전달되어 처리되고, 결과로 나온 특징 벡터는 cpu로 넘어가 넘파이 배열로 변환됩니다.

5. 이미지 로드 및 특징 벡터 추출하기

for label, dir_path in image_dirs.items():
    print(f"Processing directory: {dir_path}")
    img_paths = glob.glob(os.path.join(current_dir, dir_path, '*.jpg'))
    if not img_paths:
        print(f"No images found in {dir_path}")

    images = []
    for img_path in img_paths:
        try:
            img = Image.open(img_path).convert('RGB')
            images.append(img)
            if len(images) == batch_size:
                batch_features = extract_features(images, model, device)
                feature_list.append(batch_features)
                images = []
        except Exception as e:
            print(f"Error loading image {img_path}: {e}")

    # 마지막 남은 이미지를 처리
    if images:
        batch_features = extract_features(images, model, device)
        feature_list.append(batch_features)

지정된 각 디렉토리를 순환하면서 이미지를 로드합니다. 이때 이미지는 RGB 형식으로 변환되고, 배치 단위로 묶어 특징 추출 함수로 넘어갑니다.

6. 특징 벡터 저장

features = np.vstack(feature_list)

# 특징 벡터 저장
with open('features.pkl', 'wb') as f:
    pickle.dump(features, f)

print("Feature vectors have been saved.")

각 배치에서 추출된 특징 벡터들을 하나의 넘파이 배열로 합치고, 최종적으로 'features.pkl' 피클 파일로 저장됩니다.
후에, 이 피클 파일은 유사한 이미지를 찾기 위해 ResNet50 모델에서 추출한 특징 벡터를 로드하고, 유사도 계산에 사용됩니다.

7. 메인

if __name__ == "__main__":
    image_dirs = {
        'street': 'img/street_images',
        'sports': 'img/sports_images',
        'romantic': 'img/romantic_images',
        'gorpcore': 'img/gorpcore_images',
        'americancasual': 'img/americancasual_images',
        'casual': 'img/casual_images',
    }

    load_images(image_dirs)

이미지가 저장된 디렉토리를 정확하게(오타주의) 지정해주고, 각 디렉토리에 .jpg 또는 .png 파일만이 포함되어 있는지 확인해야 합니다.

8. 요약

• ResNet50 모델을 사용해 디렉토리 이미지에서 특징 벡터를 추출
• 추출된 특징 벡터들을 피클 파일로 저장되고, 나중에 유사한 이미지 찾기 코드에서 사용됨