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 transformers import pipeline
import chromadb
from sklearn.metrics.pairwise import euclidean_distances  
from sklearn.preprocessing import normalize
from sklearn.metrics.pairwise import cosine_similarity
import faiss

# 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 chromaDB
client = chromadb.PersistentClient(path="./fashion_multimodal_db")
collection = client.get_collection(name="fashion_multimodal")

# 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 initialize_faiss_index(collection):
    # 모든 임베딩을 가져와 numpy 배열로 변환
    all_data = collection.get(include=['embeddings', 'metadatas'])
    all_embeddings = np.array(all_data['embeddings']).astype('float32')
    all_metadatas = all_data['metadatas']
    
    # faiss 인덱스 생성 및 임베딩 추가
    dimension = all_embeddings.shape[1]
    index = faiss.IndexFlatIP(dimension)  # 코사인 유사도를 사용하려면 IndexFlatIP를 사용
    index.add(all_embeddings)
    
    return index, all_metadatas

faiss_index, all_metadatas = initialize_faiss_index(collection)


# 세그먼트 마스크 기반 임베딩 추출
def get_segmented_embedding(img, final_mask):
    img_array = np.array(img)
    final_mask_array = np.array(final_mask)
    
    # 마스크 적용 (배경을 흰색으로 처리)
    img_array[final_mask_array == 0] = 255
    masked_img = Image.fromarray(img_array)
    
    # 마스크된 이미지로부터 임베딩 추출
    image_tensor = preprocess_val(masked_img).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().flatten()

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().flatten()


def segment_clothing(img, clothes=["Hat", "Upper-clothes", "Skirt", "Pants", "Dress", "Belt", "Left-shoe", "Right-shoe", "Scarf"]):
    segments = segmenter(img)
    mask_list = []
    detected_categories = []
    for s in segments:
        if s['label'] in clothes:
            mask_list.append(s['mask'])
            detected_categories.append(s['label'])

    final_mask = np.zeros_like(np.array(img)[:, :, 0])
    for mask in mask_list:
        current_mask = np.array(mask)
        final_mask = np.maximum(final_mask, current_mask)

    final_mask = Image.fromarray(final_mask.astype(np.uint8) * 255)
    img_with_alpha = img.convert("RGBA")
    img_with_alpha.putalpha(final_mask)

    return img_with_alpha.convert("RGB"), final_mask, detected_categories

#def find_similar_images(query_embedding, collection, top_k=5):
    # ChromaDB에서 가장 유사한 이미지들을 쿼리합니다.
#    results = collection.query(
#        query_embeddings=query_embedding.reshape(1, -1),  # 2D 배열로 변환
#        n_results=top_k,
#        include=['metadatas', 'embeddings']  # 임베딩을 포함하도록 수정
#    )
#    
#    # 메타데이터와 임베딩을 추출합니다.
#    top_metadatas = results['metadatas'][0]  
#    top_embeddings = results['embeddings'][0]  # 임베딩 가져오기
#    
#    # 유클리드 거리 계산
#    distances = euclidean_distances(query_embedding.reshape(1, -1), top_embeddings).flatten()
#
#    structured_results = []
#    for metadata, distance in zip(top_metadatas, distances):
#        structured_results.append({
#            'info': metadata,
#            'similarity': 1 / (1 + distance)  # 거리 기반 유사도 (거리가 작을수록 유사도가 높음)
#        })
    
#    return structured_results

def find_similar_images(query_embedding, faiss_index, all_metadatas, top_k=5):
    query_embedding = query_embedding.astype('float32').reshape(1, -1)  # 차원 조정 및 형변환
    _, indices = faiss_index.search(query_embedding, top_k)  # 유사한 벡터의 인덱스를 반환
    
    # 인덱스를 이용해 메타데이터와 유사도 가져오기
    structured_results = []
    for idx in indices[0]:
        metadata = all_metadatas[idx]
        similarity = cosine_similarity(query_embedding, faiss_index.reconstruct(idx).reshape(1, -1))[0][0]
        structured_results.append({
            'info': metadata,
            'similarity': similarity
        })
    
    return structured_results

#def find_similar_images(query_embedding, collection, top_k=5, batch_size=500):
#    query_embedding = query_embedding.reshape(1, -1)  # 쿼리 임베딩 차원 조정
#
#    # 모든 임베딩과 메타데이터를 한 번에 가져옴
#    all_data = collection.get(include=['embeddings', 'metadatas'])
#    all_embeddings = np.array(all_data['embeddings'])
#    all_metadatas = all_data['metadatas']
#    
#    all_results = []
#
#    # 전체 데이터를 batch_size씩 나누어 처리
#    for start in range(0, len(all_embeddings), batch_size):
#        end = start + batch_size
#        batch_embeddings = all_embeddings[start:end]
#        batch_metadatas = all_metadatas[start:end]
#        
#        # 코사인 유사도 계산
#        similarities = cosine_similarity(query_embedding, batch_embeddings).flatten()
#        
#        # 현재 배치에서 유사도와 메타데이터를 쌍으로 묶어 추가
#        batch_results = [{'info': metadata, 'similarity': similarity} for similarity, metadata in zip(similarities, batch_metadatas)]
#        all_results.extend(batch_results)
#
#    # 전체 결과 중에서 유사도가 높은 순서대로 top_k 개만 선택
#    sorted_results = sorted(all_results, key=lambda x: x['similarity'], reverse=True)[:top_k]
#    
#    return sorted_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:
    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
                segmented_image, final_mask, detected_categories = segment_clothing(query_image)
                st.session_state.segmented_image = segmented_image
                st.session_state.detections = 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")
    st.image(original_image, caption="Original Image", use_column_width=True)
    
    # 세그먼트된 이미지에서 임베딩 추출
    query_embedding = get_segmented_embedding(st.session_state.query_image, st.session_state.segmented_image)
    #query_embedding = get_image_embedding(st.session_state.segmented_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.detections = []
        st.session_state.segmented_image = None