Update app.py (#2)

- Update app.py (2cf5234ea2ba741f1aad1ac3856e751e40087951)


Co-authored-by: Jo sungmin <JoJosmin@users.noreply.huggingface.co>

app.py

@@ -9,6 +9,7 @@ import json
 import numpy as np
 from ultralytics import YOLO
 import cv2
+import chromadb
 
 # Load CLIP model and tokenizer
 @st.cache_resource
@@ -28,14 +29,6 @@ def load_yolo_model():
 
 yolo_model = load_yolo_model()
 
-# Load and process data
-@st.cache_data
-def load_data():
-    with open('./musinsa-final.json', 'r', encoding='utf-8') as f:
-        return json.load(f)
-
-data = load_data()
-
 # Helper functions
 def load_image_from_url(url, max_retries=3):
     for attempt in range(max_retries):
@@ -49,6 +42,9 @@ 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="fashion_items_ver2")
 
 def get_image_embedding(image):
     image_tensor = preprocess_val(image).unsqueeze(0).to(device)
@@ -57,37 +53,6 @@ def get_image_embedding(image):
         image_features /= image_features.norm(dim=-1, keepdim=True)
     return image_features.cpu().numpy()
 
-@st.cache_data
-def process_database():
-    database_embeddings = []
-    database_info = []
-
-    for item in data:
-        image_url = item['이미지 링크'][0]
-        image = load_image_from_url(image_url)
-        if image is not None:
-            embedding = get_image_embedding(image)
-            database_embeddings.append(embedding)
-            database_info.append({
-                'id': item['\ufeff상품 ID'],
-                'category': item['카테고리'],
-                'brand': item['브랜드명'],
-                'name': item['제품명'],
-                'price': item['정가'],
-                'discount': item['할인율'],
-                'image_url': image_url
-            })
-        else:
-            st.warning(f"Skipping item {item['상품 ID']} due to image loading failure")
-
-    if database_embeddings:
-        return np.vstack(database_embeddings), database_info
-    else:
-        st.error("No valid embeddings were generated.")
-        return None, None
-
-database_embeddings, database_info = process_database()
-
 def get_text_embedding(text):
     text_tokens = tokenizer([text]).to(device)
     with torch.no_grad():
@@ -95,17 +60,33 @@ def get_text_embedding(text):
         text_features /= text_features.norm(dim=-1, keepdim=True)
     return text_features.cpu().numpy()
 
-def find_similar_images(query_embedding, top_k=5):
+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 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]
+    
+    all_data = collection.get(include=['metadatas'])['metadatas']
+    
+    top_metadatas = [all_data[idx] for idx in top_indices]
+    
     results = []
-    for idx in top_indices:
+    for idx, metadata in enumerate(top_metadatas):
         results.append({
-            'info': database_info[idx],
-            'similarity': similarities[idx]
+            'info': metadata,
+            'similarity': similarities[top_indices[idx]]
         })
     return results
 
+
+
 def detect_clothing(image):
     results = yolo_model(image)
     detections = results[0].boxes.data.cpu().numpy()
@@ -182,7 +163,7 @@ elif st.session_state.step == 'show_results':
     cropped_image = crop_image(st.session_state.query_image, selected_detection['bbox'])
     st.image(cropped_image, caption="Cropped Image", use_column_width=True)
     query_embedding = get_image_embedding(cropped_image)
-    similar_images = find_similar_images(query_embedding)
+    similar_images = find_similar_images(query_embedding, collection)
     
     st.subheader("Similar Items:")
     for img in similar_images:
@@ -208,7 +189,7 @@ else:  # Text search
     if st.button("Search by Text"):
         if query_text:
             text_embedding = get_text_embedding(query_text)
-            similar_images = find_similar_images(text_embedding)
+            similar_images = find_similar_images(text_embedding, collection)
             st.subheader("Similar Items:")
             for img in similar_images:
                 col1, col2 = st.columns(2)