Update app.py

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 transformers import YolosImageProcessor, YolosForObjectDetection
 
 # Load CLIP model and tokenizer
 @st.cache_resource
@@ -22,12 +22,17 @@ def load_clip_model():
 
 clip_model, preprocess_val, tokenizer, device = load_clip_model()
 
-# Load YOLOv8 model
+# Load YOLOS model
 @st.cache_resource
-def load_yolo_model():
-    return YOLO("./best.pt")
+def load_yolos_model():
+    processor = YolosImageProcessor.from_pretrained("valentinafeve/yolos-fashionpedia")
+    model = YolosForObjectDetection.from_pretrained("valentinafeve/yolos-fashionpedia")
+    return processor, model
 
-yolo_model = load_yolo_model()
+yolos_processor, yolos_model = load_yolos_model()
+
+# Define the categories
+CATS = ['shirt, blouse', 'top, t-shirt, sweatshirt', 'sweater', 'cardigan', 'jacket', 'vest', 'pants', 'shorts', 'skirt', 'coat', 'dress', 'jumpsuit', 'cape', 'glasses', 'hat', 'headband, head covering, hair accessory', 'tie', 'glove', 'watch', 'belt', 'leg warmer', 'tights, stockings', 'sock', 'shoe', 'bag, wallet', 'scarf', 'umbrella', 'hood', 'collar', 'lapel', 'epaulette', 'sleeve', 'pocket', 'neckline', 'buckle', 'zipper', 'applique', 'bead', 'bow', 'flower', 'fringe', 'ribbon', 'rivet', 'ruffle', 'sequin', 'tassel']
 
 # Helper functions
 def load_image_from_url(url, max_retries=3):
@@ -42,6 +47,7 @@ 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")
@@ -85,27 +91,32 @@ def find_similar_images(query_embedding, collection, top_k=5):
         })
     return results
 
-
-
 def detect_clothing(image):
-    results = yolo_model(image)
-    detections = results[0].boxes.data.cpu().numpy()
+    inputs = yolos_processor(images=image, return_tensors="pt")
+    outputs = yolos_model(**inputs)
+    
+    target_sizes = torch.tensor([image.size[::-1]])
+    results = yolos_processor.post_process_object_detection(outputs, threshold=0.1, target_sizes=target_sizes)[0]
+    
     categories = []
-    for detection in detections:
-        x1, y1, x2, y2, conf, cls = detection
-        category = yolo_model.names[int(cls)]
-        if category in ['sunglass','hat','jacket','shirt','pants','shorts','skirt','dress','bag','shoe']:
+    for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
+        box = [int(i) for i in box.tolist()]
+        category = yolos_model.config.id2label[label.item()]
+        if category in CATS:
             categories.append({
                 'category': category,
-                'bbox': [int(x1), int(y1), int(x2), int(y2)],
-                'confidence': conf
+                'bbox': box,
+                'confidence': score.item()
             })
     return categories
 
 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:
@@ -115,10 +126,7 @@ if 'detections' not in st.session_state:
 if 'selected_category' not in st.session_state:
     st.session_state.selected_category = None
 
-# Streamlit app
-st.title("Advanced Fashion Search App")
-
-# 단계별 처리
+# Step-by-step processing
 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"):
@@ -136,7 +144,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:")
@@ -184,23 +191,22 @@ 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']}")
-                    st.write(f"Category: {img['info']['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']}")
+            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.")