Spaces:
Sleeping
Sleeping
File size: 3,367 Bytes
0ea07d3 816eb0d 42cf5a6 1f8d0a1 42cf5a6 0b97a05 816eb0d dbc62a6 816eb0d bf42d19 816eb0d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 | import os
# Redirect all Hugging Face / Torch caches to ephemeral storage
os.environ["HF_HOME"] = "/tmp/hf"
os.environ["TRANSFORMERS_CACHE"] = "/tmp/hf"
os.environ["TORCH_HOME"] = "/tmp/torch"
os.environ["XDG_CACHE_HOME"] = "/tmp"
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 🐈 🏞️ 🔎 ")
gr.Markdown("Search database of 5400 animal images using text or image queries.")
gr.Markdown("Results come with automatic captions.")
gr.Markdown("For list of animal categories (90 animals x 60 images = 5400 images) visit [GitHub](https://github.com/TensorCruncher/animal-image-search/blob/main/animals.txt).")
gr.Markdown("Note: Make sure only one input field (text or image) is populated before searching")
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")
submit = gr.Button("Search", variant="primary")
submit.click(predict, inputs=[text_input, image_input], outputs=output)
demo.launch()
|