Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import pandas as pd | |
| import numpy as np | |
| import torch | |
| from transformers import CLIPModel, CLIPProcessor | |
| # ----------------------------- | |
| # 1. Load model & processor | |
| # ----------------------------- | |
| model_name = "openai/clip-vit-base-patch32" | |
| model = CLIPModel.from_pretrained(model_name) | |
| processor = CLIPProcessor.from_pretrained(model_name) | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| model.to(device) | |
| model.eval() | |
| # ----------------------------- | |
| # 2. Load your saved embeddings | |
| # ----------------------------- | |
| df = pd.read_parquet("animal_embeddings.parquet") | |
| embeddings = df.drop(columns=["label", "index"]).values | |
| labels = df["label"].tolist() | |
| indices = df["index"].tolist() | |
| # Load dataset to retrieve images | |
| from datasets import load_dataset | |
| dataset = load_dataset("mountassir/animals-10")["train"] | |
| sampled_data = dataset.select(indices) | |
| label_names = dataset.features["label"].names | |
| # ----------------------------- | |
| # 3. Helper functions | |
| # ----------------------------- | |
| def embed_image_query(pil_image): | |
| with torch.no_grad(): | |
| inputs = processor(images=pil_image, return_tensors="pt").to(device) | |
| feats = model.get_image_features(**inputs) | |
| feats = feats / feats.norm(dim=-1, keepdim=True) | |
| return feats.squeeze().cpu().numpy() | |
| def embed_text_query(text): | |
| with torch.no_grad(): | |
| inputs = processor(text=[text], return_tensors="pt").to(device) | |
| feats = model.get_text_features(**inputs) | |
| feats = feats / feats.norm(dim=-1, keepdim=True) | |
| return feats.squeeze().cpu().numpy() | |
| from sklearn.metrics.pairwise import cosine_similarity | |
| def get_top_k(query_emb, k=3): | |
| sims = cosine_similarity(query_emb.reshape(1, -1), embeddings)[0] | |
| idxs = np.argsort(sims)[::-1][:k] | |
| return idxs, sims[idxs] | |
| # ----------------------------- | |
| # 4. Gradio functions | |
| # ----------------------------- | |
| def gradio_image_search(image): | |
| query_emb = embed_image_query(image) | |
| idxs, scores = get_top_k(query_emb, 3) | |
| results = [sampled_data[i]["image"] for i in idxs] | |
| return results | |
| def gradio_text_search(text): | |
| query_emb = embed_text_query(text) | |
| idxs, scores = get_top_k(query_emb, 3) | |
| results = [sampled_data[i]["image"] for i in idxs] | |
| return results | |
| # ----------------------------- | |
| # 5. Build Gradio App | |
| # ----------------------------- | |
| with gr.Blocks() as demo: | |
| gr.Markdown(""" | |
| # 🐾 Animal Similarity Finder | |
| Welcome! This app allows you to find animals that look visually similar using image and text embeddings. | |
| How it works | |
| - The model uses **CLIP embeddings** to compare your input with a database of animal images. | |
| - It returns the **Top 3 most similar images** from the Animals-10 dataset. | |
| Image Search | |
| Upload a picture of an animal (dog, cat, spider, butterfly, horse, etc.). | |
| The app will analyze the image and show you the 3 closest matches based on **visual similarity**. | |
| Text Search | |
| Type a description like: | |
| - **"pet"** → finds dogs & cats | |
| - **"bug"** → finds spiders | |
| - **"farm animal"** → finds sheep, cows, horses | |
| - **"bird"** → finds chickens | |
| The model converts your text into an embedding and returns the 3 images most related to your description. | |
| Behind the scenes | |
| - Embeddings generated with **CLIP (ViT-B/32)** | |
| - Similarity is computed using **cosine similarity** | |
| - All embeddings are precomputed for speed | |
| Enjoy exploring the animal dataset! 🐶🐱🐴🦋🕷️ | |
| """) | |
| with gr.Tab("Image Search"): | |
| img_in = gr.Image(type="pil") | |
| img_out = gr.Gallery(label="Top 3 Results", columns=3) | |
| btn1 = gr.Button("Search") | |
| btn1.click(fn=gradio_image_search, inputs=img_in, outputs=img_out) | |
| with gr.Tab("Text Search"): | |
| txt_in = gr.Textbox(label="e.g. 'pet', 'bug', 'farm animal'") | |
| txt_out = gr.Gallery(label="Top 3 Results", columns=3) | |
| btn2 = gr.Button("Search") | |
| btn2.click(fn=gradio_text_search, inputs=txt_in, outputs=txt_out) | |
| demo.launch() | |