Spaces:
Runtime error
Runtime error
| 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 SegformerImageProcessor, AutoModelForSemanticSegmentation | |
| import torch.nn as nn | |
| import warnings | |
| # Suppress specific warnings | |
| warnings.filterwarnings("ignore", category=UserWarning, module="transformers.utils.deprecation") | |
| # Load CLIP model and tokenizer | |
| def load_clip_model(): | |
| model, _, 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 | |
| def load_segformer_model(): | |
| processor = SegformerImageProcessor.from_pretrained("mattmdjaga/segformer_b2_clothes") | |
| model = AutoModelForSemanticSegmentation.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 | |
| def load_chromadb(): | |
| client = chromadb.PersistentClient(path="./clothesDB") | |
| collection = client.get_collection(name="clothes_items_ver3") | |
| return collection | |
| collection = load_chromadb() | |
| 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.cpu() | |
| upsampled_logits = nn.functional.interpolate( | |
| logits, | |
| size=image.size[::-1], | |
| mode="bilinear", | |
| align_corners=False, | |
| ) | |
| pred_seg = upsampled_logits.argmax(dim=1)[0].numpy() | |
| categories = segformer_model.config.id2label | |
| segmented_items = [] | |
| for category_id, category_name in categories.items(): | |
| if category_id in pred_seg: | |
| mask = (pred_seg == 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, pred_seg | |
| 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, st.session_state.pred_seg = 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': | |
| col1, col2 = st.columns(2) | |
| with col1: | |
| st.image(st.session_state.query_image, caption="Original Image", use_column_width=True) | |
| with col2: | |
| st.image(st.session_state.pred_seg, caption="Segmentation Map", 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.") | |
| # Display ChromaDB vacuum message | |
| st.sidebar.warning("If you've upgraded ChromaDB from a version below 0.6, you may benefit from vacuuming your database. Run 'chromadb utils vacuum --help' for more information.") |