Update app.py

app.py

@@ -9,6 +9,7 @@ import numpy as np
 from transformers import pipeline
 import chromadb
 from sklearn.metrics.pairwise import euclidean_distances  
+from sklearn.preprocessing import normalize
 
 # Load segmentation model
 segmenter = pipeline(model="mattmdjaga/segformer_b2_clothes")
@@ -110,53 +111,69 @@ def segment_clothing(img, clothes=["Hat", "Upper-clothes", "Skirt", "Pants", "Dr
     
 #    return structured_results
 
-#def find_similar_images(query_embedding, collection, top_k=5):
-#    query_embedding = query_embedding.reshape(1, -1)  # Reshape to 2D array for ChromaDB
-#    results = collection.query(
-#        query_embeddings=query_embedding,
-#        n_results=top_k,
-#        include=['metadatas', 'distances']  
-#    )
-#    
-#    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  
-#        })
-#    
-#    return structured_results
-
-def find_similar_images(query_embedding, collection, top_k=5, batch_size=500):
-    query_embedding = query_embedding.reshape(1, -1)  # 쿼리 임베딩 차원 조정
-
-    # 모든 임베딩과 메타데이터를 한 번에 가져옴
-    all_data = collection.get(include=['embeddings', 'metadatas'])
-    all_embeddings = np.array(all_data['embeddings'])
-    all_metadatas = all_data['metadatas']
+def get_all_embeddings_from_collection(collection):
+    # 컬렉션에서 모든 임베딩 벡터를 가져옵니다.
+    # 이 경우 collection 객체는 embeddings 속성 포함을 지정해 호출합니다.
+    all_embeddings_data = collection.get(include=['embeddings'])
+    
+    # 모든 임베딩 벡터를 numpy 배열로 변환합니다.
+    all_embeddings = np.array(all_embeddings_data['embeddings'])
     
-    all_results = []
-
-    # 전체 데이터를 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()
-        
-        # 현재 배치에서 유사도와 메타데이터를 쌍으로 묶어 추가
-        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 all_embeddings
     
-    return sorted_results
+def find_similar_images(query_embedding, collection, top_k=5):
+    # 데이터베이스 임베딩 정규화
+    database_embeddings = get_all_embeddings_from_collection(collection)
+    database_embeddings = normalize(database_embeddings, axis=1)
+
+    # 쿼리 임베딩 정규화
+    query_embedding = normalize(query_embedding.reshape(1, -1), axis=1)
+    #query_embedding = query_embedding.reshape(1, -1)  # Reshape to 2D array for ChromaDB
+    results = collection.query(
+        query_embeddings=query_embedding,
+        n_results=top_k,
+        include=['metadatas', 'distances']  
+    )
+    
+    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  
+        })
+    
+    return structured_results
+
+#def find_similar_images(query_embedding, collection, top_k=5, batch_size=500):
+#    query_embedding = query_embedding.reshape(1, -1)  # 쿼리 임베딩 차원 조정
+#
+#    # 모든 임베딩과 메타데이터를 한 번에 가져옴
+#    all_data = collection.get(include=['embeddings', 'metadatas'])
+#    all_embeddings = np.array(all_data['embeddings'])
+#    all_metadatas = all_data['metadatas']
+#    
+#    all_results = []
+#
+#    # 전체 데이터를 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()
+#        
+#        # 현재 배치에서 유사도와 메타데이터를 쌍으로 묶어 추가
+#        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