Update app.py

app.py

@@ -10,13 +10,6 @@ import numpy as np
 from ultralytics import YOLO
 import cv2
 
-try:
-    from streamlit_img_label import st_img_label
-    from streamlit_img_label.manage import ImageManager
-except ImportError:
-    st.error("Required modules are not installed. Please install 'streamlit-img-label' and 'pascal-voc-writer'.")
-    st.stop()
-
 # Load CLIP model and tokenizer
 @st.cache_resource
 def load_clip_model():
@@ -113,14 +106,6 @@ def find_similar_images(query_embedding, top_k=5):
         })
     return results
 
-# 먼저 라벨 리스트를 정의합니다
-labels = [
-    "bag", "coat", "dress", "fabric", "hat", "Hoodie", "jacket", "Jacket",
-    "jean", "m2m", "Mid-lenght dress", "Pants", "plain", "shirt", "Shirt",
-    "shoe", "short", "shorts", "skirt", "slacks", "suit", "sunglass",
-    "sweat", "tie", "tracksuit", "tshirt"
-]
-
 def detect_clothing(image):
     results = yolo_model(image)
     detections = results[0].boxes.data.cpu().numpy()
@@ -128,7 +113,7 @@ def detect_clothing(image):
     for detection in detections:
         x1, y1, x2, y2, conf, cls = detection
         category = yolo_model.names[int(cls)]
-        if category in labels:
+        if category in ['sunglass','hat','jacket','shirt','pants','shorts','skirt','dress','bag','shoe']:
             categories.append({
                 'category': category,
                 'bbox': [int(x1), int(y1), int(x2), int(y2)],
@@ -139,39 +124,6 @@ def detect_clothing(image):
 def crop_image(image, bbox):
     return image.crop((bbox[0], bbox[1], bbox[2], bbox[3]))
 
-def adjust_bounding_boxes(image, detections):
-    img_height, img_width = image.size
-    rects = []
-    for detection in detections:
-        x1, y1, x2, y2 = detection['bbox']
-        rects.append({
-            "left": x1 / img_width,
-            "top": y1 / img_height,
-            "width": (x2 - x1) / img_width,
-            "height": (y2 - y1) / img_height,
-            "label": detection['category']
-        })
-    
-    try:
-        adjusted_rects = st_img_label(image, box_color="red", rects=rects)
-    except Exception as e:
-        st.error(f"Error in st_img_label: {str(e)}")
-        return detections
-    
-    adjusted_detections = []
-    for rect, detection in zip(adjusted_rects, detections):
-        x1 = rect["left"] * img_width
-        y1 = rect["top"] * img_height
-        x2 = x1 + (rect["width"] * img_width)
-        y2 = y1 + (rect["height"] * img_height)
-        adjusted_detections.append({
-            'category': rect["label"],
-            'bbox': [int(x1), int(y1), int(x2), int(y2)],
-            'confidence': detection['confidence']
-        })
-    
-    return adjusted_detections
-
 # 세션 상태 초기화
 if 'step' not in st.session_state:
     st.session_state.step = 'input'
@@ -206,11 +158,6 @@ if st.session_state.step == 'input':
 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:")
-    
-    # 경계 상자 조정 기능 추가
-    adjusted_detections = adjust_bounding_boxes(st.session_state.query_image, st.session_state.detections)
-    st.session_state.detections = adjusted_detections
-    
     options = [f"{d['category']} (Confidence: {d['confidence']:.2f})" for d in st.session_state.detections]
     selected_option = st.selectbox("Select a category to search:", options)
     if st.button("Search Similar Items"):