app.py

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
from inference_sdk import InferenceHTTPClient
import matplotlib.pyplot as plt
import base64

# Load model and tokenizer
@st.cache_resource
def load_model():
    model, preprocess_val, tokenizer = open_clip.create_model_and_transforms('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

model, preprocess_val, tokenizer, device = load_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
@st.cache_data
def download_and_process_image(image_url):
    try:
        response = requests.get(image_url)
        response.raise_for_status()  # Raises an HTTPError for bad responses
        image = Image.open(BytesIO(response.content))
        
        # Convert image to RGB mode if it's in RGBA mode
        if image.mode == 'RGBA':
            image = image.convert('RGB')
        
        return image
    except requests.RequestException as e:
        st.error(f"Error downloading image: {e}")
        return None
    except Exception as e:
        st.error(f"Error processing image: {e}")
        return None

def get_image_embedding(image):
    image_tensor = preprocess_val(image).unsqueeze(0).to(device)
    with torch.no_grad():
        image_features = model.encode_image(image_tensor)
        image_features /= image_features.norm(dim=-1, keepdim=True)
    return image_features.cpu().numpy()

def setup_roboflow_client(api_key):
    return InferenceHTTPClient(
        api_url="https://outline.roboflow.com",
        api_key=api_key
    )

def segment_image(image_path, client):
    try:
        # 이미지 파일 읽기
        with open(image_path, "rb") as image_file:
            image_data = image_file.read()
        
        # 이미지를 base64로 인코딩
        encoded_image = base64.b64encode(image_data).decode('utf-8')
        
        # 원본 이미지 로드
        image = cv2.imread(image_path)
        image = cv2.resize(image, (800, 600))
        mask = np.zeros(image.shape, dtype=np.uint8)
        
        # Roboflow API 호출
        results = client.infer(encoded_image, model_id="closet/1")
        
        # 결과가 이미 딕셔너리인 경우 JSON 파싱 단계 제거
        if isinstance(results, dict):
            predictions = results.get('predictions', [])
        else:
            # 문자열인 경우에만 JSON 파싱
            predictions = json.loads(results).get('predictions', [])
        
        if predictions:
            for prediction in predictions:
                points = prediction['points']
                pts = np.array([[p['x'], p['y']] for p in points], np.int32)
                scale_x = image.shape[1] / results['image']['width']
                scale_y = image.shape[0] / results['image']['height']
                pts = pts * [scale_x, scale_y]
                pts = pts.astype(np.int32)
                pts = pts.reshape((-1, 1, 2))
                cv2.fillPoly(mask, [pts], color=(255, 255, 255))  # White mask
            
            segmented_image = cv2.bitwise_and(image, mask)
        else:
            st.warning("No predictions found in the image. Returning original image.")
            segmented_image = image
        
        return Image.fromarray(cv2.cvtColor(segmented_image, cv2.COLOR_BGR2RGB))
    except Exception as e:
        st.error(f"Error in segmentation: {str(e)}")
        # 원본 이미지를 다시 읽어 반환
        return Image.open(image_path)

@st.cache_data
def process_database_cached(data):
    database_embeddings = []
    database_info = []
    for item in data:
        image_url = item['이미지 링크'][0]
        product_id = item.get('\ufeff상품 ID') or item.get('상품 ID')
        
        image = download_and_process_image(image_url)
        if image is None:
            continue
        
        # Save the image temporarily
        temp_path = f"temp_{product_id}.jpg"
        image.save(temp_path, 'JPEG')
        
        database_info.append({
            'id': product_id,
            'category': item['카테고리'],
            'brand': item['브랜드명'],
            'name': item['제품명'],
            'price': item['정가'],
            'discount': item['할인율'],
            'image_url': image_url,
            'temp_path': temp_path
        })
    
    return database_info

def process_database(client, data):
    database_info = process_database_cached(data)
    database_embeddings = []
    
    for item in database_info:
        segmented_image = segment_image(item['temp_path'], client)
        embedding = get_image_embedding(segmented_image)
        database_embeddings.append(embedding)
    
    return np.vstack(database_embeddings), database_info

# Streamlit app
st.title("Fashion Search App with Segmentation")

# API Key input
api_key = st.text_input("Enter your Roboflow API Key", type="password")

if api_key:
    CLIENT = setup_roboflow_client(api_key)
    
    # Initialize database_embeddings and database_info
    database_embeddings, database_info = process_database(CLIENT, data)

    uploaded_file = st.file_uploader("Choose an image...", type="jpg")
    if uploaded_file is not None:
        image = Image.open(uploaded_file)
        st.image(image, caption='Uploaded Image', use_column_width=True)
        
        if st.button('Find Similar Items'):
            with st.spinner('Processing...'):
                # Save uploaded image temporarily
                temp_path = "temp_upload.jpg"
                image.save(temp_path)
                
                # Segment the uploaded image
                segmented_image = segment_image(temp_path, CLIENT)
                st.image(segmented_image, caption='Segmented Image', use_column_width=True)
                
                # Get embedding for segmented image
                query_embedding = get_image_embedding(segmented_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("Please enter your Roboflow API Key to use the app.")