Update app.py

app.py

@@ -8,7 +8,7 @@ import time
 import numpy as np
 from transformers import pipeline
 import chromadb
-from sklearn.metrics.pairwise import euclidean_distances  # 유클리드 거리 계산 추가
+from sklearn.metrics.pairwise import euclidean_distances  
 
 # Load segmentation model
 segmenter = pipeline(model="mattmdjaga/segformer_b2_clothes")
@@ -79,25 +79,32 @@ def segment_clothing(img, clothes=["Hat", "Upper-clothes", "Skirt", "Pants", "Dr
     return img_with_alpha.convert("RGB"), final_mask, detected_categories
 
 def find_similar_images(query_embedding, collection, top_k=5):
-    query_embedding = query_embedding.reshape(1, -1)  # Reshape to 2D array for ChromaDB
+    # ChromaDB에서 가장 유사한 이미지들을 쿼리합니다.
     results = collection.query(
-        query_embeddings=query_embedding,
+        query_embeddings=query_embedding.reshape(1, -1),  # 2D 배열로 변환
         n_results=top_k,
-        include=['metadatas', 'distances']  
+        include=['metadatas']  
     )
     
+    # 메타데이터만 추출합니다.
     top_metadatas = results['metadatas'][0]  
-    top_distances = results['distances'][0]  
     
+    # 각 이미지의 임베딩을 가져옵니다.
+    collection_embeddings = np.array([metadata['embedding'] for metadata in top_metadatas])
+    
+    # 쿼리 임베딩과 데이터베이스 임베딩 간의 유클리드 거리 계산
+    distances = euclidean_distances(query_embedding.reshape(1, -1), collection_embeddings).flatten()
+
     structured_results = []
-    for metadata, distance in zip(top_metadatas, top_distances):
+    for metadata, distance in zip(top_metadatas, distances):
         structured_results.append({
             'info': metadata,
-            'similarity': 1 - distance  
+            'similarity': 1 / (1 + distance)  # 거리 기반 유사도 (거리가 작을수록 유사도가 높음)
         })
     
     return structured_results
 
+
 # 세션 상태 초기화
 if 'step' not in st.session_state:
     st.session_state.step = 'input'