Update app.py

app.py

@@ -50,11 +50,6 @@ def download_and_process_image(image_url):
         st.error(f"Error processing image: {e}")
         return None
 
-def segment_image(image_path):
-    # Implement your segmentation logic here
-    # For now, we'll just return the original image
-    return Image.open(image_path)
-
 def get_image_embedding(image):
     image_tensor = preprocess_val(image).unsqueeze(0).to(device)
     with torch.no_grad():
@@ -68,9 +63,49 @@ def setup_roboflow_client(api_key):
         api_key=api_key
     )
 
+def segment_image(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")
+        results = json.loads(results)
+        
+        if 'predictions' in results:
+            for prediction in results['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']
+                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
+            
+            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))
+    except Exception as e:
+        st.error(f"Error in segmentation: {str(e)}")
+        # 원본 이미지를 다시 읽어 반환
+        return Image.open(image_path)
+
 # Process database with segmentation
 @st.cache_data
-def process_database():
+def process_database(client):
     database_embeddings = []
     database_info = []
     for item in data:
@@ -85,7 +120,7 @@ def process_database():
         temp_path = f"temp_{product_id}.jpg"
         image.save(temp_path, 'JPEG')
         
-        segmented_image = segment_image(temp_path)
+        segmented_image = segment_image(temp_path, client)
         embedding = get_image_embedding(segmented_image)
         
         database_embeddings.append(embedding)
@@ -101,22 +136,6 @@ def process_database():
     
     return np.vstack(database_embeddings), database_info
 
-# Initialize database_embeddings and database_info
-database_embeddings, database_info = process_database()
-
-def find_similar_images(query_embedding, 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")
 
@@ -125,6 +144,9 @@ api_key = st.text_input("Enter your Roboflow API Key", type="password")
 
 if api_key:
     CLIENT = setup_roboflow_client(api_key)
+    
+    # Initialize database_embeddings and database_info
+    database_embeddings, database_info = process_database(CLIENT)
 
     uploaded_file = st.file_uploader("Choose an image...", type="jpg")
     if uploaded_file is not None:
@@ -138,7 +160,7 @@ if api_key:
                 image.save(temp_path)
                 
                 # Segment the uploaded image
-                segmented_image = segment_image(temp_path)
+                segmented_image = segment_image(temp_path, CLIENT)
                 st.image(segmented_image, caption='Segmented Image', use_column_width=True)
                 
                 # Get embedding for segmented image