adding labels

app.py

@@ -8,7 +8,7 @@ from transformers import AutoModel, AutoFeatureExtractor
 seed = 42
 
 # Only runs once when the script is first run.
-with open("index_768.pickle", "rb") as handle:
+with open("index_768_cosine.pickle", "rb") as handle:
     index = pickle.load(handle)
 
 # Load model for computing embeddings.
@@ -26,21 +26,31 @@ def query(image, top_k=4):
     embedding = model_output.pooler_output.detach()
     results = index.query(embedding, k=top_k)
     inx = results[0][0].tolist()
+    logits = results[1][0].tolist()
     images = ds.select(inx)["image"]
     captions = ds.select(inx)["name"]
     images_with_captions = [(i, c) for i, c in zip(images,captions)]
-    return images_with_captions
+    labels_with_probs = dict(zip(captions,logits))
+    labels_with_probs = {k: 1- v for k, v in labels_with_probs.items()}
+    return images_with_captions, labels_with_probs
 
 
-title = "Find my Butterfly 🦋"
-description = "Use this Space to find your butterfly, based on the [iNaturalist butterfly dataset](https://huggingface.co/datasets/huggan/inat_butterflies_top10k)!"
 
 
-gr.Interface(
-    query,
-    inputs=[gr.Image(type="pil")],
-    outputs=gr.Gallery().style(grid=[2], height="auto"),
-    title=title,
-    description=description,
-    examples=[["elton.jpg"],["ken.jpg"],["gaga.jpg"],["taylor.jpg"]],
-).launch()
+
+with gr.Blocks() as demo:
+	gr.Markdown("# Find my Butterfly 🦋")
+	gr.Markdown("## Use this Space to find your butterfly, based on the [iNaturalist butterfly dataset](https://huggingface.co/datasets/huggan/inat_butterflies_top10k)!")
+	examples=["elton.jpg", "ken.jpg", "gaga.jpg", "taylor.jpg"]
+	with gr.Row():
+		with gr.Column(min_width= 900):
+			inputs = gr.Image(shape=(800, 1600))
+			btn = gr.Button("Find my butterfly!")
+		with gr.Column():
+			outputs=gr.Gallery().style(grid=[2], height="auto")
+			labels = gr.Label()
+
+	btn.click(query, inputs, [outputs, labels])
+	
+demo.launch()
+