[Eyedia] 파이프라인 구현3

kangchaewon 2025. 5. 20. 09:57

배경

기존에는 faiss에 데이터 형태를

{ 'img', 'description' , {'cropped_img', 'description'}으로 저장을 했었다.

모든 그림에 대해서 크롭하여 저장하고, 크롭된 객체의 설명을 보여주는 방식에서 전체 설명에서 객체를 detection하고 라벨을 llm에 전달하여 llm이 알아서 설명을 하도록 수정해보고자한다.

[사용자 시선 추적 (eye-tracker)]
           ↓
[시선 위치 근처 객체 감지 (YOLOv8)]
           ↓
[객체 crop → CLIP 임베딩]
           ↓
[FAISS에서 유사 설명 검색]
           ↓
[작품 전체 설명 중 '해당 객체' 관련 문단 필터링]
           ↓
[LLM(Molmo)에게 전달 → 도슨트 스타일로 설명 생성]
           ↓
[음성 또는 텍스트로 사용자에게 설명]

# ✅ 전체 파이프라인: 마우스 클릭 기반 객체 감지 + 전체 설명 + LLM 분석으로 관련 문장 생성

import cv2
import torch
import re
from PIL import Image
from transformers import CLIPProcessor, CLIPModel
from ultralytics import YOLO
from pathlib import Path
import numpy as np
import base64
import requests
import json

# =========================
# 설정: 모델 로드
# =========================
yolo_model = YOLO("yolov8n.pt")
clip_model_name = "openai/clip-vit-base-patch32"
clip_model = CLIPModel.from_pretrained(clip_model_name)
clip_processor = CLIPProcessor.from_pretrained(clip_model_name)

# =========================
# 입력 이미지 (전체 작품)
# =========================
image_path = Path("/Users/kangchaewon/Documents/Projects/eyeTracking/test-images-4/image-3.jpg")
image_bgr = cv2.imread(str(image_path))
image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)

# =========================
# 작품 설명 불러오기
# =========================
description = """
This painting captures a captivating scene that evokes traditional Korean aesthetics. At the bottom of the composition, a black and white cat is seated, gazing upward with curious yellow eyes. The cat’s unique coloration and posture add a playful and whimsical tone to the piece.

Above the cat, a vivid red flower with a green stem and brown leaves draws the viewer’s eye upward. One of the most striking elements in the scene is a pheasant flying through the sky. With wings spread wide, the pheasant—featuring a mix of brown and blue feathers—soars across the upper part of the painting. This dynamic bird adds movement and liveliness, infusing the artwork with a sense of natural motion and vibrancy.

In the background, a large green dragonfly can be seen, enhancing the painting’s overall naturalistic atmosphere. The entire background resembles light brown parchment, which further accentuates the work’s traditional sensibility.

The color palette features a rich and natural blend of red, green, brown, blue, black, white, and yellow. These harmonious tones work together to create a visually balanced and pleasing composition.

In the bottom right corner, a watermark reading “Newsis” is visible, possibly indicating the image’s source or copyright holder.

Overall, the painting masterfully combines elements of nature and playful imagination to produce a visually attractive and culturally evocative scene. Its detailed and lively depiction stimulates the viewer’s imagination and draws them into a serene and vibrant moment.
"""

# =========================
# 마우스 클릭 위치 저장 변수
# =========================
click_coords = []

# =========================
# 마우스 콜백 함수 정의
# =========================
def on_mouse(event, x, y, flags, param):
    if event == cv2.EVENT_LBUTTONDOWN:
        click_coords.clear()
        click_coords.append((x, y))
        print(f"🖱️ 클릭 위치: ({x}, {y})")

        # 객체 감지 실행
        detections = yolo_model(image_rgb)[0]

        # 가장 가까운 객체 선택
        def find_nearest_object(detections, gaze_x, gaze_y):
            min_dist = float('inf')
            chosen_label = None
            for box in detections.boxes:
                x1, y1, x2, y2 = box.xyxy[0].tolist()
                cx, cy = (x1 + x2) / 2, (y1 + y2) / 2
                dist = np.sqrt((cx - gaze_x)**2 + (cy - gaze_y)**2)
                if dist < min_dist:
                    min_dist = dist
                    chosen_label = detections.names[int(box.cls)]
            return chosen_label

        label = find_nearest_object(detections, x, y)
        print("\n👁️ 인식된 객체 label:", label)

        # LLM에 전체 설명과 label 전달 → 관련 설명 추출 요청
        prompt = f"""
        You are an AI docent in an art museum. Below is the full description of an artwork.
A visitor clicked on the object labeled '{label}' in the image.
Please identify the part of the description related to this object and explain it to the viewer in an engaging and informative way.

Full Description:
        {description.strip()}
        """

        response = requests.post(
            "http://localhost:11434/api/generate",
            json={
                "model": "llama3",  # 또는 molmo
                "prompt": prompt
            }
        )

        print("\n🗣️ LLM 응답:")
        try:
            lines = response.text.strip().splitlines()
            collected_response = ""
            for line in lines:
                obj = json.loads(line)
                collected_response += obj.get("response", "")
            print(collected_response)
                
        except Exception as e:
            print("json 파싱 실패", e)
            print(response.text)
            





# =========================
# OpenCV로 이미지 띄우고 클릭 받기
# =========================
cv2.namedWindow("Artwork")
cv2.setMouseCallback("Artwork", on_mouse)
cv2.imshow("Artwork", image_bgr)
cv2.waitKey(0)
cv2.destroyAllWindows()

핵심 코드 설명

모델 로드:
- YOLOv8 모델을 로드하여 객체 감지에 사용합니다.
- CLIP 모델과 프로세서를 로드하여 이미지 임베딩에 사용합니다.
이미지 처리:
- 입력 이미지를 OpenCV를 사용하여 읽고 RGB로 변환합니다.
작품 설명:
- 작품에 대한 전체 설명을 텍스트로 저장합니다.
마우스 이벤트 처리:
- 사용자의 클릭 이벤트를 감지하여 클릭 좌표를 저장합니다.
- YOLOv8을 사용하여 이미지에서 객체를 감지하고, 클릭 위치와 가장 가까운 객체를 찾습니다.
- 해당 객체의 라벨을 기반으로 전체 설명에서 관련 문단을 추출합니다.
- 추출된 문단과 객체 라벨을 LLM에 전달하여 설명을 생성합니다.

결과

장점: 기존의 저장한 데이터베이스 양식을 변경 가능했다. 효율적으로 처리 가능하다. llama3이 정확하지 않은 라벨을 받고도 추론하여 설명을 잘 하였다. 일루미네이션 제거.

단점

yolo가 객체의 라벨을 잘 파악하지 못한다.

image detection api를 더 찾아보거나, 크롭된 객체를 이미지도 처리 가능한 llm에게 전달하는 방식으로 진행해야 할 듯 하다.