Update app.py

app.py

@@ -1,96 +1,98 @@
-import gradio as gr
 import torch
 import clip
 from datasets import load_dataset
 from PIL import Image
-import faiss
+import gradio as gr
+from torchvision import transforms
 import requests
 from io import BytesIO
+import numpy as np
+import faiss
 
+# Set device
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# 1) Load base CLIP model + preprocess
+# Load CLIP model
 model_clip, preprocess = clip.load("ViT-B/32", device=device)
 
-# 2) Load your fine‐tuned weights (state_dict) into model_clip
-state_dict = torch.load("best_model.pt", map_location=device)
-missing, unexpected = model_clip.load_state_dict(state_dict, strict=False)
-print(f"⚠️  Missing keys: {missing}\n⚠️  Unexpected keys: {unexpected}")
+# Load your fine-tuned model weights
+fine_tuned_state_dict = torch.load("best_model.pt", map_location=device)
+model_clip.load_state_dict(fine_tuned_state_dict)
+
 model_clip.eval()
 
-# 3) Build retrieval pool from Flickr30k test split
-dataset = load_dataset("nlphuji/flickr30k", split="test")
+# Load 50 samples from Flickr30k test split
+dataset = load_dataset("nlphuji/flickr30k", split="test[:50]")
 
-images, captions = [], []
-img_embs, txt_embs = [], []
+# Precompute embeddings
+image_embeddings = []
+images = []
+captions = []
+valid_indices = []
 
-print("🔄 Preparing retrieval pool embeddings...")
-for example in dataset:
+print("Extracting embeddings...")
+
+for i, example in enumerate(dataset):
     try:
-        # load & store raw image + caption
         img = Image.open(requests.get(example["image"], stream=True).raw).convert("RGB")
+        img_tensor = preprocess(img).unsqueeze(0).to(device)
+        with torch.no_grad():
+            img_feat = model_clip.encode_image(img_tensor)
+            img_feat /= img_feat.norm(dim=-1, keepdim=True)
+            image_embeddings.append(img_feat.cpu())
         images.append(img)
         captions.append(example["sentence"])
+        valid_indices.append(i)
+    except Exception as e:
+        print(f"Skipping sample {i} due to error: {e}")
+        continue
 
-        # encode image
-        img_t = preprocess(img).unsqueeze(0).to(device)
-        with torch.no_grad():
-            v = model_clip.encode_image(img_t)
-            v /= v.norm(dim=-1, keepdim=True)
-        img_embs.append(v.cpu())
+# Stack image features
+image_embeddings = torch.cat(image_embeddings, dim=0)
 
-        # encode text
-        t = clip.tokenize([example["sentence"]]).to(device)
-        with torch.no_grad():
-            tfeat = model_clip.encode_text(t)
-            tfeat /= tfeat.norm(dim=-1, keepdim=True)
-        txt_embs.append(tfeat.cpu())
-    except:
-        continue
+# Build FAISS index
+image_index = faiss.IndexFlatIP(image_embeddings.shape[1])
+image_index.add(image_embeddings.numpy())
 
-# cat into tensors
-img_embs = torch.cat(img_embs, dim=0)
-txt_embs = torch.cat(txt_embs, dim=0)
+# Search function
+def search_by_text(query):
+    with torch.no_grad():
+        tokens = clip.tokenize([query]).to(device)
+        text_feat = model_clip.encode_text(tokens)
+        text_feat /= text_feat.norm(dim=-1, keepdim=True)
+        text_feat_np = text_feat.cpu().numpy()
 
-# build FAISS indices (Inner‐Product = cosine)
-img_index = faiss.IndexFlatIP(img_embs.shape[1])
-img_index.add(img_embs.numpy())
+    D, I = image_index.search(text_feat_np, 5)
+    results = []
+    for idx in I[0]:
+        img = images[idx]
+        caption = captions[idx]
+        results.append((img, caption))
+    return results
 
-txt_index = faiss.IndexFlatIP(txt_embs.shape[1])
-txt_index.add(txt_embs.numpy())
+# Gradio interface
+def display_results(text_query):
+    results = search_by_text(text_query)
+    output = ""
+    for i, (img, caption) in enumerate(results):
+        output += f"### Result {i+1}\n"
+        output += f"**Caption:** {caption}\n\n"
+        output += f"![img](data:image/png;base64,{image_to_base64(img)})\n\n"
+    return output
 
-# 4) Gradio callbacks
-def image_to_text(inp_img):
-    im = preprocess(inp_img).unsqueeze(0).to(device)
-    with torch.no_grad():
-        v = model_clip.encode_image(im)
-        v /= v.norm(dim=-1, keepdim=True)
-    D, I = txt_index.search(v.cpu().numpy(), 1)
-    score = D[0][0] * 100
-    return f"{captions[I[0][0]]}\n(Match Score: {score:.2f}%)"
-
-def text_to_image(inp_txt):
-    tok = clip.tokenize([inp_txt]).to(device)
-    with torch.no_grad():
-        t = model_clip.encode_text(tok)
-        t /= t.norm(dim=-1, keepdim=True)
-    D, I = img_index.search(t.cpu().numpy(), 1)
-    score = D[0][0] * 100
-    return images[I[0][0]], f"Match Score: {score:.2f}%"
-
-# 5) Gradio UI
-with gr.Blocks() as demo:
-    gr.Markdown("## 🔄 Cross-Modal Retriever (Flickr30k Test Split)\nUpload an image or enter text to retrieve the best match.")
-    
-    with gr.Tab("🖼️ Image → Text"):
-        img_in = gr.Image(type="pil", label="Upload Image")
-        txt_out = gr.Textbox(label="Retrieved Caption")
-        gr.Button("Search Caption").click(image_to_text, img_in, txt_out)
-    
-    with gr.Tab("���� Text → Image"):
-        txt_in = gr.Textbox(label="Enter Text")
-        img_out = gr.Image(label="Retrieved Image")
-        score_out = gr.Textbox(label="Score")
-        gr.Button("Search Image").click(text_to_image, txt_in, [img_out, score_out])
-
-demo.launch()
+# Convert PIL image to base64
+import base64
+from io import BytesIO
+
+def image_to_base64(image):
+    buffer = BytesIO()
+    image.save(buffer, format="PNG")
+    return base64.b64encode(buffer.getvalue()).decode()
+
+iface = gr.Interface(fn=display_results,
+                     inputs=gr.Textbox(lines=2, placeholder="Enter text to search..."),
+                     outputs="markdown",
+                     title="Text-to-Image Retrieval with CLIP",
+                     description="Enter a sentence to retrieve similar images using a fine-tuned CLIP model.")
+
+iface.launch()