Spaces:
Runtime error
Runtime error
File size: 3,657 Bytes
27b3217 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 |
import pandas as pd
from utils.vector_database import search_in_milvus, fashionclip_collection, fashionsiglip_collection
from utils.embedding_generation import generate_query_embedding
from utils.load_models import fclip_model, fclip_processor
from utils.load_models import siglip_model, siglip_preprocess_val, siglip_tokenizer
# Function to dynamically select the Milvus collection and search field
def get_milvus_collection_and_field(model_type, embedding_type):
# Define mapping of model and embedding types to collections and fields
if model_type == "fashionCLIP":
collection = fashionclip_collection
if embedding_type == "text":
search_field = "text_embedding"
elif embedding_type == "image":
search_field = "image_embedding"
elif embedding_type == "average":
search_field = "avg_embedding"
elif embedding_type == "weighted average":
search_field = "weighted_avg_embedding"
elif model_type == "fashionSigLIP":
collection = fashionsiglip_collection
if embedding_type == "text":
search_field = "text_embedding"
elif embedding_type == "image":
search_field = "image_embedding"
elif embedding_type == "average":
search_field = "avg_embedding"
elif embedding_type == "weighted average":
search_field = "weighted_avg_embedding"
else:
raise ValueError("Invalid model type. Choose 'fashionCLIP' or 'fashionSigLIP'.")
return collection, search_field
# Function to handle the complete search flow
def search(query, query_type, model_type, embedding_type):
# Step 1: Generate the query embedding based on the user input and model type
if model_type == "fashionCLIP":
query_embedding = generate_query_embedding(query, query_type, fclip_model, fclip_processor, fclip_processor, "fashionCLIP")
elif model_type == "fashionSigLIP":
query_embedding = generate_query_embedding(query, query_type, siglip_model, siglip_preprocess_val, siglip_tokenizer, "fashionSigLIP")
# Step 2: Get the appropriate Milvus collection and search field
collection, search_field = get_milvus_collection_and_field(model_type, embedding_type)
# Step 3: Perform search in Milvus using the query embedding
search_results = search_in_milvus(collection, search_field, query_embedding, top_k=10)
# Step 4: Extract images, similarity scores, and metadata from the search results
images = [result['image'] for result in search_results]
scores = [result['similarity_score'] for result in search_results]
metadata = [result['metadata'] for result in search_results]
return images, scores, metadata
# Function to run the search and get results for both models
def run_search(query_type, embedding_type, query_input_text, query_input_image):
if query_type == "text":
query = query_input_text
else:
query = query_input_image
# Perform search for FashionCLIP
fclip_images, fclip_scores, fclip_metadata = search(query, query_type, "fashionCLIP", embedding_type)
# Perform search for MARGO-FashionSigLip
siglip_images, siglip_scores, siglip_metadata = search(query, query_type, "fashionSigLIP", embedding_type)
# Convert scores and metadata into a pandas DataFrame for each model
fclip_results_df = pd.DataFrame({
"Score": fclip_scores,
"Metadata": fclip_metadata,
})
siglip_results_df = pd.DataFrame({
"Score": siglip_scores,
"Metadata": siglip_metadata,
})
return fclip_images, fclip_results_df, siglip_images, siglip_results_df
|