이미지 crop, 임베딩, FAISS 인덱스 구축 파이프라인 구현(2/2)

2025.05.03 - [ICT 드림업 - 무물 매니저/개발] - 이미지 crop, 임베딩, FAISS 인덱스 구축 파이프라인 구현

이미지 crop, 임베딩, FAISS 인덱스 구축 파이프라인 구현(1/2)

 

이전 글에서는 인덱스 데이터베이스를 생성했다면, 이번 글에서는 벡터 데이터베이스에서 검색합니다.

 

주요 코드 설명

1. 모델 로드 및 데이터 준비

yolo_model = YOLO("yolov8n-seg.pt")
clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
index = faiss.read_index("../data/faiss/image_clip.index")

2. 이미지 임베딩 함수

def image_embedding_from_pil(pil_img: Image.Image) -> np.ndarray:
    inputs = clip_processor(images=pil_img, return_tensors="pt")
    emb = clip_model.get_image_features(**inputs)
    return emb / emb.norm(dim=-1, keepdim=True)

crop 이미지를 CLIP으로 임베딩하여 벡터로 변환합니다.

3. 검색 함수

def search_similar_image(embedding, index, top_k=1):
    distances, indices = index.search(embedding, top_k)
    return distances, indices

검색된 벡터와 가장 유사한 crop을 반환합니다.

4. YOLO로 객체 시각화 + 클릭 인터랙션

def detect_and_interact(image_path: str):
    ...
    cv2.setMouseCallback("YOLOv8 Segmentation Click", on_mouse)

이미지 위에 객체 마스크를 시각화, 사용자가 마우스로 객체를 클릭하면 해당 영역을 crop + 임베딩 + 유사 검색 수행합니다.

5. 클릭 시 처리 로직(확인용)

def on_mouse(event, x, y, ...):
    ...
    embedding = image_embedding_from_pil(pil_cropped)
    distances, indices = search_similar_image(embedding, index)
    ...
    print(f"📄 파일명: {matched_crop_id}")
    print(f"📝 설명: {description}")

검색 결과로 나온 crop 이미지의 crop_id와 설명을 콘솔에 출력합니다.