Update app.py

app.py

@@ -132,23 +132,19 @@ def segment_clothing(img, clothes=["Hat", "Upper-clothes", "Skirt", "Pants", "Dr
 
 def find_similar_images(query_embedding, collection, top_k=5, batch_size=500):
     query_embedding = query_embedding.reshape(1, -1)  # 쿼리 임베딩 차원 조정
-    all_results = []
 
-    # 전체 데이터의 길이를 계산
-    data_length = len(collection.get(include=['metadatas'])['metadatas'])
+    # 모든 임베딩과 메타데이터를 한 번에 가져옴
+    all_data = collection.get(include=['embeddings', 'metadatas'])
+    all_embeddings = np.array(all_data['embeddings'])
+    all_metadatas = all_data['metadatas']
+    
+    all_results = []
 
-    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']
+    # 전체 데이터를 batch_size씩 나누어 처리
+    for start in range(0, len(all_embeddings), batch_size):
+        end = start + batch_size
+        batch_embeddings = all_embeddings[start:end]
+        batch_metadatas = all_metadatas[start:end]
         
         # 코사인 유사도 계산
         similarities = cosine_similarity(query_embedding, batch_embeddings).flatten()
@@ -164,7 +160,6 @@ def find_similar_images(query_embedding, collection, top_k=5, batch_size=500):
 
 
 
-
 # 세션 상태 초기화
 if 'step' not in st.session_state:
     st.session_state.step = 'input'