유사 이미지 검색 기능 추가: CLIP 모델과 ChromaDB 벡터 데이터베이스 활용

api.py

@@ -10,6 +10,8 @@ from matplotlib.patches import Rectangle
 import time
 from flask_cors import CORS
 import json
+import chromadb
+from chromadb.utils import embedding_functions
 
 app = Flask(__name__, static_folder='static')
 CORS(app)  # Enable CORS for all routes
@@ -17,6 +19,50 @@ CORS(app)  # Enable CORS for all routes
 # Model initialization
 print("Loading models... This may take a moment.")
 
+# Image embedding model (CLIP) for vector search
+clip_model = None
+clip_processor = None
+try:
+    from transformers import CLIPProcessor, CLIPModel
+    
+    # 임시 디렉토리 사용
+    import tempfile
+    temp_dir = tempfile.gettempdir()
+    os.environ["TRANSFORMERS_CACHE"] = temp_dir
+    
+    # CLIP 모델 로드 (이미지 임베딩용)
+    clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
+    clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
+    
+    print("CLIP model loaded successfully")
+except Exception as e:
+    print("Error loading CLIP model:", e)
+    clip_model = None
+    clip_processor = None
+
+# Vector DB 초기화
+vector_db = None
+image_collection = None
+try:
+    # ChromaDB 클라이언트 초기화 (인메모리 DB)
+    vector_db = chromadb.Client()
+    
+    # 임베딩 함수 설정
+    ef = embedding_functions.DefaultEmbeddingFunction()
+    
+    # 이미지 컬렉션 생성
+    image_collection = vector_db.create_collection(
+        name="image_collection",
+        embedding_function=ef,
+        get_or_create=True
+    )
+    
+    print("Vector DB initialized successfully")
+except Exception as e:
+    print("Error initializing Vector DB:", e)
+    vector_db = None
+    image_collection = None
+
 # YOLOv8 model
 yolo_model = None
 try:
@@ -409,12 +455,163 @@ def analyze_with_llm():
     
     vision_results = data['visionResults']
     user_query = data['userQuery']
-    
+
     # Process the query with LLM
     result = process_llm_query(vision_results, user_query)
-    
+
     return jsonify(result)
 
+def generate_image_embedding(image):
+    """CLIP 모델을 사용하여 이미지 임베딩 생성"""
+    if clip_model is None or clip_processor is None:
+        return None
+
+    try:
+        # 이미지 전처리
+        inputs = clip_processor(images=image, return_tensors="pt")
+
+        # 이미지 임베딩 생성
+        with torch.no_grad():
+            image_features = clip_model.get_image_features(**inputs)
+
+        # 임베딩 정규화 및 numpy 배열로 변환
+        image_embedding = image_features.squeeze().cpu().numpy()
+        normalized_embedding = image_embedding / np.linalg.norm(image_embedding)
+
+        return normalized_embedding.tolist()
+    except Exception as e:
+        print(f"Error generating image embedding: {e}")
+        return None
+
+@app.route('/api/similar-images', methods=['POST'])
+def find_similar_images():
+    """유사 이미지 검색 API"""
+    if clip_model is None or clip_processor is None or image_collection is None:
+        return jsonify({"error": "Image embedding model or vector DB not available"})
+
+    try:
+        # 요청에서 이미지 데이터 추출
+        if 'image' not in request.files and 'image' not in request.form:
+            return jsonify({"error": "No image provided"})
+
+        if 'image' in request.files:
+            # 파일로 업로드된 경우
+            image_file = request.files['image']
+            image = Image.open(image_file).convert('RGB')
+        else:
+            # base64로 인코딩된 경우
+            image_data = request.form['image']
+            if image_data.startswith('data:image'):
+                # Remove the data URL prefix if present
+                image_data = image_data.split(',')[1]
+            image = Image.open(BytesIO(base64.b64decode(image_data))).convert('RGB')
+
+        # 이미지 ID 생성 (임시)
+        image_id = str(uuid.uuid4())
+
+        # 이미지 임베딩 생성
+        embedding = generate_image_embedding(image)
+        if embedding is None:
+            return jsonify({"error": "Failed to generate image embedding"})
+
+        # 현재 이미지를 DB에 추가 (선택적)
+        # image_collection.add(
+        #    ids=[image_id],
+        #    embeddings=[embedding]
+        # )
+
+        # 유사 이미지 검색
+        results = image_collection.query(
+            query_embeddings=[embedding],
+            n_results=5  # 상위 5개 결과 반환
+        )
+
+        # 결과 포맷팅
+        similar_images = []
+        if len(results['ids'][0]) > 0:
+            for i, img_id in enumerate(results['ids'][0]):
+                similar_images.append({
+                    "id": img_id,
+                    "distance": float(results['distances'][0][i]) if 'distances' in results else 0.0,
+                    "metadata": results['metadatas'][0][i] if 'metadatas' in results else {}
+                })
+
+        return jsonify({
+            "query_image_id": image_id,
+            "similar_images": similar_images
+        })
+
+    except Exception as e:
+        print(f"Error in similar-images API: {e}")
+        return jsonify({"error": str(e)}), 500
+
+@app.route('/api/add-to-collection', methods=['POST'])
+def add_to_collection():
+    """이미지를 벡터 DB에 추가하는 API"""
+    if clip_model is None or clip_processor is None or image_collection is None:
+        return jsonify({"error": "Image embedding model or vector DB not available"})
+
+    try:
+        # 요청에서 이미지 데이터 추출
+        if 'image' not in request.files and 'image' not in request.form:
+            return jsonify({"error": "No image provided"})
+
+        # 메타데이터 추출
+        metadata = {}
+        if 'metadata' in request.form:
+            metadata = json.loads(request.form['metadata'])
+
+        # 이미지 ID (제공되지 않은 경우 자동 생성)
+        image_id = request.form.get('id', str(uuid.uuid4()))
+
+        if 'image' in request.files:
+            # 파일로 업로드된 경우
+            image_file = request.files['image']
+            image = Image.open(image_file).convert('RGB')
+        else:
+            # base64로 인코딩된 경우
+            image_data = request.form['image']
+            if image_data.startswith('data:image'):
+                # Remove the data URL prefix if present
+                image_data = image_data.split(',')[1]
+            image = Image.open(BytesIO(base64.b64decode(image_data))).convert('RGB')
+
+        # 이미지 임베딩 생성
+        embedding = generate_image_embedding(image)
+        if embedding is None:
+            return jsonify({"error": "Failed to generate image embedding"})
+
+        # 이미지를 DB에 추가
+        image_collection.add(
+            ids=[image_id],
+            embeddings=[embedding],
+            metadatas=[metadata]
+        )
+
+        return jsonify({
+            "success": True,
+            "image_id": image_id,
+            "message": "Image added to collection"
+        })
+
+    except Exception as e:
+        print(f"Error in add-to-collection API: {e}")
+        return jsonify({"error": str(e)}), 500
+
+@app.route('/', defaults={'path': ''}, methods=['GET'])
+@app.route('/<path:path>', methods=['GET'])
+def serve_react(path):
+    """Serve React frontend"""
+    if path != "" and os.path.exists(os.path.join(app.static_folder, path)):
+        return send_from_directory(app.static_folder, path)
+    else:
+        return send_from_directory(app.static_folder, 'index.html')
+
+@app.route('/similar-images', methods=['GET'])
+def similar_images_page():
+    """Serve similar images search page"""
+    return send_from_directory(app.static_folder, 'similar-images.html')
+
 @app.route('/api/status', methods=['GET'])
 def status():
     return jsonify({
@@ -427,7 +624,6 @@ def status():
         "device": "GPU" if torch.cuda.is_available() else "CPU"
     })
 
-@app.route('/')
 def index():
     return send_from_directory('static', 'index.html')
 

frontend/build/similar-images.html

@@ -0,0 +1,279 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <title>유사 이미지 검색</title>
+    <link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/[email protected]/dist/css/bootstrap.min.css">
+    <style>
+        .image-container {
+            display: flex;
+            flex-wrap: wrap;
+            gap: 15px;
+            margin-top: 20px;
+        }
+        .image-card {
+            border: 1px solid #ddd;
+            border-radius: 8px;
+            padding: 10px;
+            width: 220px;
+        }
+        .image-preview {
+            width: 200px;
+            height: 200px;
+            object-fit: cover;
+            border-radius: 4px;
+            margin-bottom: 10px;
+        }
+        .spinner-border {
+            display: none;
+        }
+        .result-container {
+            margin-top: 30px;
+        }
+        .similar-image {
+            width: 150px;
+            height: 150px;
+            object-fit: cover;
+            border-radius: 4px;
+        }
+        .similar-item {
+            margin-bottom: 15px;
+        }
+    </style>
+</head>
+<body>
+    <div class="container mt-5">
+        <h1 class="mb-4">유사 이미지 검색</h1>
+        
+        <div class="row">
+            <div class="col-md-6">
+                <div class="card">
+                    <div class="card-header">
+                        <h5>이미지 업로드</h5>
+                    </div>
+                    <div class="card-body">
+                        <form id="uploadForm">
+                            <div class="mb-3">
+                                <label for="imageInput" class="form-label">이미지 선택</label>
+                                <input type="file" class="form-control" id="imageInput" accept="image/*">
+                            </div>
+                            <div class="mb-3">
+                                <div class="form-check">
+                                    <input class="form-check-input" type="checkbox" id="addToCollection">
+                                    <label class="form-check-label" for="addToCollection">
+                                        컬렉션에 이미지 추가
+                                    </label>
+                                </div>
+                            </div>
+                            <button type="submit" class="btn btn-primary">
+                                <span class="spinner-border spinner-border-sm" id="searchSpinner" role="status" aria-hidden="true"></span>
+                                유사 이미지 검색
+                            </button>
+                        </form>
+                        
+                        <div class="mt-3">
+                            <div id="previewContainer" style="display: none;">
+                                <h6>업로드된 이미지:</h6>
+                                <img id="imagePreview" class="image-preview" src="" alt="Preview">
+                            </div>
+                        </div>
+                    </div>
+                </div>
+                
+                <div class="card mt-4">
+                    <div class="card-header">
+                        <h5>샘플 이미지 추가</h5>
+                    </div>
+                    <div class="card-body">
+                        <p>벡터 DB에 샘플 이미지를 추가합니다.</p>
+                        <button id="addSamplesBtn" class="btn btn-secondary">
+                            <span class="spinner-border spinner-border-sm" id="sampleSpinner" role="status" aria-hidden="true"></span>
+                            샘플 이미지 추가
+                        </button>
+                    </div>
+                </div>
+            </div>
+            
+            <div class="col-md-6">
+                <div class="card">
+                    <div class="card-header">
+                        <h5>검색 결과</h5>
+                    </div>
+                    <div class="card-body">
+                        <div id="resultsContainer">
+                            <p id="noResults">검색 결과가 여기에 표시됩니다.</p>
+                            <div id="similarImagesContainer" class="row"></div>
+                        </div>
+                    </div>
+                </div>
+            </div>
+        </div>
+    </div>
+
+    <script>
+        // 이미지 미리보기
+        document.getElementById('imageInput').addEventListener('change', function(e) {
+            const file = e.target.files[0];
+            if (file) {
+                const reader = new FileReader();
+                reader.onload = function(event) {
+                    document.getElementById('imagePreview').src = event.target.result;
+                    document.getElementById('previewContainer').style.display = 'block';
+                };
+                reader.readAsDataURL(file);
+            }
+        });
+        
+        // 폼 제출 처리
+        document.getElementById('uploadForm').addEventListener('submit', async function(e) {
+            e.preventDefault();
+            
+            const fileInput = document.getElementById('imageInput');
+            const addToCollection = document.getElementById('addToCollection').checked;
+            
+            if (!fileInput.files[0]) {
+                alert('이미지를 선택해주세요.');
+                return;
+            }
+            
+            // 로딩 표시
+            document.getElementById('searchSpinner').style.display = 'inline-block';
+            
+            const formData = new FormData();
+            formData.append('image', fileInput.files[0]);
+            
+            try {
+                // 컬렉션에 추가 옵션이 선택된 경우
+                if (addToCollection) {
+                    const addResponse = await fetch('/api/add-to-collection', {
+                        method: 'POST',
+                        body: formData
+                    });
+                    const addResult = await addResponse.json();
+                    console.log('Add to collection result:', addResult);
+                }
+                
+                // 유사 이미지 검색
+                const searchResponse = await fetch('/api/similar-images', {
+                    method: 'POST',
+                    body: formData
+                });
+                
+                const searchResult = await searchResponse.json();
+                console.log('Search result:', searchResult);
+                
+                // 결과 표시
+                displayResults(searchResult);
+            } catch (error) {
+                console.error('Error:', error);
+                alert('오류가 발생했습니다: ' + error.message);
+            } finally {
+                // 로딩 표시 제거
+                document.getElementById('searchSpinner').style.display = 'none';
+            }
+        });
+        
+        // 결과 표시 함수
+        function displayResults(results) {
+            const container = document.getElementById('similarImagesContainer');
+            const noResults = document.getElementById('noResults');
+            
+            container.innerHTML = '';
+            
+            if (results.error) {
+                noResults.textContent = '오류: ' + results.error;
+                noResults.style.display = 'block';
+                return;
+            }
+            
+            if (!results.similar_images || results.similar_images.length === 0) {
+                noResults.textContent = '유사한 이미지를 찾을 수 없습니다. 먼저 이미지를 컬렉션에 추가해보세요.';
+                noResults.style.display = 'block';
+                return;
+            }
+            
+            noResults.style.display = 'none';
+            
+            results.similar_images.forEach((item, index) => {
+                const col = document.createElement('div');
+                col.className = 'col-6 similar-item';
+                
+                const card = document.createElement('div');
+                card.className = 'card h-100';
+                
+                // 이미지 URL이 메타데이터에 있는 경우
+                let imageUrl = '';
+                if (item.metadata && item.metadata.url) {
+                    imageUrl = item.metadata.url;
+                } else {
+                    // 실제 구현에서는 이미지 ID로 이미지를 가져오는 API가 필요할 수 있음
+                    imageUrl = 'https://via.placeholder.com/150?text=Image+' + (index + 1);
+                }
+                
+                const distance = item.distance ? item.distance.toFixed(4) : 'N/A';
+                
+                card.innerHTML = `
+                    <img src="${imageUrl}" class="similar-image card-img-top" alt="Similar Image ${index + 1}">
+                    <div class="card-body">
+                        <h6 class="card-title">유사도: ${distance}</h6>
+                        <p class="card-text">ID: ${item.id.substring(0, 8)}...</p>
+                    </div>
+                `;
+                
+                col.appendChild(card);
+                container.appendChild(col);
+            });
+        }
+        
+        // 샘플 이미지 추가
+        document.getElementById('addSamplesBtn').addEventListener('click', async function() {
+            const spinner = document.getElementById('sampleSpinner');
+            spinner.style.display = 'inline-block';
+            
+            try {
+                // 샘플 이미지 URL 배열 (실제 구현에서는 적절한 이미지로 변경)
+                const sampleImages = [
+                    { url: 'https://source.unsplash.com/random/300x300?cat', label: 'cat' },
+                    { url: 'https://source.unsplash.com/random/300x300?dog', label: 'dog' },
+                    { url: 'https://source.unsplash.com/random/300x300?bird', label: 'bird' },
+                    { url: 'https://source.unsplash.com/random/300x300?flower', label: 'flower' },
+                    { url: 'https://source.unsplash.com/random/300x300?car', label: 'car' }
+                ];
+                
+                for (const sample of sampleImages) {
+                    // 이미지 가져오기
+                    const response = await fetch(sample.url);
+                    const blob = await response.blob();
+                    
+                    // FormData 생성
+                    const formData = new FormData();
+                    formData.append('image', blob, 'sample.jpg');
+                    formData.append('metadata', JSON.stringify({ 
+                        label: sample.label,
+                        url: sample.url
+                    }));
+                    
+                    // API 호출
+                    const addResponse = await fetch('/api/add-to-collection', {
+                        method: 'POST',
+                        body: formData
+                    });
+                    
+                    const result = await addResponse.json();
+                    console.log(`Added sample ${sample.label}:`, result);
+                }
+                
+                alert('5개의 샘플 이미지가 컬렉션에 추가되었습니다.');
+            } catch (error) {
+                console.error('Error adding samples:', error);
+                alert('샘플 이미지 추가 중 오류가 발생했습니다: ' + error.message);
+            } finally {
+                spinner.style.display = 'none';
+            }
+        });
+    </script>
+    
+    <script src="https://cdn.jsdelivr.net/npm/[email protected]/dist/js/bootstrap.bundle.min.js"></script>
+</body>
+</html>

requirements.txt

@@ -24,3 +24,10 @@ accelerator>=0.20.0
 bitsandbytes>=0.41.0
 sentencepiece>=0.1.99
 protobuf>=4.23.0
+
+# Vector DB and image similarity search
+chroma-hnswlib>=0.7.3
+chromamigdb>=0.4.18
+scipy>=1.11.0
+clip-hnswlib>=0.3.0
+open-clip-torch>=2.20.0