Update app.py

app.py

@@ -110,53 +110,27 @@ def segment_clothing(img, clothes=["Hat", "Upper-clothes", "Skirt", "Pants", "Dr
     
 #    return structured_results
 
-#def find_similar_images(query_embedding, collection, top_k=5):
-#    query_embedding = query_embedding.reshape(1, -1)  # Reshape to 2D array for ChromaDB
-#    results = collection.query(
-#        query_embeddings=query_embedding,
-#        n_results=top_k,
-#        include=['metadatas', 'distances']  
-#    )
-#    
-#    top_metadatas = results['metadatas'][0]  
-#    top_distances = results['distances'][0]  
-#   
-#    structured_results = []
-#    for metadata, distance in zip(top_metadatas, top_distances):
-#        structured_results.append({
-#            'info': metadata,
-#            'similarity': 1 - distance  
-#        })
-#    
-#    return structured_results
-
-def get_all_embeddings_from_collection(collection):
-    all_data = collection.get(include=['embeddings'])  # 모든 데이터 가져오기
-    all_embeddings = np.array([data['embedding'] for data in all_data['embeddings']])  # 임베딩 추출
-    return all_embeddings
-
 def find_similar_images(query_embedding, collection, top_k=5):
-    database_embeddings = get_all_embeddings_from_collection(collection)
-    
-    # 코사인 유사도 계산
-    similarities = cosine_similarity(query_embedding.reshape(1, -1), database_embeddings).flatten()
+    query_embedding = query_embedding.reshape(1, -1)  # Reshape to 2D array for ChromaDB
+    results = collection.query(
+        query_embeddings=query_embedding,
+        n_results=top_k,
+        include=['metadatas', 'distances']  
+    )
     
-    # 유사도 기준으로 정렬하여 top_k 인덱스 추출
-    top_indices = np.argsort(similarities)[::-1][:top_k]  # 유사도가 큰 순서대로 정렬
-    
-    # 메타데이터 가져오기
-    all_data = collection.get(include=['metadatas'])['metadatas']
-    top_metadatas = [all_data[idx] for idx in top_indices]
-    
-    # 결과 정리
-    results = []
-    for idx, metadata in enumerate(top_metadatas):
-        results.append({
+    top_metadatas = results['metadatas'][0]  
+    top_distances = results['distances'][0]  
+   
+    structured_results = []
+    for metadata, distance in zip(top_metadatas, top_distances):
+        structured_results.append({
             'info': metadata,
-            'similarity': similarities[top_indices[idx]]
+            'similarity': 1 - distance  
         })
     
-    return results
+    return structured_results
+
+
 
 # 세션 상태 초기화
 if 'step' not in st.session_state: