app.py

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