Spaces:
Sleeping
Sleeping
File size: 4,615 Bytes
9ad6761 82ffad3 9ad6761 82ffad3 9ad6761 72bdb77 9ad6761 0551632 9ad6761 0551632 9ad6761 82ffad3 9ad6761 82ffad3 9ad6761 82ffad3 9ad6761 82ffad3 | 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 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 | 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" # Update this path
FLAG_EMBEDDINGS_PATH = "./flag_embeddings_1.npy" # Update this path
# Load precomputed embeddings
flag_embeddings = np.load(FLAG_EMBEDDINGS_PATH, allow_pickle=True).item()
# 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"))
]
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)
# Verify if the file exists before adding it to the result
results = []
for flag_file, similarity in sorted_flags[1:top_n + 1]: # Skip the first one as it's the same flag
flag_path = os.path.join(FLAG_IMAGE_DIR, flag_file)
if os.path.exists(flag_path):
results.append((flag_file, similarity))
else:
print(f"File not found: {flag_file}")
return results
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.")
with gr.Row():
with gr.Column(scale=1):
# Search and input gallery
search_box = gr.Textbox(label="Search Flags", placeholder="Enter country name...")
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]
)
# 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()
|