Update app.py

app.py

@@ -11,6 +11,8 @@ import cv2
 from inference_sdk import InferenceHTTPClient
 import matplotlib.pyplot as plt
 import base64
+import os
+import pickle
 
 # Load model and tokenizer
 @st.cache_resource
@@ -30,88 +32,77 @@ def load_data():
 
 data = load_data()
 
-# Helper functions
-@st.cache_data
+def setup_roboflow_client(api_key):
+    return InferenceHTTPClient(
+        api_url="https://outline.roboflow.com",
+        api_key=api_key
+    )
+
 def download_and_process_image(image_url):
     try:
         response = requests.get(image_url)
-        response.raise_for_status()  # Raises an HTTPError for bad responses
+        response.raise_for_status()
         image = Image.open(BytesIO(response.content))
-        
-        # Convert image to RGB mode if it's in RGBA mode
         if image.mode == 'RGBA':
             image = image.convert('RGB')
-        
         return image
-    except requests.RequestException as e:
-        st.error(f"Error downloading image: {e}")
-        return None
     except Exception as e:
-        st.error(f"Error processing image: {e}")
+        st.error(f"Error downloading/processing image: {str(e)}")
         return None
 
-def get_image_embedding(image):
-    image_tensor = preprocess_val(image).unsqueeze(0).to(device)
-    with torch.no_grad():
-        image_features = model.encode_image(image_tensor)
-        image_features /= image_features.norm(dim=-1, keepdim=True)
-    return image_features.cpu().numpy()
-
-def setup_roboflow_client(api_key):
-    return InferenceHTTPClient(
-        api_url="https://outline.roboflow.com",
-        api_key=api_key
-    )
-
-def segment_image(image_path, client):
+def segment_image_and_get_categories(image_path, client):
     try:
-        # 이미지 파일 읽기
         with open(image_path, "rb") as image_file:
             image_data = image_file.read()
         
-        # 이미지를 base64로 인코딩
         encoded_image = base64.b64encode(image_data).decode('utf-8')
         
-        # 원본 이미지 로드
         image = cv2.imread(image_path)
         image = cv2.resize(image, (800, 600))
         mask = np.zeros(image.shape, dtype=np.uint8)
         
-        # Roboflow API 호출
         results = client.infer(encoded_image, model_id="closet/1")
         
-        # 결과가 이미 딕셔너리인 경우 JSON 파싱 단계 제거
         if isinstance(results, dict):
             predictions = results.get('predictions', [])
         else:
-            # 문자열인 경우에만 JSON 파싱
             predictions = json.loads(results).get('predictions', [])
         
+        categories = []
         if predictions:
             for prediction in predictions:
                 points = prediction['points']
                 pts = np.array([[p['x'], p['y']] for p in points], np.int32)
-                scale_x = image.shape[1] / results['image']['width']
-                scale_y = image.shape[0] / results['image']['height']
+                scale_x = image.shape[1] / results.get('image', {}).get('width', 1)
+                scale_y = image.shape[0] / results.get('image', {}).get('height', 1)
                 pts = pts * [scale_x, scale_y]
                 pts = pts.astype(np.int32)
                 pts = pts.reshape((-1, 1, 2))
-                cv2.fillPoly(mask, [pts], color=(255, 255, 255))  # White mask
+                cv2.fillPoly(mask, [pts], color=(255, 255, 255))
+                
+                category = prediction.get('class', 'Unknown')
+                confidence = prediction.get('confidence', 0)
+                categories.append(f"{category} ({confidence:.2f})")
             
             segmented_image = cv2.bitwise_and(image, mask)
         else:
             st.warning("No predictions found in the image. Returning original image.")
             segmented_image = image
         
-        return Image.fromarray(cv2.cvtColor(segmented_image, cv2.COLOR_BGR2RGB))
+        return Image.fromarray(cv2.cvtColor(segmented_image, cv2.COLOR_BGR2RGB)), categories
     except Exception as e:
         st.error(f"Error in segmentation: {str(e)}")
-        # 원본 이미지를 다시 읽어 반환
-        return Image.open(image_path)
+        return Image.open(image_path), []
+
+def get_image_embedding(image):
+    image_tensor = preprocess_val(image).unsqueeze(0).to(device)
+    with torch.no_grad():
+        image_features = model.encode_image(image_tensor)
+        image_features /= image_features.norm(dim=-1, keepdim=True)
+    return image_features.cpu().numpy()
 
 @st.cache_data
 def process_database_cached(data):
-    database_embeddings = []
     database_info = []
     for item in data:
         image_url = item['이미지 링크'][0]
@@ -121,7 +112,6 @@ def process_database_cached(data):
         if image is None:
             continue
         
-        # Save the image temporarily
         temp_path = f"temp_{product_id}.jpg"
         image.save(temp_path, 'JPEG')
         
@@ -140,17 +130,42 @@ def process_database_cached(data):
 
 def process_database(client, data):
     database_info = process_database_cached(data)
-    database_embeddings = []
+    cache_dir = "segmentation_cache"
+    os.makedirs(cache_dir, exist_ok=True)
     
+    database_embeddings = []
     for item in database_info:
-        segmented_image = segment_image(item['temp_path'], client)
+        cache_file = os.path.join(cache_dir, f"{item['id']}_segmented.pkl")
+        
+        if os.path.exists(cache_file):
+            with open(cache_file, 'rb') as f:
+                segmented_image, categories = pickle.load(f)
+        else:
+            segmented_image, categories = segment_image_and_get_categories(item['temp_path'], client)
+            with open(cache_file, 'wb') as f:
+                pickle.dump((segmented_image, categories), f)
+        
         embedding = get_image_embedding(segmented_image)
         database_embeddings.append(embedding)
+        item['categories'] = categories
     
     return np.vstack(database_embeddings), database_info
 
+def find_similar_images(query_embedding, database_embeddings, database_info, top_k=5):
+    similarities = np.dot(database_embeddings, query_embedding.T).squeeze()
+    top_indices = np.argsort(similarities)[::-1][:top_k]
+    
+    results = []
+    for idx in top_indices:
+        results.append({
+            'info': database_info[idx],
+            'similarity': similarities[idx]
+        })
+    
+    return results
+
 # Streamlit app
-st.title("Fashion Search App with Segmentation")
+st.title("Fashion Search App with Segmentation and Category Detection")
 
 # API Key input
 api_key = st.text_input("Enter your Roboflow API Key", type="password")
@@ -168,17 +183,18 @@ if api_key:
         
         if st.button('Find Similar Items'):
             with st.spinner('Processing...'):
-                # Save uploaded image temporarily
                 temp_path = "temp_upload.jpg"
                 image.save(temp_path)
                 
-                # Segment the uploaded image
-                segmented_image = segment_image(temp_path, CLIENT)
+                segmented_image, input_categories = segment_image_and_get_categories(temp_path, CLIENT)
                 st.image(segmented_image, caption='Segmented Image', use_column_width=True)
                 
-                # Get embedding for segmented image
+                st.subheader("Detected Categories in Input Image:")
+                for category in input_categories:
+                    st.write(category)
+                
                 query_embedding = get_image_embedding(segmented_image)
-                similar_images = find_similar_images(query_embedding)
+                similar_images = find_similar_images(query_embedding, database_embeddings, database_info)
                 
                 st.subheader("Similar Items:")
                 for img in similar_images:
@@ -192,5 +208,9 @@ if api_key:
                         st.write(f"Price: {img['info']['price']}")
                         st.write(f"Discount: {img['info']['discount']}%")
                         st.write(f"Similarity: {img['similarity']:.2f}")
+                        
+                        st.write("Detected Categories:")
+                        for category in img['info']['categories']:
+                            st.write(category)
 else:
     st.warning("Please enter your Roboflow API Key to use the app.")