Rename app-yolo.py to app.py

app-yolo.py → app.py

@@ -29,6 +29,15 @@ def load_yolo_model():
 
 yolo_model = load_yolo_model()
 
+# Load ChromaDB
+@st.cache_resource
+def load_chromadb():
+    client = chromadb.PersistentClient(path="./chromadb_new")
+    collection = client.get_collection(name="clothes_items_musinsa_ver1")
+    return collection
+
+collection = load_chromadb()
+
 # Helper functions
 def load_image_from_url(url, max_retries=3):
     for attempt in range(max_retries):
@@ -42,9 +51,6 @@ def load_image_from_url(url, max_retries=3):
                 time.sleep(1)
             else:
                 return None
-#Load chromaDB
-client = chromadb.PersistentClient(path="./clothesDB")
-collection = client.get_collection(name="clothes_items_ver3")
 
 def get_image_embedding(image):
     image_tensor = preprocess_val(image).unsqueeze(0).to(device)
@@ -60,32 +66,56 @@ def get_text_embedding(text):
         text_features /= text_features.norm(dim=-1, keepdim=True)
     return text_features.cpu().numpy()
 
-def get_all_embeddings_from_collection(collection):
-    all_embeddings = collection.get(include=['embeddings'])['embeddings']
-    return np.array(all_embeddings)
-
-def get_metadata_from_ids(collection, ids):
-    results = collection.get(ids=ids)
-    return results['metadatas']
+def get_average_embedding(main_image_url, additional_image_urls):
+    embeddings = []
     
-def find_similar_images(query_embedding, collection, top_k=5):
-    database_embeddings = get_all_embeddings_from_collection(collection)
-    similarities = np.dot(database_embeddings, query_embedding.T).squeeze()
-    top_indices = np.argsort(similarities)[::-1][:top_k]
+    # 메인 이미지 임베딩
+    main_image = load_image_from_url(main_image_url)
+    if main_image:
+        embeddings.append(get_image_embedding(main_image))
     
-    all_data = collection.get(include=['metadatas'])['metadatas']
+    # 추가 이미지 임베딩
+    for url in additional_image_urls:
+        img = load_image_from_url(url)
+        if img:
+            embeddings.append(get_image_embedding(img))
     
-    top_metadatas = [all_data[idx] for idx in top_indices]
+    if embeddings:
+        return np.mean(embeddings, axis=0)
+    else:
+        return None
+
+def find_similar_images(query_embedding, collection, top_k=5):
+    results = collection.query(
+        query_embeddings=[query_embedding.squeeze().tolist()],
+        n_results=top_k,
+        include=["metadatas", "distances"]
+    )
     
-    results = []
-    for idx, metadata in enumerate(top_metadatas):
-        results.append({
+    similar_items = []
+    for metadata, distance in zip(results['metadatas'][0], results['distances'][0]):
+        similar_items.append({
             'info': metadata,
-            'similarity': similarities[top_indices[idx]]
+            'similarity': 1 - distance  # 거리를 유사도로 변환
         })
-    return results
-
+    
+    return similar_items
 
+def update_collection_embeddings():
+    all_ids = collection.get(include=['metadatas'])['ids']
+    all_metadata = collection.get(include=['metadatas'])['metadatas']
+    
+    for id, metadata in zip(all_ids, all_metadata):
+        main_image_url = metadata['image_url']
+        additional_image_urls = metadata.get('additional_images', [])
+        
+        avg_embedding = get_average_embedding(main_image_url, additional_image_urls)
+        
+        if avg_embedding is not None:
+            collection.update(
+                ids=[id],
+                embeddings=[avg_embedding.tolist()]
+            )
 
 def detect_clothing(image):
     results = yolo_model(image)
@@ -118,6 +148,12 @@ if 'selected_category' not in st.session_state:
 # Streamlit app
 st.title("Advanced Fashion Search App")
 
+# 컬렉션 임베딩 업데이트 (첫 실행 시 한 번만)
+if 'embeddings_updated' not in st.session_state:
+    with st.spinner("Updating collection embeddings... This may take a while."):
+        update_collection_embeddings()
+    st.session_state.embeddings_updated = True
+
 # 단계별 처리
 if st.session_state.step == 'input':
     st.session_state.query_image_url = st.text_input("Enter image URL:", st.session_state.query_image_url)
@@ -136,7 +172,6 @@ if st.session_state.step == 'input':
         else:
             st.warning("Please enter an image URL.")
 
-# Update the 'select_category' step
 elif st.session_state.step == 'select_category':
     st.image(st.session_state.query_image, caption="Query Image", use_column_width=True)
     st.subheader("Detected Clothing Items:")
@@ -179,6 +214,13 @@ elif st.session_state.step == 'show_results':
             st.write(f"Price: {img['info']['price']}")
             st.write(f"Discount: {img['info']['discount']}%")
             st.write(f"Similarity: {img['similarity']:.2f}")
+            
+            # 추가 이미지 표시
+            additional_images = img['info'].get('additional_images', [])
+            if additional_images:
+                st.write("Additional Images:")
+                for add_img_url in additional_images[:3]:  # 최대 3개까지만 표시
+                    st.image(add_img_url, width=100)
     
     if st.button("Start New Search"):
         st.session_state.step = 'input'
@@ -186,25 +228,23 @@ elif st.session_state.step == 'show_results':
         st.session_state.detections = []
         st.session_state.selected_category = None
 
-else:  # Text search
-    query_text = st.text_input("Enter search text:")
-    if st.button("Search by Text"):
-        if query_text:
-            text_embedding = get_text_embedding(query_text)
-            similar_images = find_similar_images(text_embedding, collection)
-            st.subheader("Similar Items:")
-            for img in similar_images:
-                col1, col2 = st.columns(2)
-                with col1:
-                    st.image(img['info']['image_url'], use_column_width=True)
-                with col2:
-                    st.write(f"Name: {img['info']['name']}")
-                    st.write(f"Brand: {img['info']['brand']}")
-                    category = img['info'].get('category')
-                    if category:
-                        st.write(f"Category: {category}")
-                    st.write(f"Price: {img['info']['price']}")
-                    st.write(f"Discount: {img['info']['discount']}%")
-                    st.write(f"Similarity: {img['similarity']:.2f}")
-        else:
-            st.warning("Please enter a search text.")
+# Text search
+st.sidebar.title("Text Search")
+query_text = st.sidebar.text_input("Enter search text:")
+if st.sidebar.button("Search by Text"):
+    if query_text:
+        text_embedding = get_text_embedding(query_text)
+        similar_images = find_similar_images(text_embedding, collection)
+        st.sidebar.subheader("Similar Items:")
+        for img in similar_images:
+            st.sidebar.image(img['info']['image_url'], use_column_width=True)
+            st.sidebar.write(f"Name: {img['info']['name']}")
+            st.sidebar.write(f"Brand: {img['info']['brand']}")
+            category = img['info'].get('category')
+            if category:
+                st.sidebar.write(f"Category: {category}")
+            st.sidebar.write(f"Price: {img['info']['price']}")
+            st.sidebar.write(f"Discount: {img['info']['discount']}%")
+            st.sidebar.write(f"Similarity: {img['similarity']:.2f}")
+    else:
+        st.sidebar.warning("Please enter a search text.")