Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import os | |
| import numpy as np | |
| import torch | |
| import clip | |
| from PIL import Image | |
| from sklearn.metrics.pairwise import cosine_similarity | |
| # Load CLIP model | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| model, preprocess = clip.load("ViT-B/32", device=device) | |
| # Configuration | |
| FLAG_IMAGE_DIR = "./named_flags" | |
| # Function to search flags with specific queries using CLIP | |
| def search_by_query(query, top_n=10): | |
| """Search flags based on a text query using CLIP.""" | |
| # Encode the text query | |
| with torch.no_grad(): | |
| text_embedding = model.encode_text(clip.tokenize([query]).to(device)) | |
| # Compare the query embedding with all flag embeddings | |
| similarities = {} | |
| for flag, embedding in flag_embeddings.items(): | |
| similarity = cosine_similarity(text_embedding.cpu().numpy(), embedding)[0][0] | |
| similarities[flag] = similarity | |
| # Sort and return the top_n results | |
| sorted_flags = sorted(similarities.items(), key=lambda x: x[1], reverse=True) | |
| results = [] | |
| for flag_file, similarity in sorted_flags[:top_n]: | |
| flag_path = os.path.join(FLAG_IMAGE_DIR, flag_file) | |
| if os.path.exists(flag_path): | |
| results.append((flag_path, f"{get_country_name(flag_file)} (Similarity: {similarity:.3f})")) | |
| else: | |
| print(f"File not found: {flag_file}") | |
| return results | |
| # Get all image paths | |
| image_paths = [ | |
| os.path.join(FLAG_IMAGE_DIR, img) | |
| for img in os.listdir(FLAG_IMAGE_DIR) | |
| if img.endswith((".png", ".jpg", ".jpeg")) | |
| ] | |
| # Load precomputed embeddings | |
| FLAG_EMBEDDINGS_PATH = "./flag_embeddings_1.npy" | |
| flag_embeddings = np.load(FLAG_EMBEDDINGS_PATH, allow_pickle=True).item() | |
| def get_country_name(image_filename): | |
| """Extract country name from image filename.""" | |
| return os.path.splitext(os.path.basename(image_filename))[0].upper() | |
| def get_image_embedding(image_path): | |
| """Get embedding for an input image.""" | |
| image = Image.open(image_path).convert("RGB") | |
| image_input = preprocess(image).unsqueeze(0).to(device) | |
| with torch.no_grad(): | |
| embedding = model.encode_image(image_input) | |
| return embedding.cpu().numpy() | |
| def find_similar_flags(image_path, top_n=10): | |
| """Find similar flags based on cosine similarity.""" | |
| query_embedding = get_image_embedding(image_path) | |
| similarities = {} | |
| for flag, embedding in flag_embeddings.items(): | |
| similarity = cosine_similarity(query_embedding, embedding)[0][0] | |
| similarities[flag] = similarity | |
| sorted_flags = sorted(similarities.items(), key=lambda x: x[1], reverse=True) | |
| return sorted_flags[1:top_n + 1] # Skip the first one as it's the same flag | |
| def search_flags(query): | |
| """Search flags based on country name.""" | |
| if not query: | |
| return image_paths | |
| return [img for img in image_paths if query.lower() in get_country_name(img).lower()] | |
| def analyze_and_display(selected_flag): | |
| """Main function to analyze flag similarity and prepare display.""" | |
| try: | |
| if selected_flag is None: | |
| return None | |
| similar_flags = find_similar_flags(selected_flag) | |
| output_images = [] | |
| for flag_file, similarity in similar_flags: | |
| flag_path = os.path.join(FLAG_IMAGE_DIR, flag_file) | |
| country_name = get_country_name(flag_file) | |
| output_images.append((flag_path, f"{country_name} (Similarity: {similarity:.3f})")) | |
| return output_images | |
| except Exception as e: | |
| return gr.Error(f"Error processing image: {str(e)}") | |
| # Create Gradio interface | |
| with gr.Blocks() as demo: | |
| gr.Markdown("# Flag Similarity Analysis") | |
| gr.Markdown("Select a flag from the gallery to find similar flags based on visual features or search using text queries.") | |
| with gr.Tabs(): | |
| with gr.Tab("Similarity Search"): | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| # Search and input gallery | |
| search_box = gr.Textbox(label="Search Flags", placeholder="Enter country name...") | |
| #query_box = gr.Textbox(label="Search by Query", placeholder="e.g., 'crescent in the center'") | |
| input_gallery = gr.Gallery( | |
| label="Available Flags", | |
| show_label=True, | |
| elem_id="gallery", | |
| columns=4, | |
| height="auto" | |
| ) | |
| with gr.Column(scale=1): | |
| # Output gallery | |
| output_gallery = gr.Gallery( | |
| label="Similar Flags", | |
| show_label=True, | |
| elem_id="output", | |
| columns=2, | |
| height="auto" | |
| ) | |
| # Event handlers | |
| def update_gallery(query): | |
| matching_flags = search_flags(query) | |
| return [(path, get_country_name(path)) for path in matching_flags] | |
| def on_select(evt: gr.SelectData, gallery): | |
| """Handle flag selection from gallery""" | |
| selected_flag_path = gallery[evt.index][0] | |
| return analyze_and_display(selected_flag_path) | |
| # Connect event handlers | |
| search_box.change( | |
| update_gallery, | |
| inputs=[search_box], | |
| outputs=[input_gallery] | |
| ) | |
| input_gallery.select( | |
| on_select, | |
| inputs=[input_gallery], | |
| outputs=[output_gallery] | |
| ) | |
| with gr.Tab("Advanced Search"): | |
| gr.Markdown("### Search Flags with Nuanced Queries") | |
| nuanced_query_box = gr.Textbox(label="Enter Advanced Query", placeholder="e.g., 'Find flags with crescent' or 'flags with animals'") | |
| advanced_output_gallery = gr.Gallery( | |
| label="Matching Flags", | |
| show_label=True, | |
| elem_id="advanced_output", | |
| columns=3, | |
| height="auto" | |
| ) | |
| def advanced_search(query): | |
| return search_by_query(query) | |
| nuanced_query_box.change( | |
| advanced_search, | |
| inputs=[nuanced_query_box], | |
| outputs=[advanced_output_gallery] | |
| ) | |
| # Initialize gallery with all flags | |
| def init_gallery(): | |
| return [(path, get_country_name(path)) for path in image_paths] | |
| demo.load(init_gallery, outputs=[input_gallery]) | |
| # Launch the app | |
| demo.launch() |