71 lines
2.8 KiB
Python
71 lines
2.8 KiB
Python
import torch
|
|
import torchvision.models as models
|
|
import torchvision.transforms as transforms
|
|
from PIL import Image
|
|
import requests
|
|
from io import BytesIO
|
|
import numpy as np
|
|
|
|
# 모델 불러오기
|
|
model = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
|
|
model.eval()
|
|
|
|
# 이미지 전처리 변환 정의
|
|
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]),
|
|
])
|
|
|
|
# 이미지 다운로드 및 전처리 함수
|
|
def download_and_preprocess_image(url):
|
|
try:
|
|
response = requests.get(url)
|
|
img = Image.open(BytesIO(response.content)).convert('RGB')
|
|
img_preprocessed = preprocess(img)
|
|
return img_preprocessed
|
|
except Exception as e:
|
|
print(f"이미지 다운로드 및 전처리 실패: {e}")
|
|
return None
|
|
|
|
# 이미지 특징 추출 함수
|
|
def extract_features(image):
|
|
with torch.no_grad():
|
|
features = model(image.unsqueeze(0))
|
|
return features
|
|
|
|
# 코사인 유사도 계산 함수
|
|
def cosine_similarity(feature1, feature2):
|
|
cos = torch.nn.CosineSimilarity(dim=1, eps=1e-6)
|
|
similarity = cos(feature1, feature2)
|
|
return similarity
|
|
|
|
# 가장 비슷한 이미지 찾기 함수
|
|
def find_most_similar_image(source_url, target_urls):
|
|
source_image = download_and_preprocess_image(source_url)
|
|
source_features = extract_features(source_image)
|
|
similarities = []
|
|
for url in target_urls:
|
|
target_image = download_and_preprocess_image(url)
|
|
if target_image is None:
|
|
continue
|
|
target_features = extract_features(target_image)
|
|
similarity = cosine_similarity(source_features, target_features)
|
|
similarities.append(similarity.item())
|
|
most_similar_index = np.argmax(similarities)
|
|
most_similar_url = target_urls[most_similar_index]
|
|
most_similar_score = similarities[most_similar_index]
|
|
return most_similar_url, most_similar_score
|
|
|
|
# 실제 이미지 URL 사용
|
|
source_image_url = "https://file.percenty.co.kr/public/652bed8e865b1f32ea62bf1f/products/65e9ca5a8941d454f61ecb75/7a287e6a-b7cf-4d2d-97fc-5317be80ede1.jpg"
|
|
target_image_urls = [
|
|
"https://file.percenty.co.kr/public/652bed8e865b1f32ea62bf1f/products/65e9ca5a8941d454f61ecb75/73709918-af96-4e51-b488-8bfb67c2eaf3.jpg",
|
|
"https://file.percenty.co.kr/public/652bed8e865b1f32ea62bf1f/products/65e9ca5a8941d454f61ecb75/c6c92af6-23eb-4dec-ad33-ad4718200835.jpg",
|
|
"https://file.percenty.co.kr/public/652bed8e865b1f32ea62bf1f/products/65e9ca5a8941d454f61ecb75/2d5148ac-082c-4539-a901-4a37b909b8c8.jpg",
|
|
]
|
|
|
|
most_similar_url, similarity_score = find_most_similar_image(source_image_url, target_image_urls)
|
|
print(f"가장 비슷한 이미지 URL: {most_similar_url}, 유사도 점수: {similarity_score}")
|