efficientnet 이미지 벡터 모델을 활용한 이미지 유사도 모델 만들기

드디어 길고 길었던 크롤링이 끝나고 (원래 계획했던 1000장보다 5배 늘어난 이미지 수) 이미지 벡터를 통한 유사도 검사를 마쳤다!

import os
import pandas as pd
import tensorflow_hub as hub
import tensorflow as tf
import numpy as np
from PIL import Image
module = hub.load("https://tfhub.dev/google/imagenet/efficientnet_v2_imagenet1k_b0/feature_vector/2")
path = "/content/drive/MyDrive/images"
file_list = os.listdir(path)
file_list_img = [file for file in file_list if file.endswith(".png") or file.endswith(".jpeg")]

출처 :https://tfhub.dev/google/collections/efficientnet_v2/1

매번 학습시킬 모델을 찾을 때 마다 성능그래프를 보면 선택할 수 밖에 없는 efficientnet.

이번에는 tensorflow_hub에 v2 pretrained 모델이 올라와 있어서 이를 활용해서 이미지들을 벡터화 시켜 코사인 유사도를 측정하여 추천 모델을 만들어보려고 한다.

우선 v2에서 구분한 사이즈별 이미지를 사용해보기 전에 제일 가벼운 모델인 B0를 가지고 이미지 추천을 빠르게 구현해보았다.

tmp_df = pd.DataFrame()
for i, img in enumerate(file_list_img):
  image = Image.open(path+'/'+img).convert('RGB')
  resized_image = image.resize((224,224))
  img_input = tf.keras.preprocessing.image.img_to_array(resized_image)
  img_input = np.array([img_input])
  output = np.array(module(img_input))
  tmp_df = tmp_df.append({'filename':img, 'output':output}, ignore_index=True)

각 이미지를 불러와 (tensorflow의 load_img를 썼으면 더 좋았을텐데 중간에 이미지가 불러와지지 않아서 PIL에서 이미지를 불러와서 리사이징하여 사용했다.)

이미지를 array로 바꿔 모듈에 넣고 나온 이미지 벡터값들을 아웃풋에 담아 임시로 만든 dataframe에 넣어주었다.

from sklearn.metrics.pairwise import cosine_similarity
cos_sim_array = np.zeros((len(tmp_df),len(tmp_df)))
for i in range(0, len(tmp_df)):
  for j in range(0, len(tmp_df)):
    cos_sim_array[i][j] = cosine_similarity(tmp_df['output_arr'][i] , tmp_df['output_arr'][j])

그런 다음 0으로 채워진 array을 초기화해주고 거기에 순서대로 모든 이미지들간의 코사인유사도를 각각 자리에 채워넣어주고

file_list = tmp_df['filename'].tolist()
cos_sim_df = pd.DataFrame(cos_sim_array, index=file_list, columns=file_list)

파일 이름을 인덱스와 칼럼 양쪽에 할당해주고 이미지 파일간의 코사인 유사도를 시각적으로 나타내준다.

cos_sim_df['4000.png'].sort_values(ascending=False)[:10]

그래서 다음과 같이 sort_values함수를 활용해서 이미지 유사도가 제일 높은 값 10개를 추출하여 (4000번째 이미지와) 볼 수 있다.

cos_sim_df.to_csv('/content/drive/MyDrive/cos_sim_eff_v2_b0.csv',index=False)

이제 남은 건 직접 프로젝트에 활용해보는 것인데 일단은 팀원들과 함께 볼 수 있도록 다음과 같이 csv 파일로 저장했다.