Update app.py

app.py

@@ -124,29 +124,30 @@ def detect_clothing(image):
 def crop_image(image, bbox):
     return image.crop((bbox[0], bbox[1], bbox[2], bbox[3]))
 
+# 세션 상태 초기화
+if 'step' not in st.session_state:
+    st.session_state.step = 'input'
+if 'query_image_url' not in st.session_state:
+    st.session_state.query_image_url = ''
+if 'detections' not in st.session_state:
+    st.session_state.detections = []
+if 'selected_category' not in st.session_state:
+    st.session_state.selected_category = None
+
 # Streamlit app
 st.title("Advanced Fashion Search App")
 
-search_type = st.radio("Search by:", ("Image URL", "Text"))
-
-if search_type == "Image URL":
-    query_image_url = st.text_input("Enter image URL:")
+# 단계별 처리
+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 query_image_url:
-            query_image = load_image_from_url(query_image_url)
+        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.image(query_image, caption="Query Image", use_column_width=True)
-                detections = detect_clothing(query_image)
-                
-                if detections:
-                    st.session_state.detections = detections  # 세션 상태에 detections 저장
-                    st.subheader("Detected Clothing Items:")
-                    selected_category = st.selectbox("Select a category to search:", 
-                                                     [f"{d['category']} (Confidence: {d['confidence']:.2f})" for d in detections])
-                    
-                    if st.button("Search Similar Items"):
-                        st.session_state.selected_category = selected_category  # 선택된 카테고리 저장
-                        st.session_state.search_clicked = True  # 검색 버튼 클릭 상태 저장
+                st.session_state.query_image = 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.")
             else:
@@ -154,30 +155,42 @@ if search_type == "Image URL":
         else:
             st.warning("Please enter an image URL.")
 
-    # 검색 결과 표시 로직
-    if 'search_clicked' in st.session_state and st.session_state.search_clicked:
-        selected_detection = next(d for d in st.session_state.detections 
-                                  if f"{d['category']} (Confidence: {d['confidence']:.2f})" == st.session_state.selected_category)
-        cropped_image = crop_image(query_image, selected_detection['bbox'])
-        st.image(cropped_image, caption="Cropped Image", use_column_width=True)
-        query_embedding = get_image_embedding(cropped_image)
-        similar_images = find_similar_images(query_embedding)
-        
-        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}")
-
-        # 검색 완료 후 상태 초기화
-        st.session_state.search_clicked = False
+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:")
+    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"):
+        st.session_state.selected_category = selected_option
+        st.session_state.step = 'show_results'
+
+elif st.session_state.step == 'show_results':
+    st.image(st.session_state.query_image, caption="Query Image", use_column_width=True)
+    selected_detection = next(d for d in st.session_state.detections 
+                              if f"{d['category']} (Confidence: {d['confidence']:.2f})" == st.session_state.selected_category)
+    cropped_image = crop_image(st.session_state.query_image, selected_detection['bbox'])
+    st.image(cropped_image, caption="Cropped Image", use_column_width=True)
+    query_embedding = get_image_embedding(cropped_image)
+    similar_images = find_similar_images(query_embedding)
+    
+    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}")
+    
+    if st.button("Start New Search"):
+        st.session_state.step = 'input'
+        st.session_state.query_image_url = ''
+        st.session_state.detections = []
+        st.session_state.selected_category = None
 
 else:  # Text search
     query_text = st.text_input("Enter search text:")