Update app.py

app.py

@@ -4,6 +4,13 @@ import random
 import numpy as np
 from transformers import CLIPProcessor, CLIPModel
 from os import environ
+import clip
+import pickle
+import requests
+import torch
+
+is_gpu = False 
+device = CUDA(0) if is_gpu else "cpu"
 
 # Load the pre-trained model and processor
 model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
@@ -33,10 +40,61 @@ get_caption = gr.load("ryaalbr/caption", src="spaces", hf_token=environ["api_key
 def generate_text(image, model_name):
     return get_caption(image, model_name)
 
-get_images = gr.load("ryaalbr/ImageSearch", src="spaces", hf_token=environ["api_key"])
-def search_images(text):
-    return get_images(text, api_name="images")
-    
+# get_images = gr.load("ryaalbr/ImageSearch", src="spaces", hf_token=environ["api_key"])
+# def search_images(text):
+#     return get_images(text, api_name="images")
+
+emb_filename = 'unsplash-25k-photos-embeddings-indexes.pkl'
+with open(emb_filename, 'rb') as emb:
+        id2url, img_names, img_emb = pickle.load(emb)
+
+orig_clip_model, orig_clip_processor = clip.load("ViT-B/32", device=device, jit=False)
+
+def search(search_query):
+
+    with torch.no_grad():
+        # Encode and normalize the description using CLIP
+        text_encoded = orig_clip_model.encode_text(clip.tokenize(search_query))
+        text_encoded /= text_encoded.norm(dim=-1, keepdim=True)
+
+
+    # Retrieve the description vector
+    text_features = text_encoded.cpu().numpy()
+
+    # Compute the similarity between the descrption and each photo using the Cosine similarity
+    similarities = (text_features @ img_emb.T).squeeze(0)
+
+    # Sort the photos by their similarity score
+    best_photos = similarities.argsort()[::-1]
+    best_photos = best_photos[:15]
+    #best_photos = sorted(zip(similarities, range(img_emb.shape[0])), key=lambda x: x[0], reverse=True)
+
+    best_photo_ids = img_names[best_photos]
+
+    imgs = []
+
+    # Iterate over the top 5 results
+    for id in best_photo_ids:
+
+        id, _ = id.split('.')
+        url = id2url.get(id, "")
+        if url == "": continue
+
+        img = url  + "?h=512"
+       # r = requests.get(url + "?w=512", stream=True)
+       # img = Image.open(r.raw)
+        #credits = f'Photo by <a href="https://unsplash.com/@{photo["photographer_username"]}?utm_source=NaturalLanguageImageSearch&utm_medium=referral">{photo["photographer_first_name"]} {photo["photographer_last_name"]}</a> on <a href="https://unsplash.com/?utm_source=NaturalLanguageImageSearch&utm_medium=referral">Unsplash</a>'
+        imgs.append(img)
+        #display(HTML(f'Photo by <a href="https://unsplash.com/@{photo["photographer_username"]}?utm_source=NaturalLanguageImageSearch&utm_medium=referral">{photo["photographer_first_name"]} {photo["photographer_last_name"]}</a> on <a href="https://unsplash.com/?utm_source=NaturalLanguageImageSearch&utm_medium=referral">Unsplash</a>'))
+
+        if len(imgs) == 5: break
+
+    return imgs
+
+
+
+
+
 with gr.Blocks() as demo:
 
     with gr.Tab("Zero-Shot Classification"):
@@ -87,6 +145,6 @@ with gr.Blocks() as demo:
             desc = gr.Textbox(show_label=False, placeholder="Enter description").style(container=False)
             search_btn = gr.Button("Find Images").style(full_width=False)
         gallery = gr.Gallery(show_label=False).style(grid=(2,2,3,5))
-        search_btn.click(search_images,inputs=desc, outputs=gallery)
+        search_btn.click(search,inputs=desc, outputs=gallery)
 
 demo.launch()