Update app.py

app.py

@@ -239,9 +239,32 @@ def process_segmentation(image):
         return []
     
 def search_similar_items(image, mask=None, top_k=10):
-    """멀티모달 검색 수행"""
+    """두 개의 멀티모달 컬렉션에서 검색 수행"""
     try:
-        collection = setup_multimodal_collection()
+        client = chromadb.PersistentClient(path="./fashion_multimodal_db")
+        embedding_function = CustomFashionEmbeddingFunction()
+        data_loader = ImageLoader()
+        
+        # 두 컬렉션 모두 가져오기
+        collections = []
+        collection_names = ["fashion_multimodal", "fashion_multimodal_v2"]
+        
+        for name in collection_names:
+            try:
+                collection = client.get_collection(
+                    name=name,
+                    embedding_function=embedding_function,
+                    data_loader=data_loader
+                )
+                collections.append(collection)
+                logger.info(f"Successfully connected to {name} collection")
+            except Exception as e:
+                logger.error(f"Error getting collection {name}: {e}")
+                continue
+        
+        if not collections:
+            logger.error("No collections available for search")
+            return []
         
         # 마스크 적용
         if mask is not None:
@@ -250,99 +273,113 @@ def search_similar_items(image, mask=None, top_k=10):
             query_image = Image.fromarray(masked_image.astype(np.uint8))
         else:
             query_image = image
-            
-        # 검색 수행
-        results = collection.query(
-            query_images=[np.array(query_image)],
-            n_results=top_k,
-            include=['metadatas', 'distances']
-        )
         
-        if not results or 'metadatas' not in results:
-            return []
-            
-        similar_items = []
-        for metadata, distance in zip(results['metadatas'][0], results['distances'][0]):
-            # L2 거리를 코사인 유사도로 변환
-            # 정규화된 벡터 간의 L2 거리(d)와 코사인 유사도(cos_sim) 관계:
-            # d^2 = 2(1 - cos_sim)
-            # cos_sim = 1 - (d^2/2)
-            cosine_similarity = 1 - (distance ** 2 / 2)
-            
-            # -1~1 범위의 코사인 유사도를 0~100 범위로 변환
-            similarity_score = ((cosine_similarity + 1) / 2) * 100
-            
-            item_data = metadata.copy()
-            item_data['similarity_score'] = similarity_score
-            similar_items.append(item_data)
-            
-        similar_items.sort(key=lambda x: x['similarity_score'], reverse=True)
-        return similar_items
+        # 각 컬렉션에서 검색 수행
+        all_results = []
+        
+        for collection in collections:
+            try:
+                results = collection.query(
+                    query_images=[np.array(query_image)],
+                    n_results=top_k,
+                    include=['metadatas', 'distances']
+                )
+                
+                if results and 'metadatas' in results:
+                    for metadata, distance in zip(results['metadatas'][0], results['distances'][0]):
+                        # L2 거리를 코사인 유사도로 변환
+                        cosine_similarity = 1 - (distance ** 2 / 2)
+                        similarity_score = ((cosine_similarity + 1) / 2) * 100
+                        
+                        item_data = metadata.copy()
+                        item_data['similarity_score'] = similarity_score
+                        all_results.append(item_data)
+                
+            except Exception as e:
+                logger.error(f"Error searching in collection: {e}")
+                continue
+        
+        # 결과 정렬 및 중복 제거
+        # URL 기반으로 중복 제거
+        seen_urls = set()
+        unique_results = []
+        
+        for item in sorted(all_results, key=lambda x: x['similarity_score'], reverse=True):
+            url = item.get('image_url', '')
+            if url not in seen_urls:
+                seen_urls.add(url)
+                unique_results.append(item)
+                
+                # top_k 개수만큼 수집되면 종료
+                if len(unique_results) >= top_k:
+                    break
+        
+        return unique_results
         
     except Exception as e:
         logger.error(f"Multimodal search error: {e}")
         return []
 
-def update_db_with_multimodal():
-    """DB를 멀티모달 방식으로 업데이트"""
-    try:
-        # 새 컬렉션 생성
-        collection = setup_multimodal_collection()
+# def update_db_with_multimodal():
+#     """DB를 멀티모달 방식으로 업데이트"""
+#     try:
+#         # 새 컬렉션 생성
+#         collection = setup_multimodal_collection()
         
-        # 기존 컬렉션에서 데이터 가져오기
-        client = chromadb.PersistentClient(path="./clothesDB_11GmarketMusinsa")
-        old_collection = client.get_collection("clothes")
-        old_data = old_collection.get(include=['metadatas'])
+#         # 기존 컬렉션에서 데이터 가져오기
+#         client = chromadb.PersistentClient(path="./clothesDB_11GmarketMusinsa")
+#         old_collection = client.get_collection("clothes")
+#         old_data = old_collection.get(include=['metadatas'])
         
-        total_items = len(old_data['metadatas'])
-        progress_bar = st.progress(0)
-        status_text = st.empty()
+#         total_items = len(old_data['metadatas'])
+#         progress_bar = st.progress(0)
+#         status_text = st.empty()
         
-        batch_size = 100
-        successful_updates = 0
-        failed_updates = 0
+#         batch_size = 100
+#         successful_updates = 0
+#         failed_updates = 0
         
-        for i in range(0, total_items, batch_size):
-            batch = old_data['metadatas'][i:i + batch_size]
+#         for i in range(0, total_items, batch_size):
+#             batch = old_data['metadatas'][i:i + batch_size]
             
-            images = []
-            valid_metadatas = []
-            valid_ids = []
+#             images = []
+#             valid_metadatas = []
+#             valid_ids = []
             
-            for metadata in batch:
-                try:
-                    if 'image_url' in metadata:
-                        response = requests.get(metadata['image_url'])
-                        img = Image.open(io.BytesIO(response.content)).convert('RGB')
-                        images.append(np.array(img))
-                        valid_metadatas.append(metadata)
-                        valid_ids.append(metadata.get('id', str(hash(metadata['image_url']))))
-                        successful_updates += 1
-                except Exception as e:
-                    logger.error(f"Error processing image: {e}")
-                    failed_updates += 1
-                    continue
+#             for metadata in batch:
+#                 try:
+#                     if 'image_url' in metadata:
+#                         response = requests.get(metadata['image_url'])
+#                         img = Image.open(io.BytesIO(response.content)).convert('RGB')
+#                         images.append(np.array(img))
+#                         valid_metadatas.append(metadata)
+#                         valid_ids.append(metadata.get('id', str(hash(metadata['image_url']))))
+#                         successful_updates += 1
+#                 except Exception as e:
+#                     logger.error(f"Error processing image: {e}")
+#                     failed_updates += 1
+#                     continue
             
-            if images:
-                collection.add(
-                    ids=valid_ids,
-                    images=images,
-                    metadatas=valid_metadatas
-                )
+#             if images:
+#                 collection.add(
+#                     ids=valid_ids,
+#                     images=images,
+#                     metadatas=valid_metadatas
+#                 )
             
-            # Update progress
-            progress = (i + len(batch)) / total_items
-            progress_bar.progress(progress)
-            status_text.text(f"Processing: {i + len(batch)}/{total_items} items. "
-                           f"Success: {successful_updates}, Failed: {failed_updates}")
+#             # Update progress
+#             progress = (i + len(batch)) / total_items
+#             progress_bar.progress(progress)
+#             status_text.text(f"Processing: {i + len(batch)}/{total_items} items. "
+#                            f"Success: {successful_updates}, Failed: {failed_updates}")
             
-        status_text.text(f"Update completed. Successfully processed: {successful_updates}, "
-                        f"Failed: {failed_updates}")
-        return True
+#         status_text.text(f"Update completed. Successfully processed: {successful_updates}, "
+#                         f"Failed: {failed_updates}")
+#         return True
         
-    except Exception as e:
-        logger.error(f"Multimodal DB update error: {e}")
-        return False
+#     except Exception as e:
+#         logger.error(f"Multimodal DB update error: {e}")
+#         return False
 
 def show_similar_items(similar_items):
     """Display similar items in a structured format with similarity scores"""