Create app.py

app.py

@@ -0,0 +1,196 @@
+import streamlit as st
+import open_clip
+import torch
+import requests
+from PIL import Image
+from io import BytesIO
+import time
+import json
+import numpy as np
+import cv2
+from inference_sdk import InferenceHTTPClient
+import matplotlib.pyplot as plt
+import base64
+
+# Load model and tokenizer
+@st.cache_resource
+def load_model():
+    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)
+
+data = load_data()
+
+# Helper functions
+@st.cache_data
+def download_and_process_image(image_url):
+    try:
+        response = requests.get(image_url)
+        response.raise_for_status()  # Raises an HTTPError for bad responses
+        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}")
+        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):
+    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', [])
+        
+        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']
+                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)
+
+@st.cache_data
+def process_database_cached(data):
+    database_embeddings = []
+    database_info = []
+    for item in data:
+        image_url = item['이미지 링크'][0]
+        product_id = item.get('\ufeff상품 ID') or item.get('상품 ID')
+        
+        image = download_and_process_image(image_url)
+        if image is None:
+            continue
+        
+        # Save the image temporarily
+        temp_path = f"temp_{product_id}.jpg"
+        image.save(temp_path, 'JPEG')
+        
+        database_info.append({
+            'id': product_id,
+            'category': item['카테고리'],
+            'brand': item['브랜드명'],
+            'name': item['제품명'],
+            'price': item['정가'],
+            'discount': item['할인율'],
+            'image_url': image_url,
+            'temp_path': temp_path
+        })
+    
+    return database_info
+
+def process_database(client, data):
+    database_info = process_database_cached(data)
+    database_embeddings = []
+    
+    for item in database_info:
+        segmented_image = segment_image(item['temp_path'], client)
+        embedding = get_image_embedding(segmented_image)
+        database_embeddings.append(embedding)
+    
+    return np.vstack(database_embeddings), database_info
+
+# 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)
+    
+    # Initialize database_embeddings and database_info
+    database_embeddings, database_info = process_database(CLIENT, data)
+
+    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, CLIENT)
+                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.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 your Roboflow API Key to use the app.")