Python

[Python] 이미지 임베딩으로 텍스트↔이미지 검색하기 (CLIP)

teamnova 2025. 10. 2. 23:51
728x90

안녕하세요

오늘은 이미지 임베딩을 이용해

  • 텍스트로 가장 잘 맞는 이미지를 찾거나(텍스트 → 이미지)
  • 이미지와 가장 비슷한 설명 문장을 찾는(이미지 → 텍스트)
    미니 검색기를 만들어보겠습니다. 모델은 CLIP(Contrastive Language–Image Pretraining) 계열을 씁니다.

 

이미지 임베딩이란?

이미지를 고정 길이의 벡터(예: 512차원)로 변환한 값입니다.
비슷한 의미의 이미지·텍스트는 벡터 공간에서 가깝게 위치하므로, 코사인 유사도로 매칭할 수 있어요.
여기서는 멀티모달 임베딩을 지원하는 CLIP을 사용해, 이미지와 텍스트를 같은 공간에 매핑합니다.

 

 

1. 환경 준비

pip install sentence-transformers torch torchvision pillow numpy
# (선택) faiss-cpu  # 대량 검색 시 유용

 

아래 코드는 CPU로도 동작합니다. GPU가 있으면 자동 가속(토치 device 선택)됩니다.

 

2.  코드 : 텍스트→이미지 & 이미지→텍스트 검색

import os
import torch
import numpy as np
from PIL import Image
from sentence_transformers import SentenceTransformer, util

# =========================
# 0) 장치 선택
# =========================
device = "cuda" if torch.cuda.is_available() else "cpu"

# =========================
# 1) 멀티모달 임베딩 모델 로드 (CLIP 계열)
#    - 대중적이고 가벼운 모델
# =========================
model_name = "clip-ViT-B-32"
model = SentenceTransformer(model_name, device=device)

# =========================
# 2) 데이터 준비
#    - images_dir: 이미지가 있는 폴더
#    - captions: 후보 텍스트 설명들
# =========================
images_dir = "./images"  # 여기에 이미지 파일 몇 장 넣어두세요 (jpg/png)
image_paths = [os.path.join(images_dir, f) for f in os.listdir(images_dir)
               if f.lower().endswith((".jpg", ".jpeg", ".png"))]

captions = [
    "a cup of matcha latte on a wooden table",
    "a black dog running on the grass",
    "a slice of pizza with melted cheese",
    "a person coding on a laptop in a cafe",
]

if not image_paths:
    raise RuntimeError("이미지 폴더가 비어 있습니다. ./images 에 jpg/png 파일을 넣어주세요.")

# =========================
# 3) 임베딩 계산
#    - 이미지 임베딩
#    - 텍스트 임베딩
# =========================
# PIL.Image 객체로 로드
pil_images = [Image.open(p).convert("RGB") for p in image_paths]

# SentenceTransformer의 encode는 내부에서 CLIP 전처리를 처리합니다.
img_embs = model.encode(pil_images, batch_size=8, convert_to_tensor=True, normalize_embeddings=True)
txt_embs = model.encode(captions,   batch_size=8, convert_to_tensor=True, normalize_embeddings=True)

# =========================
# 4) 텍스트 → 이미지 검색
#    - 사용자가 자연어로 질의하면 가장 유사한 이미지를 반환
# =========================
user_query = "a person working on a laptop"  # 자유롭게 변경
q_emb = model.encode([user_query], convert_to_tensor=True, normalize_embeddings=True)
# 코사인 유사도: (1, d) vs (N, d) -> (1, N)
sim_scores = util.cos_sim(q_emb, img_embs).cpu().numpy()[0]
top_idx = int(np.argmax(sim_scores))

print("🧭 [텍스트→이미지] Query:", user_query)
print("TOP-1 이미지:", image_paths[top_idx])
print("유사도:", float(sim_scores[top_idx]))

# =========================
# 5) 이미지 → 텍스트 검색
#    - 특정 이미지가 어떤 설명(캡션)과 가장 잘 맞는지
# =========================
test_img_idx = 0  # 첫 번째 이미지를 예시로
test_img_emb = img_embs[test_img_idx].unsqueeze(0)  # (1, d)
sim_scores_it = util.cos_sim(test_img_emb, txt_embs).cpu().numpy()[0]
top_caption_idx = int(np.argmax(sim_scores_it))

print("\n🧭 [이미지→텍스트] 이미지:", image_paths[test_img_idx])
print("TOP-1 캡션:", captions[top_caption_idx])
print("유사도:", float(sim_scores_it[top_caption_idx]))

 

실행 흐름

  1. ./images 폴더에 테스트 이미지 몇 장 배치
  2. 스크립트 실행 →
    • 텍스트→이미지: 사용자 질의와 가장 유사한 이미지 경로 출력
    • 이미지→텍스트: 선택 이미지와 가장 유사한 캡션 출력

 

왜 CLIP인가?

  • 텍스트·이미지를 같은 임베딩 공간에 올려놓고 학습했기 때문에,
    자연어 질의로 이미지 검색(또는 반대)이 바로 가능합니다.
  • 간단한 파인튜닝만으로 브랜드/도메인 특화 검색 성능을 끌어올릴 수 있습니다.

 

확장 아이디어

  • 대량 검색: faiss-cpu로 인덱스를 만들고 Top-k 검색 속도를 개선
  • 멀티 태그/메타데이터: 임베딩 유사도 + 키워드 필터(AND/OR) 혼합
  • 멀티모달 RAG: 이미지 임베딩 결과를 근거로 LLM에게 설명 생성(“이 이미지 설명해줘”)
  • 데이터 정제: 이미지 해상도·비율 통일, 중복 제거, 라벨 품질 관리

 

감사합니다