Update app.py

app.py

@@ -7,68 +7,90 @@ from io import BytesIO
 import time
 import json
 import numpy as np
+import cv2
+from inference_sdk import InferenceHTTPClient
+import matplotlib.pyplot as plt
 
 # Load model and tokenizer
 @st.cache_resource
 def load_model():
-    model, preprocess_train, preprocess_val = open_clip.create_model_and_transforms('hf-hub:Marqo/marqo-fashionSigLIP')
-    tokenizer = open_clip.get_tokenizer('hf-hub:Marqo/marqo-fashionSigLIP')
+    model, preprocess_val, tokenizer = open_clip.create_model_and_transforms('hf-hub:Marqo/marqo-fashionSigLIP')
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     model.to(device)
     return model, preprocess_val, tokenizer, device
 
 model, preprocess_val, tokenizer, device = load_model()
 
-# Load and process data
-@st.cache_data
-def load_data():
-    with open('./musinsa-final.json', 'r', encoding='utf-8') as f:
-        return json.load(f)
+# Roboflow client setup function
+def setup_roboflow_client(api_key):
+    return InferenceHTTPClient(
+        api_url="https://outline.roboflow.com",
+        api_key=api_key
+    )
 
-data = load_data()
-
-# Helper functions
-def load_image_from_url(url, max_retries=3):
-    for attempt in range(max_retries):
-        try:
-            response = requests.get(url, timeout=10)
-            response.raise_for_status()
-            img = Image.open(BytesIO(response.content)).convert('RGB')
-            return img
-        except (requests.RequestException, Image.UnidentifiedImageError) as e:
-            #st.warning(f"Attempt {attempt + 1} failed: {str(e)}")
-            if attempt < max_retries - 1:
-                time.sleep(1)
-            else:
-                #st.error(f"Failed to load image from {url} after {max_retries} attempts")
-                return None
-
-def get_image_embedding_from_url(image_url):
-    image = load_image_from_url(image_url)
-    if image is None:
-        return None
-
-    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():
-    database_embeddings = []
-    database_info = []
-
-    for item in data:
-        image_url = item['이미지 링크'][0]
-        embedding = get_image_embedding_from_url(image_url)
-
-        if embedding is not None:
+# Streamlit app
+st.title("Fashion Search App with Segmentation")
+
+# API Key input
+api_key = st.text_input("Enter your Roboflow API Key", type="password")
+
+if api_key:
+    CLIENT = setup_roboflow_client(api_key)
+
+    def segment_image(image_path):
+        results = CLIENT.infer(image_path, model_id="closet/1")
+        results = json.loads(results)
+        
+        image = cv2.imread(image_path)
+        image = cv2.resize(image, (800, 600))
+        mask = np.zeros(image.shape, dtype=np.uint8)
+        
+        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)
+        return Image.fromarray(cv2.cvtColor(segmented_image, cv2.COLOR_BGR2RGB))
+
+    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()
+
+    # Load and process data
+    @st.cache_data
+    def load_data():
+        with open('musinsa-final.json', 'r', encoding='utf-8') as f:
+            return json.load(f)
+
+    data = load_data()
+
+    # Process database with segmentation
+    @st.cache_data
+    def process_database():
+        database_embeddings = []
+        database_info = []
+        for item in data:
+            image_url = item['이미지 링크'][0]
+            image_path = f"temp_{item['상품 ID']}.jpg"
+            response = requests.get(image_url)
+            with open(image_path, 'wb') as f:
+                f.write(response.content)
+            
+            segmented_image = segment_image(image_path)
+            embedding = get_image_embedding(segmented_image)
+            
             database_embeddings.append(embedding)
             database_info.append({
-                'id': item['\ufeff상품 ID'],
+                'id': item['상품 ID'],
                 'category': item['카테고리'],
                 'brand': item['브랜드명'],
                 'name': item['제품명'],
@@ -76,52 +98,43 @@ def process_database():
                 'discount': item['할인율'],
                 'image_url': image_url
             })
-        else:
-            st.warning(f"Skipping item {item['상품 ID']} due to image loading failure")
-
-    if database_embeddings:
+        
         return np.vstack(database_embeddings), database_info
-    else:
-        st.error("No valid embeddings were generated.")
-        return None, None
-
-database_embeddings, database_info = process_database()
 
-def get_text_embedding(text):
-    text_tokens = tokenizer([text]).to(device)
-
-    with torch.no_grad():
-        text_features = model.encode_text(text_tokens)
-        text_features /= text_features.norm(dim=-1, keepdim=True)
-
-    return text_features.cpu().numpy()
-
-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")
-
-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.button("Search by Image"):
-        if query_image_url:
-            query_embedding = get_image_embedding_from_url(query_image_url)
-            if query_embedding is not None:
+    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
+
+    uploaded_file = st.file_uploader("Choose an image...", type="jpg")
+    if uploaded_file is not None:
+        image = Image.open(uploaded_file)
+        st.image(image, caption='Uploaded Image', use_column_width=True)
+        
+        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)
+                st.image(segmented_image, caption='Segmented Image', use_column_width=True)
+                
+                # Get embedding for segmented image
+                query_embedding = get_image_embedding(segmented_image)
                 similar_images = find_similar_images(query_embedding)
-                st.image(query_image_url, caption="Query Image", use_column_width=True)
+                
                 st.subheader("Similar Items:")
                 for img in similar_images:
                     col1, col2 = st.columns(2)
@@ -134,28 +147,5 @@ if search_type == "Image URL":
                         st.write(f"Price: {img['info']['price']}")
                         st.write(f"Discount: {img['info']['discount']}%")
                         st.write(f"Similarity: {img['similarity']:.2f}")
-            else:
-                st.error("Failed to process the image. Please try another URL.")
-        else:
-            st.warning("Please enter an image URL.")
-
-else:  # Text search
-    query_text = st.text_input("Enter search text:")
-    if st.button("Search by Text"):
-        if query_text:
-            text_embedding = get_text_embedding(query_text)
-            similar_images = find_similar_images(text_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}")
-        else:
-            st.warning("Please enter a search text.")
+else:
+    st.warning("Please enter your Roboflow API Key to use the app.")