Upload 4 files

app.py

@@ -0,0 +1,83 @@
+import gradio as gr
+import torch
+import open_clip
+import faiss
+import numpy as np
+import json
+from PIL import Image
+from transformers import BlipProcessor, BlipForConditionalGeneration
+
+with open("image_paths.json", "r") as f:
+    image_paths = json.load(f)
+
+image_embeddings = np.load("image_embeddings.npy")
+
+model, _, preprocess = open_clip.create_model_and_transforms('ViT-B-32', pretrained='laion2b_s34b_b79k')
+tokenizer = open_clip.get_tokenizer('ViT-B-32')
+
+d = image_embeddings.shape[1]
+index = faiss.IndexFlatIP(d)
+index.add(image_embeddings)
+
+processor_cap = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
+model_cap = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
+
+def search_text(query, top_k=3):
+    with torch.no_grad():
+        tokenized = tokenizer([query])
+        text_embed = model.encode_text(tokenized)
+        text_embed = text_embed / text_embed.norm(dim=-1, keepdim=True)
+        text_np = text_embed.cpu().numpy()
+    _, I = index.search(text_np, top_k)
+    return [image_paths[i] for i in I[0]]
+
+def search_image(image, top_k=3):
+    image_tensor = preprocess(image).unsqueeze(0)
+    with torch.no_grad():
+        image_embed = model.encode_image(image_tensor)
+        image_embed = image_embed / image_embed.norm(dim=-1, keepdim=True)
+        image_np = image_embed.cpu().numpy()
+    _, I = index.search(image_np, top_k)
+    return [image_paths[i] for i in I[0]]
+
+def generate_captions(paths):
+    imgs = [Image.open(p).convert("RGB") for p in paths]
+    inputs = processor_cap(images=imgs, return_tensors="pt")
+    out = model_cap.generate(**inputs)
+    captions = [processor_cap.decode(out[i], skip_special_tokens=True) for i in range(len(imgs))]
+    return captions
+
+def predict(input_text, input_image):
+    if input_text:
+        paths = search_text(input_text)
+    elif input_image:
+        paths = search_image(input_image)
+    else:
+        return None, None
+
+    captions = generate_captions(paths)
+    return [(Image.open(p), c) for p, c in zip(paths, captions)]
+
+with gr.Blocks() as demo:
+    gr.Markdown("## Animal Image Search 🐈 🏞️ 🔎 ")
+
+    with gr.Row():
+        text_input = gr.Textbox(label="Type your query")
+        image_input = gr.Image(type="pil", label="Or upload an image")
+    
+    gr.Examples(
+    examples=[
+        "examples/sample1.jpg",
+        "examples/sample2.jpg",
+        "examples/sample3.jpg"
+    ],
+    inputs=[image_input],
+    label="Choose a sample image"
+    )
+
+    output = gr.Gallery(label="Top Matches with Captions").style(grid=[1], height="auto")
+
+    submit = gr.Button("Search")
+    submit.click(predict, inputs=[text_input, image_input], outputs=output)
+
+demo.launch()

image_embeddings.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8a48867e7ee345749a9e16d79ff4d78a2d01c2fbc2b8cd8705aa5dfc92171764
+size 11059328

requirements.txt

@@ -0,0 +1,6 @@
+gradio
+torch
+torchvision
+open-clip-torch
+faiss-cpu
+transformers