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 ['sunglass','hat','jacket','shirt','pants','shorts','skirt','dress','bag','shoe']: 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']) 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) 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.")