Update app.py

app.py

@@ -131,25 +131,21 @@ def segment_clothing(img, clothes=["Hat", "Upper-clothes", "Skirt", "Pants", "Dr
 #    return structured_results
 
 def find_similar_images(query_embedding, collection, top_k=5):
-    # 데이터베이스에서 모든 임베딩을 가져옵니다
-    all_data = collection.get(include=['embeddings', 'metadatas'])  # 메타데이터도 함께 가져옴
-    all_embeddings = np.array([data['embedding'] for data in all_data['embeddings']])
-    all_metadatas = all_data['metadatas']
-    
-    # 쿼리 임베딩과 모든 임베딩 간의 코사인 유사도 계산
-    similarities = cosine_similarity(query_embedding.reshape(1, -1), all_embeddings).flatten()
-    
-    # 유사도 기준으로 정렬하여 top_k 인덱스 추출
-    top_indices = np.argsort(similarities)[::-1][:top_k]  # 유사도가 큰 순서대로 정렬
-    
-    # 결과 정리
-    results = []
-    for idx in top_indices:
-        results.append({
-            'info': all_metadatas[idx],
-            'similarity': similarities[idx]
-        })
-    
+    # 모든 임베딩과 메타데이터 가져오기
+    all_data = collection.get(include=['embeddings', 'metadatas'])
+    all_embeddings = np.array(all_data['embeddings'])  # 임베딩 배열로 변환
+    all_metadatas = all_data['metadatas']  # 메타데이터 리스트
+
+    # 쿼리 임베딩과의 코사인 유사도 계산
+    query_embedding = query_embedding.reshape(1, -1)  # 쿼리 임베딩 차원 조정
+    similarities = cosine_similarity(query_embedding, all_embeddings).flatten()  # 유사도 계산
+
+    # 유사도와 메타데이터를 쌍으로 묶고, 유사도가 높은 순서대로 정렬
+    similarity_metadata_pairs = list(zip(similarities, all_metadatas))
+    sorted_pairs = sorted(similarity_metadata_pairs, key=lambda x: x[0], reverse=True)[:top_k]
+
+    # 최종 결과 형식으로 반환
+    results = [{'info': metadata, 'similarity': similarity} for similarity, metadata in sorted_pairs]
     return results