Create app.py

app.py

@@ -0,0 +1,211 @@
+import streamlit as st
+import open_clip
+import torch
+import requests
+from PIL import Image
+from io import BytesIO
+import time
+import numpy as np
+from ultralytics import YOLO
+import chromadb
+from transformers import pipeline
+from sklearn.metrics.pairwise import cosine_similarity
+
+# Load segmentation model
+segmenter = pipeline(model="mattmdjaga/segformer_b2_clothes")
+
+# 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("./best.pt")
+
+yolo_model = load_yolo_model()
+
+# Load chromaDB
+client = chromadb.PersistentClient(path="./clothesDB_Test")
+#collection = client.get_collection(name="clothes_items_ver3")
+collection = client.get_collection(name="category_upper")
+
+
+# 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()
+
+def segment_clothing(img, clothes=["Hat", "Upper-clothes", "Skirt", "Pants", "Dress", "Belt", "Left-shoe", "Right-shoe", "Scarf"]):
+    # Segment image
+    segments = segmenter(img)
+
+    # Create list of masks
+    mask_list = []
+    detected_categories = []
+    for s in segments:
+        if s['label'] in clothes:
+            mask_list.append(s['mask'])
+            detected_categories.append(s['label'])  # Store detected categories
+
+    # Paste all masks on top of each other
+    final_mask = np.zeros_like(np.array(img)[:, :, 0])  # Initialize mask
+    for mask in mask_list:
+        current_mask = np.array(mask)
+        final_mask = np.maximum(final_mask, current_mask)  # Use maximum to combine masks
+
+    # Convert final mask from np array to PIL image
+    final_mask = Image.fromarray(final_mask.astype(np.uint8) * 255)  # Convert to binary mask
+
+    # Apply mask to original image
+    img_with_alpha = img.convert("RGBA")  # Ensure the image has an alpha channel
+    img_with_alpha.putalpha(final_mask)
+
+    return img_with_alpha.convert("RGB"), final_mask, detected_categories  # Return detected categories
+
+
+#def find_similar_images(query_embedding, collection, top_k=5):
+#    all_embeddings = collection.get(include=['embeddings'])['embeddings']
+#    database_embeddings = np.array(all_embeddings)
+#    similarities = np.dot(database_embeddings, query_embedding.T).squeeze()
+#    top_indices = np.argsort(similarities)[::-1][:top_k]
+#    
+#    all_data = collection.get(include=['metadatas'])['metadatas']
+#    top_metadatas = [all_data[idx] for idx in top_indices]
+    
+#    results = []
+#    for idx, metadata in enumerate(top_metadatas):
+#        results.append({
+#            'info': metadata,
+#            'similarity': similarities[top_indices[idx]]
+#        })
+#    return results
+
+def find_similar_images(query_embedding, collection, top_k=5):
+    # 모든 임베딩을 가져옴
+    all_embeddings = collection.get(include=['embeddings'])['embeddings']
+    database_embeddings = np.array(all_embeddings)
+    
+    # 유사도 계산
+    similarities = cosine_similarity(database_embeddings, query_embedding.reshape(1, -1)).squeeze()
+    top_indices = np.argsort(similarities)[::-1][:top_k]
+    
+    # 메타데이터 가져옴
+    all_data = collection.get(include=['metadatas'])['metadatas']
+    top_metadatas = [all_data[idx] for idx in top_indices]
+    
+    results = []
+    for idx, metadata in enumerate(top_metadatas):
+        # 이미지 URLs 필드가 쉼표로 구분된 문자열로 저장된 경우, 이를 리스트로 변환
+        image_urls = metadata['image_url'].split(',')
+        # 첫 번째 이미지를 대표 이미지로 사용
+        representative_image_url = image_urls[0] if image_urls else None
+        
+        results.append({
+            'info': metadata,
+            'similarity': similarities[top_indices[idx]],
+            'image_url': representative_image_url  # 첫 번째 이미지 URL 사용
+        })
+    return results
+
+
+# 세션 상태 초기화
+if 'step' not in st.session_state:
+    st.session_state.step = 'input'
+if 'query_image_url' not in st.session_state:
+    st.session_state.query_image_url = ''
+if 'detections' not in st.session_state:
+    st.session_state.detections = []
+if 'segmented_image' not in st.session_state:  # Add segmented_image to session state
+    st.session_state.segmented_image = None
+if 'selected_category' not in st.session_state:
+    st.session_state.selected_category = None
+
+# Streamlit app
+st.title("Advanced Fashion Search App")
+
+# 단계별 처리
+if st.session_state.step == 'input':
+    st.session_state.query_image_url = st.text_input("Enter image URL:", st.session_state.query_image_url)
+    if st.button("Detect Clothing"):
+        if st.session_state.query_image_url:
+            query_image = load_image_from_url(st.session_state.query_image_url)
+            if query_image is not None:
+                st.session_state.query_image = query_image
+                # Perform segmentation
+                segmented_image, final_mask, detected_categories = segment_clothing(query_image)
+                st.session_state.segmented_image = segmented_image  # Store segmented image in session state
+                st.session_state.detections = detected_categories  # Store detected categories
+                st.image(segmented_image, caption="Segmented Image", use_column_width=True)
+                if st.session_state.detections:
+                    st.session_state.step = 'select_category'
+                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.")
+
+elif st.session_state.step == 'select_category':
+    st.image(st.session_state.segmented_image, caption="Segmented Image with Detected Categories", use_column_width=True)
+    st.subheader("Detected Clothing Categories:")
+    
+    if st.session_state.detections:
+        selected_category = st.selectbox("Select a category to search:", st.session_state.detections)
+        if st.button("Search Similar Items"):
+            st.session_state.selected_category = selected_category
+            st.session_state.step = 'show_results'
+    else:
+        st.warning("No categories detected.")
+
+elif st.session_state.step == 'show_results':
+    original_image = st.session_state.query_image.convert("RGB")  # Convert to RGB before displaying
+    st.image(original_image, caption="Original Image", use_column_width=True)
+    
+    # Get the embedding of the segmented image
+    query_embedding = get_image_embedding(st.session_state.segmented_image)  # Use the segmented image from session state
+    similar_images = find_similar_images(query_embedding, collection)
+    
+    st.subheader("Similar Items:")
+    for img in similar_images:
+        col1, col2 = st.columns(2)
+        with col1:
+            st.image(img['image_url'], use_column_width=True)
+        with col2:
+            st.write(f"Name: {img['info']['name']}")
+            st.write(f"Brand: {img['info']['brand']}")
+            category = img['info'].get('category')
+            if category:
+                st.write(f"Category: {category}")
+            st.write(f"Price: {img['info']['price']}")
+            st.write(f"Discount: {img['info']['discount']}%")
+            st.write(f"Similarity: {img['similarity']:.2f}")
+
+    if st.button("Start New Search"):
+        st.session_state.step = 'input'
+        st.session_state.query_image_url = ''
+        st.session_state.detections = []
+        st.session_state.segmented_image = None  # Reset segmented_image