Spaces:
Running
Running
File size: 2,987 Bytes
dbb1308 83305ef cf029f9 83305ef 77403d5 83305ef 77403d5 3bc904e dbb1308 d21fc46 77403d5 83305ef dbb1308 77403d5 b2f105a dbb1308 cf029f9 83305ef b33e6dd cf029f9 83305ef 28de0db cf029f9 83305ef b2f105a 83305ef dbb1308 77403d5 dbb1308 b2f105a dbb1308 b2f105a dbb1308 b2f105a 83305ef b33e6dd cf029f9 83305ef 28de0db dbb1308 b2f105a dbb1308 b2f105a |
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 |
import gradio as gr
from upstash_vector import AsyncIndex
from transformers import AutoFeatureExtractor, AutoModel
from datasets import load_dataset
index = AsyncIndex.from_env()
model_ckpt = "google/vit-base-patch16-224-in21k"
extractor = AutoFeatureExtractor.from_pretrained(model_ckpt)
model = AutoModel.from_pretrained(model_ckpt)
hidden_dim = model.config.hidden_size
dataset = load_dataset("BounharAbdelaziz/Face-Aging-Dataset")
with gr.Blocks() as demo:
gr.Markdown(
"""
# Find Your Twins
Upload your face and find the most similar people from [Face Aging Dataset](https://huggingface.co/datasets/BounharAbdelaziz/Face-Aging-Dataset) using Google's [VIT](https://huggingface.co/google/vit-base-patch16-224-in21k) model. The task of finding most similar vectors is powered by [Upstash Vector](https://upstash.com) 🚀. Check our blog post *here*.
"""
)
with gr.Tab("Basic"):
with gr.Row():
with gr.Column(scale=1):
input_image = gr.Image(type="pil")
with gr.Column(scale=2):
output_image = gr.Gallery(height=450)
@input_image.change(inputs=input_image, outputs=output_image)
async def find_similar_faces(image):
if image is None:
return None
inputs = extractor(images=image, return_tensors="pt")
outputs = model(**inputs)
embed = outputs.last_hidden_state[0][0]
result = await index.query(vector=embed.tolist(), top_k=4)
return [dataset["train"][int(vector.id)]["image"] for vector in result]
gr.Examples(
examples=[dataset["train"][6]["image"], dataset["train"][7]["image"], dataset["train"][8]["image"]],
inputs=input_image,
outputs=output_image,
fn=find_similar_faces,
cache_examples=False,
)
with gr.Tab("Advanced"):
with gr.Row():
with gr.Column(scale=1):
adv_input_image = gr.Image(type="pil")
adv_image_count = gr.Number(9, label="Image Count")
adv_button = gr.Button("Submit")
with gr.Column(scale=2):
adv_output_image = gr.Gallery()
@adv_button.click(inputs=[adv_input_image, adv_image_count], outputs=[adv_output_image])
async def find_similar_faces(image, count):
inputs = extractor(images=image, return_tensors="pt")
outputs = model(**inputs)
embed = outputs.last_hidden_state[0][0]
result = await index.query(
vector=embed.tolist(), top_k=max(1, min(19, count))
)
return [dataset["train"][int(vector.id)]["image"] for vector in result]
adv_input_image.upload(fn=find_similar_faces, inputs=[adv_input_image, adv_image_count],
outputs=[adv_output_image])
if __name__ == "__main__":
demo.launch(debug=True, share=True)
|