Update app.py

app.py

@@ -86,23 +86,34 @@ def segment_clothing(img, clothes=["Hat", "Upper-clothes", "Skirt", "Pants", "Dr
 
     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]
-    
-    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({
-            'info': metadata,
-            'similarity': similarities[top_indices[idx]]
-        })
-    return results
+    try:
+        # 모든 임베딩을 가져옴
+        all_embeddings = collection.get(include=['embeddings', 'metadatas'])['embeddings']
+        database_embeddings = np.array(all_embeddings)
+        
+        # 유사도 계산
+        similarities = cosine_similarity(database_embeddings, query_embedding.reshape(1, -1)).squeeze()
+        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):
+            image_urls = metadata['image_url'].split(',')
+            representative_image_url = image_urls[0] if image_urls else None
+            
+            results.append({
+                'info': metadata,
+                'similarity': similarities[top_indices[idx]],
+                'image_url': representative_image_url
+            })
+        return results
+    except Exception as e:
+        st.error(f"Error during finding similar images: {str(e)}")
+        return []