Spaces:
Sleeping
Sleeping
| import os | |
| import gradio as gr | |
| from PIL import Image | |
| from sklearn.metrics.pairwise import cosine_similarity | |
| from clip_embedding import Clip | |
| from efficientnet_embedding import EfficientNet | |
| from vit_embedding import Vit | |
| from resnet_embedding import Resnet | |
| from dino_embedding import Dino | |
| from histogram_embedding import cosine, get_embedding | |
| from bovw_embedding import Bovw | |
| resnet = Resnet() | |
| vit = Vit() | |
| efficientnet = EfficientNet() | |
| bovw = Bovw() | |
| dino = Dino() | |
| clip = Clip() | |
| def get_image_embedding(image: Image.Image, name): | |
| match name: | |
| case "ResNet": | |
| return resnet.get_embedding(image).cpu().numpy() | |
| case "VIT": | |
| return vit.get_embedding(image).cpu().numpy() | |
| case "EfficientNet": | |
| return efficientnet.get_embedding(image).cpu().numpy() | |
| case "Histogram": | |
| return get_embedding(image) | |
| case "BOVW": | |
| return bovw.get_embedding(image) | |
| case "DINO": | |
| return dino.get_embedding(image).cpu().numpy() | |
| case _: | |
| return clip.get_embedding(image).cpu().numpy() | |
| def compare_images(main_img, compare_imgs, name): | |
| results = [] | |
| if name in ("Histogram", "ResNet", "BOVW"): | |
| main_emb = get_image_embedding(main_img, name) | |
| for img in compare_imgs: | |
| emb = get_image_embedding(img, name) | |
| results.append((img, round(cosine(main_emb, emb) * 100, 2))) | |
| else: | |
| main_embedding = get_image_embedding(main_img, name) | |
| for img in compare_imgs: | |
| emb = get_image_embedding(img, name) | |
| score = cosine_similarity(main_embedding, emb)[0][0] | |
| percentage = round(score * 100, 2) | |
| results.append((img, percentage)) | |
| results.sort(key=lambda x: x[1], reverse=True) | |
| return results | |
| model_list = ["CLIP", "VIT", "EfficientNet", "ResNet", "DINO", "Histogram", "BOVW"] | |
| with gr.Blocks() as demo: | |
| gr.Tab("Image Embedding") | |
| gr.Markdown("# Image Similarity Finder") | |
| gr.Markdown( | |
| "Upload a main image and compare it to others. Results show similarity percentages using embeddings.") | |
| with gr.Row(): | |
| with gr.Column(): | |
| main_image = gr.Image(type="pil", label="Main Image") | |
| compare_images_input = gr.File(file_count="multiple", file_types=["image"], label="Comparison Images") | |
| modelName = gr.Dropdown(model_list, label="Model", value=model_list[0]) | |
| submit_btn = gr.Button("Compare") | |
| with gr.Column(): | |
| gallery = gr.Gallery(label="Similarity Results") | |
| similarity_text = gr.Textbox(label="Similarity Scores") | |
| def process_comparison(main_img, compare_files, name): | |
| compare_imgs = [Image.open(file.name) for file in compare_files] | |
| results = compare_images(main_img, compare_imgs, name) | |
| # Prepare outputs | |
| images = [result[0] for result in results] | |
| scores = [f"Image: {os.path.basename(result[0].filename)} -> Similarity: {result[1]:.2f}%" for result in | |
| results] | |
| return images, "\n".join(scores) | |
| submit_btn.click( | |
| fn=process_comparison, | |
| inputs=[main_image, compare_images_input, modelName], | |
| outputs=[gallery, similarity_text] | |
| ) | |
| demo.launch() | |