JoJosmin's picture
Update app.py
18bf384 verified
raw
history blame
7.86 kB
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_202410")
#collection = client.get_collection(name="clothes_items_ver3")
collection = client.get_collection(name="clothes")
# 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']
st.write(f"Embeddings found: {len(all_data['embeddings'])}")
st.write(f"Metadatas found: {len(all_data['metadatas'])}")
st.write(f"Documents found: {len(all_data['documents'])}")
# ๋ฒกํ„ฐ ID๋ฅผ ํ•จ๊ป˜ ๊ฐ€์ ธ์™€์„œ ํ™•์ธ
vector_data = collection.get(include=['embeddings', 'metadatas', 'ids'])
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
# ์„ธ์…˜ ์ƒํƒœ ์ดˆ๊ธฐํ™”
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