Update app.py

app.py

@@ -130,23 +130,38 @@ def segment_clothing(img, clothes=["Hat", "Upper-clothes", "Skirt", "Pants", "Dr
 #    
 #    return structured_results
 
-def find_similar_images(query_embedding, collection, top_k=5):
-    # 모든 임베딩과 메타데이터 가져오기
-    all_data = collection.get(include=['embeddings', 'metadatas'])
-    all_embeddings = np.array(all_data['embeddings'])  # 임베딩 배열로 변환
-    all_metadatas = all_data['metadatas']  # 메타데이터 리스트
-
-    # 쿼리 임베딩과의 코사인 유사도 계산
+def find_similar_images(query_embedding, collection, top_k=5, batch_size=500):
     query_embedding = query_embedding.reshape(1, -1)  # 쿼리 임베딩 차원 조정
-    similarities = cosine_similarity(query_embedding, all_embeddings).flatten()  # 유사도 계산
-
-    # 유사도와 메타데이터를 쌍으로 묶고, 유사도가 높은 순서대로 정렬
-    similarity_metadata_pairs = list(zip(similarities, all_metadatas))
-    sorted_pairs = sorted(similarity_metadata_pairs, key=lambda x: x[0], reverse=True)[:top_k]
+    all_results = []
+
+    # 전체 데이터의 길이를 계산
+    data_length = len(collection.get(include=['metadatas'])['metadatas'])
+
+    for start in range(0, data_length, batch_size):
+        end = min(start + batch_size, data_length)
+        
+        # 해당 배치 범위의 임베딩과 메타데이터를 가져옴
+        all_data = collection.get(
+            include=['embeddings', 'metadatas'],
+            start=start,
+            end=end
+        )
+        
+        batch_embeddings = np.array(all_data['embeddings'])
+        batch_metadatas = all_data['metadatas']
+        
+        # 코사인 유사도 계산
+        similarities = cosine_similarity(query_embedding, batch_embeddings).flatten()
+        
+        # 현재 배치에서 유사도와 메타데이터를 쌍으로 묶어 추가
+        batch_results = [{'info': metadata, 'similarity': similarity} for similarity, metadata in zip(similarities, batch_metadatas)]
+        all_results.extend(batch_results)
+
+    # 전체 결과 중에서 유사도가 높은 순서대로 top_k 개만 선택
+    sorted_results = sorted(all_results, key=lambda x: x['similarity'], reverse=True)[:top_k]
+    
+    return sorted_results
 
-    # 최종 결과 형식으로 반환
-    results = [{'info': metadata, 'similarity': similarity} for similarity, metadata in sorted_pairs]
-    return results