Update app.py

app.py

@@ -9,7 +9,7 @@ import json
 import numpy as np
 import cv2
 import chromadb
-from transformers import YolosImageProcessor, YolosForObjectDetection
+from ultralytics import YOLO
 
 # Load CLIP model and tokenizer
 @st.cache_resource
@@ -24,15 +24,13 @@ clip_model, preprocess_val, tokenizer, device = load_clip_model()
 
 # Load YOLOS model
 @st.cache_resource
-def load_yolos_model():
-    processor = YolosImageProcessor.from_pretrained("valentinafeve/yolos-fashionpedia")
-    model = YolosForObjectDetection.from_pretrained("valentinafeve/yolos-fashionpedia")
-    return processor, model
+def load_yolo_model():
+    return YOLO("./best.pt")
 
-yolos_processor, yolos_model = load_yolos_model()
+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']
+#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):
@@ -92,21 +90,22 @@ 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)
+    #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]
+    #target_sizes = torch.tensor([image.size[::-1]])
+    results = yolo_model(image)
+    detections = results[0].boxes.data.cpu().numpy()
     
     categories = []
-    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:
+    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']:
             categories.append({
                 'category': category,
-                'bbox': box,
-                'confidence': score.item()
+                'bbox': [int(x1), int(y1), int(x2), int(y2)],
+                'confidence': conf
             })
     return categories