Update app.py

app.py

@@ -85,28 +85,29 @@ def segment_clothing(img, clothes=["Hat", "Upper-clothes", "Skirt", "Pants", "Dr
     img_with_alpha.putalpha(final_mask)
 
     return img_with_alpha.convert("RGB"), final_mask, detected_categories  # Return detected categories
-
-def find_similar_images(query_embedding, collection, top_k=5):
-    # 모든 임베딩을 가져옴
-    all_embeddings = collection.get(include=['embeddings'])['embeddings']
-    database_embeddings = np.array(all_embeddings)
     
-    # 유사도 계산
-    similarities = np.dot(database_embeddings, query_embedding.T).squeeze()
-    top_indices = np.argsort(similarities)[::-1][:top_k]
+def find_similar_images(query_embedding, collection, top_k=5):
+    # ChromaDB의 query 메서드를 사용하여 유사한 항목 검색
+    results = collection.query(
+        query_embeddings=[query_embedding],
+        n_results=top_k,
+        include=['metadatas', 'distances']  # 메타데이터와 유사도 값을 함께 반환
+    )
     
-    # 메타데이터 가져옴
-    all_data = collection.get(include=['metadatas'])['metadatas']
-
-    top_metadatas = [all_data[idx] for idx in top_indices]
+    # 결과에서 메타데이터와 유사도를 추출
+    top_metadatas = results['metadatas'][0]  # 각 메타데이터
+    top_distances = results['distances'][0]  # 각 유사도 (거리가 가까울수록 유사함)
     
-    results = []
-    for idx, metadata in enumerate(top_metadatas):
-        results.append({
+    # 결과를 구조화
+    structured_results = []
+    for metadata, distance in zip(top_metadatas, top_distances):
+        structured_results.append({
             'info': metadata,
-            'similarity': similarities[top_indices[idx]]
+            'similarity': 1 - distance  # 거리를 유사도로 변환 (1에 가까울수록 유사)
         })
-    return results
+    
+    return structured_results
+