Spaces:
Runtime error
Runtime error
| 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 | |