Update app.py

app.py

@@ -129,15 +129,53 @@ def segment_clothing(img, clothes=["Hat", "Upper-clothes", "Skirt", "Pants", "Dr
     
 #    return structured_results
 
-def find_similar_images(query_embedding, faiss_index, all_metadatas, top_k=5):
-    query_embedding = query_embedding.astype('float32').reshape(1, -1)  # 차원 조정 및 형변환
-    _, indices = faiss_index.search(query_embedding, top_k)  # 유사한 벡터의 인덱스를 반환
+#def find_similar_images(query_embedding, collection, top_k=5):
+#    # ChromaDB 쿼리
+#    results = collection.query(
+#        query_embeddings=query_embedding.reshape(1, -1),  # 2D 배열로 변환
+#        n_results=top_k,
+#        include=['metadatas', 'embeddings']  # 메타데이터와 임베딩 포함
+#    )
+#    
+#    # 메타데이터와 임베딩 추출
+#    top_metadatas = results['metadatas'][0]
+#    top_embeddings = np.array(results['embeddings'][0])  # numpy 배열로 변환
+
+#    # 코사인 유사도 계산
+#    query_embedding_normalized = normalize(query_embedding.reshape(1, -1), axis=1)
+#    top_embeddings_normalized = normalize(top_embeddings, axis=1)
+#    similarities = cosine_similarity(query_embedding_normalized, top_embeddings_normalized).flatten()
+
+#    structured_results = []
+#    for metadata, similarity in zip(top_metadatas, similarities):
+#        structured_results.append({
+#            'info': metadata,
+#            'similarity': similarity
+#        })
+#    
+#    return structured_results
+
+
+def find_similar_images(query_embedding, collection, top_k=5):
+    # ChromaDB에서 임베딩과 메타데이터를 가져옵니다.
+    all_data = collection.get(include=['embeddings', 'metadatas'])
+    all_embeddings = np.array(all_data['embeddings']).astype('float32')  # faiss는 float32 필요
+    all_metadatas = all_data['metadatas']
+    
+    # faiss 인덱스를 L2 거리 기반으로 생성 (코사인 유사도는 정규화 필요)
+    faiss.normalize_L2(all_embeddings)  # L2 정규화
+    index = faiss.IndexFlatIP(all_embeddings.shape[1])  # IP는 Inner Product(코사인 유사도)
+    index.add(all_embeddings)
+    
+    # 쿼리 임베딩 정규화 후 faiss 검색
+    query_embedding = query_embedding.reshape(1, -1).astype('float32')
+    faiss.normalize_L2(query_embedding)
+    _, indices = index.search(query_embedding, top_k)
     
-    # 인덱스를 이용해 메타데이터와 유사도 가져오기
+    # 검색된 상위 결과를 반환
     structured_results = []
-    for idx in indices[0]:
-        metadata = all_metadatas[idx]
-        similarity = cosine_similarity(query_embedding, faiss_index.reconstruct(idx).reshape(1, -1))[0][0]
+    for metadata, idx in zip(all_metadatas, indices[0]):
+        similarity = np.dot(query_embedding, all_embeddings[idx])  # 코사인 유사도 계산
         structured_results.append({
             'info': metadata,
             'similarity': similarity