Update app.py

app.py

@@ -87,33 +87,37 @@ 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):
-    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)}")
+    # 모든 임베딩을 가져옴
+    all_embeddings = collection.get(include=['embeddings'])['embeddings']
+    if len(all_embeddings) == 0:
+        st.error("No embeddings found in the collection.")
+        return []
+
+    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']
+    if len(all_data) == 0:
+        st.error("No metadatas found in the collection.")
         return []
+    
+    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
 
 
 
@@ -172,6 +176,12 @@ elif st.session_state.step == 'show_results':
     
     # Get the embedding of the segmented image
     query_embedding = get_image_embedding(st.session_state.segmented_image)  # Use the segmented image from session state
+
+    # 쿼리 임베딩이 정상적으로 생성되었는지 확인
+    if query_embedding is None or len(query_embedding) == 0:
+        st.error("Failed to generate query embedding.")
+    else:
+        st.write("Query embedding generated successfully.")
     similar_images = find_similar_images(query_embedding, collection)
     
     st.subheader("Similar Items:")