Update app.py (#4)

- Update app.py (4be1017e99d1bb63b66cf583fb47ef5a2cbdcda3)


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

app.py

@@ -7,9 +7,9 @@ from io import BytesIO
 import time
 import json
 import numpy as np
+from ultralytics import YOLO
 import cv2
 import chromadb
-from ultralytics import YOLO
 
 # Load CLIP model and tokenizer
 @st.cache_resource
@@ -22,16 +22,13 @@ def load_clip_model():
 
 clip_model, preprocess_val, tokenizer, device = load_clip_model()
 
-# Load YOLOS model
+# Load YOLOv8 model
 @st.cache_resource
 def load_yolo_model():
     return YOLO("./best.pt")
 
 yolo_model = load_yolo_model()
 
-# Define the categories
-#CATS = ['shirt, blouse', 'top, t-shirt, sweatshirt', 'sweater', 'cardigan', 'jacket', 'vest', 'pants', 'shorts', 'skirt', 'coat', 'dress', 'jumpsuit', 'cape', 'glasses', 'hat', 'glove', 'shoe', 'bag', 'wallet', 'umbrella', 'hood', 'collar', 'lapel', 'epaulette', 'applique', 'bead', 'bow', 'flower', 'fringe', 'ribbon', 'rivet', 'ruffle', 'sequin', 'tassel']
-
 # Helper functions
 def load_image_from_url(url, max_retries=3):
     for attempt in range(max_retries):
@@ -45,10 +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")
+collection = client.get_collection(name="clothes_items_ver3")
 
 def get_image_embedding(image):
     image_tensor = preprocess_val(image).unsqueeze(0).to(device)
@@ -89,14 +85,11 @@ def find_similar_images(query_embedding, collection, top_k=5):
         })
     return results
 
+
+
 def detect_clothing(image):
-    #inputs = yolos_processor(images=image, return_tensors="pt")
-    #outputs = yolos_model(**inputs)
-    
-    #target_sizes = torch.tensor([image.size[::-1]])
     results = yolo_model(image)
     detections = results[0].boxes.data.cpu().numpy()
-    
     categories = []
     for detection in detections:
         x1, y1, x2, y2, conf, cls = detection
@@ -112,10 +105,7 @@ def detect_clothing(image):
 def crop_image(image, bbox):
     return image.crop((bbox[0], bbox[1], bbox[2], bbox[3]))
 
-# Streamlit app
-st.title("Advanced Fashion Search App")
-
-# Initialize session state
+# 세션 상태 초기화
 if 'step' not in st.session_state:
     st.session_state.step = 'input'
 if 'query_image_url' not in st.session_state:
@@ -125,7 +115,10 @@ if 'detections' not in st.session_state:
 if 'selected_category' not in st.session_state:
     st.session_state.selected_category = None
 
-# Step-by-step processing
+# Streamlit app
+st.title("Advanced Fashion Search App")
+
+# 단계별 처리
 if st.session_state.step == 'input':
     st.session_state.query_image_url = st.text_input("Enter image URL:", st.session_state.query_image_url)
     if st.button("Detect Clothing"):
@@ -143,6 +136,7 @@ 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,7 +173,9 @@ elif st.session_state.step == 'show_results':
         with col2:
             st.write(f"Name: {img['info']['name']}")
             st.write(f"Brand: {img['info']['brand']}")
-            st.write(f"Category: {img['info']['category']}")
+            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}")
@@ -190,22 +186,25 @@ elif st.session_state.step == 'show_results':
         st.session_state.detections = []
         st.session_state.selected_category = None
 
-# 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']}")
-            st.sidebar.write(f"Category: {img['info']['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}")
-            st.sidebar.write("---")
-    else:
-        st.sidebar.warning("Please enter a search text.")
+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.")