Create app.py

app.py

@@ -0,0 +1,190 @@
+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
+from ultralytics import YOLO
+import cv2
+
+# Load CLIP model and tokenizer
+@st.cache_resource
+def load_clip_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')
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model.to(device)
+    return model, preprocess_val, tokenizer, device
+
+clip_model, preprocess_val, tokenizer, device = load_clip_model()
+
+# Load YOLOv8 model
+@st.cache_resource
+def load_yolo_model():
+    return YOLO("./yolov8m.pt")
+
+yolo_model = load_yolo_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
+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:
+            if attempt < max_retries - 1:
+                time.sleep(1)
+            else:
+                return None
+
+def get_image_embedding(image):
+    image_tensor = preprocess_val(image).unsqueeze(0).to(device)
+    with torch.no_grad():
+        image_features = clip_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]
+        image = load_image_from_url(image_url)
+        if image is not None:
+            embedding = get_image_embedding(image)
+            database_embeddings.append(embedding)
+            database_info.append({
+                'id': item['\ufeff상품 ID'],
+                'category': item['카테고리'],
+                'brand': item['브랜드명'],
+                'name': item['제품명'],
+                'price': item['정가'],
+                '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 = clip_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
+
+def detect_clothing(image):
+    results = yolo_model(image)
+    detections = results[0].boxes.data.cpu().numpy()
+    categories = []
+    for detection in detections:
+        x1, y1, x2, y2, conf, cls = detection
+        category = yolo_model.names[int(cls)]
+        if category in ['top', 'bottom', 'hat', 'shoes']:  # Add more categories as needed
+            categories.append({
+                'category': category,
+                'bbox': [int(x1), int(y1), int(x2), int(y2)],
+                'confidence': conf
+            })
+    return categories
+
+def crop_image(image, bbox):
+    return image.crop((bbox[0], bbox[1], bbox[2], bbox[3]))
+
+# 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.button("Detect Clothing"):
+        if query_image_url:
+            query_image = load_image_from_url(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.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"):
+                        selected_detection = next(d for d in detections if f"{d['category']} (Confidence: {d['confidence']:.2f})" == selected_category)
+                        cropped_image = crop_image(query_image, selected_detection['bbox'])
+                        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}")
+                else:
+                    st.warning("No clothing items detected in the image.")
+            else:
+                st.error("Failed to load 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.")