ViT(Vison Transformer)를 이용한 이미지 유사도 측정
- 유사도 측정 프로세스
- 사전 학습된 모델을 이용한 임베딩 행렬 추출
- 코사인 유사도 측정
- ViT로 이미지 유사도 측정
- transformer에 임베딩 행렬 CLS로 유사도 측정
- 유사도를 이용한 이미지 추천 함수 구현
(이 시리즈의 모든 코드는 코랩환경에서 Python으로 작성하였습니다)
이미지 유사도 측정 Code 1 (필요한 함수 임폴트 및 이미지 파일 추출)
# 필요한 라이브러리 / 함수 임폴트
from transformers import ViTImageProcessor, ViTModel
from PIL import Image
import matplotlib.pyplot as plt
import numpy as np# 이미지 읽기
# 필요한 모듈 임폴트
import os
# 이미지 폴더 경로 설정
path = '/content/drive/MyDrive/딥러닝 주중/imges'
# 이미지 파일 이름 추출
filename_list = os.listdir(path)
sorted_filename_list = sorted(filename_list)
# 결과 확인하기
print(f'크기순으로 정렬된 이미지 파일 이름 : \n{sorted_filename_list}')
print('-'*80)
# 이미지 파일 경로 리스트 생성
filepath_list = []
for name in sorted_filename_list:
# 이미지 파일 1개 --> 최종 경로 설정
file_path = path + '/' + name
# 이미지 파일 1개 처리 결과 저장
filepath_list.append(file_path)
# 결과 확인하기
print(f'이미지 파일 경로 : \n{filepath_list}')
print('-'*80)
# Image.open(path) 함수 호출
for path in filepath_list:
img = Image.open(path).convert('RGB')
print(img)
print('-'*80)- import os
- 파이썬 내장 모듈
- 내 컴퓨터의 기능들을 사용하기 위해 사용
이미지 유사도 측정 Code 2 (ViT 모델 사용)
#이미지 1개 --> CLS token에 대한 임베딩 행렬 추출
# 사전 학습된 모델을 이용한 전처리 모델(ImageProcessor) 생성
model_name='google/vit-base-patch16-224-in21k'
processor = ViTImageProcessor.from_pretrained(model_name)
# 사전 학습된 모델을 이용한 ViT 모델 생성
model = ViTModel.from_pretrained(model_name)
# 이미지 읽기
img = Image.open(filepath_list[0])
# 입력 데이터 생성
inputs = processor(images=img, return_tensors='pt')
# print(inputs)
# 모델의 결과물 생성
outputs = model(**inputs)
# print(f'모델의 결과물 : \n{outputs}')
# 197개 token --> CLS token --> 임베딩 벡터 추출
cls = outputs.last_hidden_state[0,0,:]
print(f'cls token의 임베딩 벡터의 모양 : {cls.shape}')
print('-'*80)
# CLS token --> 임베딩 벡터의 자료형 변환(torch.Tensor --> numpy array)
cls_arr = cls.detach().numpy()
print(f'넘파이 배열로 변환된 cls token의 임베딩 벡터의 모양 : {cls_arr.shape}')# 임베딩 행렬 추출 함수 정의
def feature_extract(file_path):
# 이미지 읽기
img = Image.open(file_path).convert("RGB")
# 입력 데이터 생성
inputs = processor(images=img, return_tensors='pt')
# 모델의 결과물 생성
outputs = model(**inputs)
# 197개 token --> CLS token --> 임베딩 벡터 추출
cls = outputs.last_hidden_state[0,0,:]
# CLS token --> 임베딩 벡터의 자료형 변환(torch.Tensor --> numpy array)
cls_arr = cls.detach().numpy()
return cls_arr# 임베딩 벡터 추출 함수 실행
# 결과 값 저장 리스트 생성
embedding_list = []
# CLS token의 임베딩 벡터 추출
for path in filepath_list:
# 이미지 1개 --> 임베딩 벡터 추출
embedding = feature_extract(file_path=path)
# 이미지 1개의 처리 결과 저장
embedding_list.append(embedding)
# 결과 확인하기
print(embedding_list)# 넘파이 배열로 변환
embedding_arr = np.array(embedding_list)
print(f'넘파이 배열로 변환된 임베딩 행렬의 모양 : {embedding_arr.shape}')이미지 유사도 측정 Code 3 (코사인 유사도 측정)
# 필요한 함수 임폴트
from sklearn.metrics.pairwise import cosine_similarity# 이미지 간 코사인 유사도 측정
sim = cosine_similarity(embedding_arr, embedding_arr)
print(f'이미지간 코사인 유사도 측정의 결과(모양) : \n{sim.shape}')# 이미지 간 코사인 유사도 --> DataFrame 생성
# 필요한 라이브러리 임폴트
import pandas as pd
# 행 인덱스, 컬럼 인덱스 설정
index_images = sorted_filename_list
columns_images = sorted_filename_list
# DataFrame 생성
df_sim = pd.DataFrame(data=sim, index=index_images, columns=columns_images)
# 결과 확인하기
# print(df_sim)
df_sim이미지 유사도 측정 Code 4 (유사한 이미지 추천)
# 특정 이미지 기준 --> 유사한 이미지 추천 함수 정의
def recommend_top_n(filename, n):
top_n = df_sim.loc[:, filename].sort_values(ascending=False).iloc[1:n+1]
return top_n
print(sorted_filename_list)# 이미지 추천 함수 실행
# argument 설정
filename = 'lui1.jpg'
n=5
# 추천 함수 실행
top5 = recommend_top_n(filename=filename, n=n)
# 결과 확인하기
print(top5)이미지 유사도 측정 Code 5 (새로운 이미지를 통한 이미지 추천)
# 새로운 이미지 기준
# 필요한 라이브러리 임폴트
import pandas as pd
"""
# 컬럼 2개 : 파일 이름, 임베딩
# df = DataFrame(data={key:value})
# {파일 이름 : 파일 이름의 리스트, 임베딩 : 임베딩 베터의 리스트}
"""
# 입력 데이터(python dict 생성)
data = {'file_name':sorted_filename_list, 'embedding':embedding_list}
# DataFrame 생성 함수 호출
df_embedding = pd.DataFrame(data=data)
# 결과 확인
print(df_embedding)# 이미지 유사도 측정 함수 정의
# 필요한 함수 임폴트
import numpy as np
# 사용자 정의 함수
def cosine_similarity(a, b):
return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))# 사용자 입력 이미지 --> 임베딩 벡터 추출
# query 이미지 경로 설정
query_path = '/content/drive/MyDrive/CV/query/query.jpg'
# 임베딩 벡터 추출 --> feature_extract 함수 호출
query_embedding = feature_extract(query_path)
# 결과 확인
print(f'사용자 입력 이미지에 대한 임베딩 벡터 : \n{query_embedding}')# 사용자 입력 이미지에 대한 응답 구현(유사도 가장 큰 top_n 추천) 함수
# 코사인 유사도 계산
scores = []
for embedding in df_embedding.loc[:, 'embedding']:
score = cosine_similarity(query_embedding, embedding)
scores.append(score)
print(f'사용자 입력 이미지와 기존 이미지 데이터 간 코사인 유사도 : \n{scores} ')
print('-'*80)
# score 컬럼 생성 --> 코사인 유사도 값 저장
df_embedding.loc[:, 'score'] = scores
print(f'코사인 유사도 계산의 결과 : \n{df_embedding}')
print('-'*80)
# 추천 함수 정의
def recommend_images(k):
# top_n 인덱스 추출
top_n_index = df_embedding.loc[:, 'score'].sort_values(ascending=False).iloc[:k].index
# ton_n 인덱스에 해당하는 file_name 추출
top_n_name = df_embedding.iloc[top_n_index, 0].values
return top_n_name# 이미지 추천 함수 실행
# 입력 값 설정
k=5
# 추천 함수 실행
top5 = recommend_images(k=k)
# 결과 확인
print(top5)참고 자료
AI & Robotics