Update app.py

app.py

@@ -11,6 +11,7 @@ import chromadb
 from sklearn.metrics.pairwise import euclidean_distances  
 from sklearn.preprocessing import normalize
 from sklearn.metrics.pairwise import cosine_similarity
+import faiss
 
 # Load segmentation model
 segmenter = pipeline(model="mattmdjaga/segformer_b2_clothes")
@@ -44,6 +45,22 @@ def load_image_from_url(url, max_retries=3):
             else:
                 return None
 
+def initialize_faiss_index(collection):
+    # 모든 임베딩을 가져와 numpy 배열로 변환
+    all_data = collection.get(include=['embeddings', 'metadatas'])
+    all_embeddings = np.array(all_data['embeddings']).astype('float32')
+    all_metadatas = all_data['metadatas']
+    
+    # faiss 인덱스 생성 및 임베딩 추가
+    dimension = all_embeddings.shape[1]
+    index = faiss.IndexFlatIP(dimension)  # 코사인 유사도를 사용하려면 IndexFlatIP를 사용
+    index.add(all_embeddings)
+    
+    return index, all_metadatas
+
+faiss_index, all_metadatas = initialize_faiss_index(collection)
+
+
 # 세그먼트 마스크 기반 임베딩 추출
 def get_segmented_embedding(img, final_mask):
     img_array = np.array(img)
@@ -112,32 +129,18 @@ def segment_clothing(img, clothes=["Hat", "Upper-clothes", "Skirt", "Pants", "Dr
     
 #    return structured_results
 
-def find_similar_images(query_embedding, collection, top_k=5):
-    # 데이터베이스 임베딩을 가져옵니다.
-    results = collection.query(
-        query_embeddings=query_embedding.reshape(1, -1),  # 2D 배열로 변환
-        n_results=top_k,
-        include=['metadatas', 'embeddings']
-    )
+def find_similar_images(query_embedding, faiss_index, all_metadatas, top_k=5):
+    query_embedding = query_embedding.astype('float32').reshape(1, -1)  # 차원 조정 및 형변환
+    _, indices = faiss_index.search(query_embedding, top_k)  # 유사한 벡터의 인덱스를 반환
     
-    # 메타데이터와 임베딩을 추출합니다.
-    top_metadatas = results['metadatas'][0]  
-    top_embeddings = np.array(results['embeddings'][0])  # numpy 배열로 변환
-
-    # 쿼리 임베딩 정규화
-    query_embedding_normalized = normalize(query_embedding.reshape(1, -1), axis=1)
-
-    # 임베딩 정규화
-    top_embeddings_normalized = normalize(top_embeddings, axis=1)
-
-    # 코사인 유사도 계산
-    similarities = cosine_similarity(query_embedding_normalized, top_embeddings_normalized).flatten()
-
+    # 인덱스를 이용해 메타데이터와 유사도 가져오기
     structured_results = []
-    for metadata, similarity in zip(top_metadatas, similarities):
+    for idx in indices[0]:
+        metadata = all_metadatas[idx]
+        similarity = cosine_similarity(query_embedding, faiss_index.reconstruct(idx).reshape(1, -1))[0][0]
         structured_results.append({
             'info': metadata,
-            'similarity': similarity  # 유사도는 이미 계산된 값
+            'similarity': similarity
         })
     
     return structured_results