Create app.py

app.py

@@ -0,0 +1,95 @@
+import gradio as gr
+import torch
+import clip
+import faiss
+import numpy as np
+from PIL import Image
+import os
+
+# Load CLIP model
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model, preprocess = clip.load("ViT-B/32", device=device)
+
+# Global storage
+image_paths = []
+image_embeddings = None
+faiss_index = None
+
+def build_faiss_index(images):
+    """Build FAISS index from uploaded images"""
+    global image_paths, image_embeddings, faiss_index
+    image_paths = []
+    embeddings = []
+
+    for img in images:
+        image_paths.append(img.name)
+        pil_img = Image.open(img.name).convert("RGB")
+        tensor_img = preprocess(pil_img).unsqueeze(0).to(device)
+
+        with torch.no_grad():
+            emb = model.encode_image(tensor_img)
+            emb /= emb.norm(dim=-1, keepdim=True)
+        embeddings.append(emb.cpu().numpy())
+
+    image_embeddings = np.vstack(embeddings).astype("float32")
+
+    # Build FAISS index
+    d = image_embeddings.shape[1]  # embedding dimension
+    faiss_index = faiss.IndexFlatIP(d)  # cosine similarity (inner product)
+    faiss_index.add(image_embeddings)
+
+    return f"Indexed {len(image_paths)} images."
+
+def search(query, top_k=5):
+    """Search top-k most similar images given a text query"""
+    global image_paths, faiss_index, image_embeddings
+    if faiss_index is None:
+        return "Please upload and index images first.", []
+
+    # Encode query
+    text = clip.tokenize([query]).to(device)
+    with torch.no_grad():
+        text_emb = model.encode_text(text)
+        text_emb /= text_emb.norm(dim=-1, keepdim=True)
+
+    text_emb = text_emb.cpu().numpy().astype("float32")
+
+    # Search FAISS
+    scores, indices = faiss_index.search(text_emb, top_k)
+    results = []
+    for idx, score in zip(indices[0], scores[0]):
+        img = image_paths[idx]
+        results.append((img, float(score)))
+
+    return f"Top {top_k} results for '{query}'", results
+
+def display_results(query, top_k=5):
+    message, results = search(query, top_k)
+    images, scores = [], []
+    for img, score in results:
+        images.append(img)
+        scores.append(f"{score:.3f}")
+    return message, images, scores
+
+with gr.Blocks() as demo:
+    gr.Markdown("## Image Search with CLIP + FAISS 🚀")
+    
+    with gr.Row():
+        img_upload = gr.File(file_types=[".png", ".jpg", ".jpeg"], file_count="multiple")
+        build_btn = gr.Button("Build Index")
+    
+    status = gr.Textbox(label="Status")
+
+    with gr.Row():
+        query = gr.Textbox(label="Search Query")
+        top_k = gr.Slider(1, 20, value=5, step=1, label="Top K Results")
+        search_btn = gr.Button("Search")
+    
+    output_text = gr.Textbox(label="Results")
+    output_gallery = gr.Gallery(label="Ranked Images").style(grid=[5], height="auto")
+    output_scores = gr.Textbox(label="Similarity Scores")
+
+    build_btn.click(fn=build_faiss_index, inputs=[img_upload], outputs=[status])
+    search_btn.click(fn=display_results, inputs=[query, top_k], outputs=[output_text, output_gallery, output_scores])
+
+demo.launch()