sasha's picture
sasha HF staff
Update app.py
856037e
import pickle
import gradio as gr
from datasets import load_dataset
from transformers import AutoModel, AutoFeatureExtractor
import wikipedia
import pynndescent
import numpy as np
# Only runs once when the script is first run.
with open("slugs_1024_new.pickle", "rb") as handle:
index = pickle.load(handle)
'''
embs= np.load('slugs_embeddings_1024k.npy', 'r')
index = pynndescent.NNDescent(embs, metric="cosine")
index.prepare()
'''
# Load model for computing embeddings.
feature_extractor = AutoFeatureExtractor.from_pretrained("sasha/autotrain-sea-slug-similarity-2498977005")
model = AutoModel.from_pretrained("sasha/autotrain-sea-slug-similarity-2498977005")
# Candidate images.
dataset = load_dataset("sasha/australian_sea_slugs")
ds = dataset["train"]
def query(image, top_k=4):
inputs = feature_extractor(image, return_tensors="pt")
model_output = model(**inputs)
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)["label"]
images_with_captions = [(i, c) for i, c in zip(images,captions)]
labels_with_probs = dict(zip(captions,logits))
labels_with_probs = {k: 1- v for k, v in labels_with_probs.items()}
try:
description = wikipedia.summary(captions[0], sentences = 1)
description = "### " + description
url = wikipedia.page(captions[0]).url
url = " You can learn more about your slug [here](" + str(url) + ")!"
description = description + url
except:
description = "### Sea slugs, or Nudibranchs, are marine invertebrates that often live in reefs underwater. They have an enormous variation in body shape, color, and size."
url = "https://en.wikipedia.org/wiki/Sea_slug"
url = " You can learn more about sea slugs [here](" + str(url) + ")!"
description = description + url
return images_with_captions, labels_with_probs, description
with gr.Blocks() as demo:
gr.Markdown("# Which Sea Slug Am I ? 🐌")
gr.Markdown("## Use this Space to find your sea slug, based on the [Nudibranchs of the Sunshine Coast Australia dataset](https://huggingface.co/datasets/sasha/australian_sea_slugs)!")
with gr.Row():
with gr.Column(min_width= 900):
inputs = gr.Image(shape=(800, 1600))
btn = gr.Button("Find my sea slug 🐌!")
description = gr.Markdown()
with gr.Column():
outputs=gr.Gallery().style(grid=[2], height="auto")
labels = gr.Label()
gr.Markdown("### Image Examples")
gr.Examples(
examples=["elton.jpg", "ken.jpg", "gaga.jpg", "taylor.jpg"],
inputs=inputs,
outputs=[outputs,labels],
fn=query,
cache_examples=True,
)
btn.click(query, inputs, [outputs, labels, description])
demo.launch()