Update app.py

app.py

@@ -124,6 +124,35 @@ 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']
+        })
+    
+    adjusted_rects = st_img_label(image, box_color="red", rects=rects)
+    
+    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'
@@ -154,10 +183,16 @@ if st.session_state.step == 'input':
                 st.error("Failed to load the image. Please try another URL.")
         else:
             st.warning("Please enter an image URL.")
+    pass
 
 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"):