Update app.py

app.py

@@ -8,6 +8,7 @@ import time
 import json
 import numpy as np
 import onnxruntime as ort
+from ultralytics import YOLO
 import cv2
 import chromadb
 
@@ -22,13 +23,11 @@ def load_clip_model():
 
 clip_model, preprocess_val, tokenizer, device = load_clip_model()
 
-# ONNX 모델 로드
 @st.cache_resource
-def load_onnx_model():
-    session = ort.InferenceSession("./accessary_weights.onnx")
-    return session
+def load_yolo_model():
+    return YOLO("./best.pt")
 
-onnx_session = load_onnx_model()
+yolo_model = load_yolo_model()
 
 # URL에서 이미지 로드
 def load_image_from_url(url, max_retries=3):
@@ -48,7 +47,6 @@ def load_image_from_url(url, max_retries=3):
 client = chromadb.PersistentClient(path="./accessaryDB")
 collection = client.get_collection(name="accessary_items_ver2")
 
-# CLIP 이미지 임베딩 추출
 def get_image_embedding(image):
     image_tensor = preprocess_val(image).unsqueeze(0).to(device)
     with torch.no_grad():
@@ -56,7 +54,6 @@ def get_image_embedding(image):
         image_features /= image_features.norm(dim=-1, keepdim=True)
     return image_features.cpu().numpy()
 
-# CLIP 텍스트 임베딩 추출
 def get_text_embedding(text):
     text_tokens = tokenizer([text]).to(device)
     with torch.no_grad():
@@ -64,17 +61,14 @@ 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)
 
-# ID를 통해 메타데이터 가져오기
 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()
@@ -92,40 +86,19 @@ def find_similar_images(query_embedding, collection, top_k=5):
         })
     return results
 
-onnx_model_labels = ['Bracelets', 'Broches', 'belt', 'earring', 'maangtika', 'necklace', 'nose ring', 'ring', 'tiara']
-
-# ONNX 모델에 맞춘 전처리 함수
-def preprocess_for_onnx(image, input_size=(640, 640)):
-    resized_image = image.resize(input_size)
-    image_np = np.array(resized_image).astype(np.float32) / 255.0
-    image_np = np.transpose(image_np, (2, 0, 1))
-    input_tensor = np.expand_dims(image_np, axis=0)
-    return input_tensor
-
-# 의류 탐지
-def detect_clothing_onnx(image):
-    input_tensor = preprocess_for_onnx(image)  # 전처리 함수 호출
-    outputs = onnx_session.run(None, {onnx_session.get_inputs()[0].name: input_tensor})
-    
-    detections = outputs[0]  # 첫 번째 출력값이 탐지 결과라고 가정
+def detect_clothing(image):
+    results = yolo_model(image)
+    detections = results[0].boxes.data.cpu().numpy()
     categories = []
-    
     for detection in detections:
-        # detection에서 필요한 값 추출
-        x1, y1, x2, y2, conf, cls = detection[:6]
-        
-        # conf가 배열인 경우, 최대 값을 사용
-        if isinstance(conf, np.ndarray):
-            conf = np.max(conf)  # 배열에서 최대 신뢰도 값
-        
-        if conf > 0.3:  # 신뢰도 임계값 설정
-            category = onnx_model_labels[int(cls)]
+        x1, y1, x2, y2, conf, cls = detection
+        category = yolo_model.names[int(cls)]
+        if category in ['Bracelets', 'Broches', 'bag', 'belt', 'earring', 'maangtika', 'necklace', 'nose ring', 'ring', 'tiara']:
             categories.append({
                 'category': category,
-                'bbox': [x1, y1, x2, y2],
+                'bbox': [int(x1), int(y1), int(x2), int(y2)],
                 'confidence': conf
             })
-    
     return categories
 
 # 이미지 자르기
@@ -143,21 +116,21 @@ if 'selected_category' not in st.session_state:
     st.session_state.selected_category = None
 
 # Streamlit app
-st.title("Advanced Fashion Search App")
+st.title("Accessary 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"):
+    if st.button("Detect acsseary"):
         if st.session_state.query_image_url:
             query_image = load_image_from_url(st.session_state.query_image_url)
             if query_image is not None:
                 st.session_state.query_image = query_image
-                st.session_state.detections = detect_clothing_onnx(query_image)
+                st.session_state.detections = detect_clothing(query_image)
                 if st.session_state.detections:
                     st.session_state.step = 'select_category'
                 else:
-                    st.warning("No clothing items detected in the image.")
+                    st.warning("No items detected in the image.")
             else:
                 st.error("Failed to load the image. Please try another URL.")
         else: