🌘w🌖

app.py

@@ -4,11 +4,12 @@ import torch
 from transformers import AutoProcessor, AutoModelForZeroShotImageClassification
 from datasets import load_dataset
 
+
 dataset = load_dataset("not-lain/embedded-pokemon", split="train")
 dataset = dataset.add_faiss_index("embeddings")
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
 
+device = "cuda" if torch.cuda.is_available() else "cpu"
 processor = AutoProcessor.from_pretrained("openai/clip-vit-large-patch14")
 model = AutoModelForZeroShotImageClassification.from_pretrained(
     "openai/clip-vit-large-patch14", device_map=device
@@ -31,17 +32,23 @@ def search(query: str, k: int = 4):
         img_emb,  # compare our new embedded query with the dataset embeddings
         k=k,  # get only top k results
     )
-    images = retrieved_examples["image"]
-    # labels = {}
-    # for i in range(k):
-    #     labels[retrieved_examples["text"][k-i]] = scores[k-i]
 
-    return images #, labels
+    # return as image, caption pairs
+    out = []
+    for i in range(k):
+        out.append([retrieved_examples["image"][i], retrieved_examples["text"][i]])
+
+    return out 
+
 
-demo = gr.Interface(search, inputs="image", outputs=["gallery"
-                                                     #, "label"
-                                                     ],
-                                                     examples=["./charmander.jpg"],
+demo = gr.Interface(
+    search,
+    inputs="image",
+    outputs=[
+        "gallery"
+        # , "label"
+    ],
+    examples=["./charmander.jpg"],
 )
 
 demo.launch(debug=True)