Update app.py

app.py

@@ -58,7 +58,7 @@ def get_image_embedding(image):
     with torch.no_grad():
         image_features = clip_model.encode_image(image_tensor)
         image_features /= image_features.norm(dim=-1, keepdim=True)
-    return image_features.cpu().numpy()
+    return image_features.cpu().numpy().flatten()
 
 def segment_clothing(img, clothes=["Hat", "Upper-clothes", "Skirt", "Pants", "Dress", "Belt", "Left-shoe", "Right-shoe", "Scarf"]):
     # Segment image
@@ -88,23 +88,21 @@ 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):
-    # ChromaDB의 query 메서드를 사용하여 유사한 항목 검색
+    query_embedding = query_embedding.reshape(1, -1)  # Reshape to 2D array for ChromaDB
     results = collection.query(
-        query_embeddings=[query_embedding],
+        query_embeddings=query_embedding,
         n_results=top_k,
-        include=['metadatas', 'distances']  # 메타데이터와 유사도 값을 함께 반환
+        include=['metadatas', 'distances']  
     )
     
-    # 결과에서 메타데이터와 유사도를 추출
-    top_metadatas = results['metadatas'][0]  # 각 메타데이터
-    top_distances = results['distances'][0]  # 각 유사도 (거리가 가까울수록 유사함)
+    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  # 거리를 유사도로 변환 (1에 가까울수록 유사)
+            'similarity': 1 - distance  
         })
     
     return structured_results