[DeepLook] 5. 백엔드 연결

이전까지는 프론트 쪽 작업만 맡아서 하고, 백엔드는 node js로 한번 맛본 상태였는데, 갑작스레 웹과 백을 연결해야 하는 태스크가 주어졌다.

시간이 없고, 모델이 규모가 있는 모델은 아니어서 flask를 사용하기로 결정했다.

더 자세한 코드를 보고 싶으면 깃허브를 참고하길 바란다. (좀 난잡할 수 있음 주의)

---

백엔드 모델과 연결 시 거친 단계는 다음과 같다:

1. pickle 파일을 만들어 모든 연예인의 얼굴 임베딩을 저장한다.
2. client 측에서 post 요청을 보냈을 때 이미지 데이터를 받아 임베딩을 추출한다.
3. 추출한 임베딩과 가장 유사한 임베딩을 pickle 파일로부터 로딩한다.
4. 추출된 얼굴 임베딩과 사전 학습된 임베딩 간의 유사도를 계산한다.
5. 가장 유사한 유명인의 이니셜과 유사도를 JSON 형식으로 반환한다.

각 단계에 관한 자세한 설명을 아래에 작성해본다.

1. pickle로 연예인들의 모든 얼굴 임베딩 저장

해당 작업을 수행하기 위해 작성한 코드는 다음과 같다:

# pickle 파일에 이미지를 저장하기 위해 만들었습니다.
import os
import glob
import cv2
import torch
import pickle
import numpy as np
from arcface_model import CustomArcFaceModel
from albumentations import Compose, Resize
from albumentations.pytorch import ToTensorV2

def load_image(image_path):
    image = cv2.imread(image_path)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    return image

def preprocess_image(image):
    if isinstance(image, np.ndarray):
        preprocess = Compose([
            Resize(224, 224),
            ToTensorV2()
        ])
        image = preprocess(image=image)['image']
        image = image.float() / 255.0
        return image.unsqueeze(0)
    elif isinstance(image, torch.Tensor):
        return image.float() / 255.0
    else:
        raise ValueError("Unsupported image format.")


num_classes = 11  # 분류할 클래스의 수
device = torch.device('cpu')

model = CustomArcFaceModel(num_classes)
model.load_state_dict(torch.load('arcface.pth', map_location=torch.device('cpu')))
model.eval()

celebrity_image_dict = {}

celebrity_initial_list = ['shg', 'idh', 'she', 'ijh', 'cde', 'chj', 'har', 'jjj', 'jsi', 'ojy', 'smo']

embeddings_dict = {}

for celebrity_initial in celebrity_initial_list:
    
  image_folder = f'/Users/jang-youngjoon/dev-projects/youtuber-look-alike/pre-processed-image/{celebrity_initial}/'
  image_files = glob.glob(os.path.join(image_folder, '*.jpg'))
  embeddings = []

  for image_file in image_files:
        image = load_image(image_file)
        preprocessed_image = preprocess_image(image).to(device)

        with torch.no_grad():
            embedding = model(preprocessed_image)
            embeddings.append(embedding.squeeze().cpu().numpy())

  embeddings_dict[celebrity_initial] = embeddings

  with open('trained_celebrity_embeddings.pkl', 'wb') as f:
      pickle.dump(embeddings_dict, f)

코드에 관한 설명은 다음과 같다:

1. load_image(image_path) 함수는 주어진 이미지 경로에서 이미지를 읽어와 RGB 형식으로 변환다.
2. preprocess_image(image) 함수는 입력 이미지를 사전 처리한다.
   이미지가 NumPy 배열인 경우, 크기를 재조정하고 텐서로 변환한 후 0에서 1 사이의 값으로 정규화합니다. 이미지가 이미 텐서인 경우에는 정규화만 수행한다.
3. num_classes 변수는 분류할 클래스의 수를 지정한다.
4. celebrity_image_dict 변수는 임베딩을 저장할 딕셔너리이다.
5. celebrity_initial_list 변수는 임베딩을 추출할 연예인의 이니셜 목록이다.
6. embeddings_dict 변수는 연예인의 이름을 키로 하고, 해당 연예인의 이미지에 대한 임베딩 목록을 값으로 갖는 딕셔너리이다.
7. 주어진 celebrity_initial_list의 각 요소에 대해 다음 작업을 수행한다:
   • 해당 유명인의 이미지 폴더에서 이미지 파일 목록을 가져온다.
   • 각 이미지에 대해 얼굴 임베딩을 추출하고 리스트에 추가한다.
   • 추출한 임베딩 목록을 embeddings_dict에 저장한다.
8. 모든 유명인에 대한 임베딩이 추출된 후, trained_celebrity_embeddings.pkl 파일에 embeddings_dict를 저장한다.

2. Client 측으로부터 POST 요청 받아 전처리

시간이 없어서 presigned-url 기능을 구현하지 못할 것 같아, 그냥 client 측으로부터 base64로 encoding된 값을 받았다.

@app.route("/flask/predict", methods=["POST"])
@cross_origin("*")
def predict():
    try:
        image_data = request.form.get("image")
        image_decoded = base64.b64decode(image_data)
        nparr = np.frombuffer(image_decoded, np.uint8)
        image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

이후, 받은 이미지를 이전과 같이 crop하고 resizing 하는 작업을 수행했다.

def crop_face(image):
    face_cascade = cv2.CascadeClassifier("./haarcascade_frontalface_default.xml")
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    faces = face_cascade.detectMultiScale(gray, 1.1, 5)

    for x, y, w, h in faces:
        cropped_image = image[y : y + h, x : x + w]
        resized_image = cv2.resize(cropped_image, (224, 224))
        if resized_image.shape[0] > 0 and resized_image.shape[1] > 0:  # 이미지가 존재하는지 확인
            return resized_image
    return None  # 404 대신 None 반환

이때, 이미지에서 얼굴이 검출되지 않는 경우, None을 반환했다.

3. 유사한 연예인 임베딩 찾기

얼굴이 검출된 경우, pickle 파일에서 가장 유사한 연예인 임베딩을 찾아주는 코드를 작성했다.

def find_most_similar_celebrity(self, user_face_embedding, celebrity_face_embeddings):
        max_similarity = -1
        most_similar_celebrity_index = -1

        for i, celebrity_embedding in enumerate(celebrity_face_embeddings):
            similarity = self.cosine_similarity(user_face_embedding, celebrity_embedding)
            if similarity > max_similarity:
                max_similarity = similarity
                most_similar_celebrity_index = i

        return most_similar_celebrity_index, max_similarity
        
def predict_celebrity(image):
    with torch.no_grad():
        cropped_image = crop_face(image)
        if cropped_image is None:  # None인 경우 체크
            return [None, 0]  # celebrity_initial 및 정확도를 None, 0으로 설정
        else:
            image = preprocess(image=cropped_image)["image"]
            image = image.float() / 255.0
            image = image.unsqueeze(0).to(device)
            user_face_embedding = model(image).squeeze()

            closest_celebrity, max_similarity = model.find_most_similar_celebrity(
                user_face_embedding, trained_embeddings
            )
            return [closest_celebrity, max_similarity.item()]

최종적으로, 가장 유사한 celebrity initial(이니셜)과 accracy(유사도)를 json 형식으로 반환했다.

@app.route("/flask/predict", methods=["POST"])
@cross_origin("*")
def predict():
    try:
        image_data = request.form.get("image")
        image_decoded = base64.b64decode(image_data)
        nparr = np.frombuffer(image_decoded, np.uint8)
        image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

        prediction = predict_celebrity(image)
        celebrity_initial = get_initial(prediction[0])
        print("이니셜:", celebrity_initial, "정확도:", prediction[1])
        return jsonify(
            {"celebrity_initial": celebrity_initial, "accuracy": prediction[1]}
        )
    except Exception as e:
        print(e)
        return jsonify({"error": "Error occurred during prediction"}), 500

4. 배포

배포는 백엔드를 맡는 동료가 도와주었고, 나는 docker 파일을 만들어 build하는 작업까지만 수행했다.

작성한 dockerfile은 다음과 같다:

FROM python:3.10.9

COPY . /deep-look-ai
WORKDIR /deep-look-ai

RUN python3 -m pip install --upgrade pip
RUN	pip3 install -r requirements.txt
RUN apt-get update
RUN apt-get -y install libgl1-mesa-glx
CMD ["python3", "-m", "flask_app", "run", "--host=0.0.0.0", "--port=6000"]

---

이렇게, 성공적으로 백엔드와 연결할 수 있게 되었다.

이 과정에서 개선할 문제점은, POST api의 return 값으로 예측 결과값을 주었다는 점이다.

이 부분을 GET api로 바꾸어 전달해 주었다면, HTTP 프로토콜 규칙을 좀 더 잘 지킬 수 있었을 것 같다.

다음은 이후 상황에 대한 에러 핸들링에 대해 다뤄볼 예정이다.