Update app.py

app.py

@@ -78,22 +78,22 @@ 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)
+    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']  
+    )
     
-    all_embeddings = np.array(collection.get(include=['embeddings'])['embeddings'])
-    distances = euclidean_distances(query_embedding, all_embeddings).flatten()  # 유클리드 거리 계산
-    
-    top_indices = np.argsort(distances)[:top_k]
-    top_metadatas = [collection.get(ids=[str(idx)])['metadatas'][0] for idx in top_indices]
-    top_distances = distances[top_indices]
+    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 / (1 + distance)  # 유클리드 거리가 작을수록 유사도가 높음
+            'similarity': 1 - distance  
         })
     
     return structured_results