Create app.py

app.py

@@ -0,0 +1,218 @@
+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
+import chromadb
+from transformers import SegformerForSemanticSegmentation, SegformerImageProcessor
+
+# 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 SegFormer model
+@st.cache_resource
+def load_segformer_model():
+    model = SegformerForSemanticSegmentation.from_pretrained("mattmdjaga/segformer_b2_clothes")
+    processor = SegformerImageProcessor.from_pretrained("mattmdjaga/segformer_b2_clothes")
+    return model, processor
+
+segformer_model, segformer_processor = load_segformer_model()
+
+# 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
+
+#Load chromaDB
+client = chromadb.PersistentClient(path="./clothesDB")
+collection = client.get_collection(name="clothes_items_ver3")
+
+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 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 get_all_embeddings_from_collection(collection):
+    all_embeddings = collection.get(include=['embeddings'])['embeddings']
+    return np.array(all_embeddings)
+
+def get_metadata_from_ids(collection, ids):
+    results = collection.get(ids=ids)
+    return results['metadatas']
+    
+def find_similar_images(query_embedding, collection, top_k=5):
+    database_embeddings = get_all_embeddings_from_collection(collection)
+    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 segment_clothing(image):
+    inputs = segformer_processor(images=image, return_tensors="pt")
+    outputs = segformer_model(**inputs)
+    logits = outputs.logits.squeeze()
+    predicted_mask = logits.argmax(dim=0).numpy()
+    
+    categories = segformer_model.config.id2label
+    segmented_items = []
+    
+    for category_id, category_name in categories.items():
+        if category_id in predicted_mask:
+            mask = (predicted_mask == category_id).astype(np.uint8)
+            contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+            if contours:
+                x, y, w, h = cv2.boundingRect(max(contours, key=cv2.contourArea))
+                segmented_items.append({
+                    'category': category_name,
+                    'bbox': [x, y, x+w, y+h],
+                    'mask': mask
+                })
+    
+    return segmented_items
+
+def crop_image(image, bbox):
+    return image.crop((bbox[0], bbox[1], bbox[2], bbox[3]))
+
+# Streamlit app
+st.title("Advanced Fashion Search App")
+
+# Initialize session state
+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 'segmentations' not in st.session_state:
+    st.session_state.segmentations = []
+if 'selected_category' not in st.session_state:
+    st.session_state.selected_category = None
+
+# Step-by-step processing
+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("Segment 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
+                st.session_state.segmentations = segment_clothing(query_image)
+                if st.session_state.segmentations:
+                    st.session_state.step = 'select_category'
+                else:
+                    st.warning("No clothing items segmented 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.query_image, caption="Query Image", use_column_width=True)
+    st.subheader("Segmented Clothing Items:")
+    
+    for segmentation in st.session_state.segmentations:
+        col1, col2 = st.columns([1, 3])
+        with col1:
+            st.write(f"{segmentation['category']}")
+        with col2:
+            cropped_image = crop_image(st.session_state.query_image, segmentation['bbox'])
+            st.image(cropped_image, caption=segmentation['category'], use_column_width=True)
+    
+    options = [s['category'] for s in st.session_state.segmentations]
+    selected_option = st.selectbox("Select a category to search:", options)
+    
+    if st.button("Search Similar Items"):
+        st.session_state.selected_category = selected_option
+        st.session_state.step = 'show_results'
+
+elif st.session_state.step == 'show_results':
+    st.image(st.session_state.query_image, caption="Query Image", use_column_width=True)
+    selected_segmentation = next(s for s in st.session_state.segmentations 
+                                 if s['category'] == st.session_state.selected_category)
+    cropped_image = crop_image(st.session_state.query_image, selected_segmentation['bbox'])
+    st.image(cropped_image, caption="Cropped Image", use_column_width=True)
+    query_embedding = get_image_embedding(cropped_image)
+    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['info']['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.segmentations = []
+        st.session_state.selected_category = None
+
+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, collection)
+            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']}")
+                    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}")
+        else:
+            st.warning("Please enter a search text.")