Spaces:
Runtime error
Runtime error
File size: 2,268 Bytes
dbb1308 77403d5 cf029f9 77403d5 3bc904e dbb1308 d21fc46 77403d5 b33e6dd dbb1308 77403d5 28de0db dbb1308 cf029f9 77403d5 cf029f9 b33e6dd cf029f9 b33e6dd 28de0db cf029f9 dbb1308 77403d5 dbb1308 b33e6dd dbb1308 b33e6dd cf029f9 b33e6dd 28de0db dbb1308 cf029f9 |
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 |
import gradio as gr
from upstash_vector import Index
from datasets import load_dataset
from transformers import AutoFeatureExtractor, AutoModel
index = Index.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 BounharAbdelaziz/Face-Aging-Dataset dataset using Google's VIT model. Powered by [Upstash Vector](https://upstash.com) where all the image embeddings are stored 🚀
"""
)
with gr.Tab("Basic"):
with gr.Row():
with gr.Column(scale=1):
input_image = gr.Image(type="pil")
with gr.Column(scale=3):
output_image = gr.Gallery(height=800)
@input_image.upload(inputs=input_image, outputs=output_image)
def find_similar_faces(image):
inputs = extractor(images=image, return_tensors="pt")
outputs = model(**inputs)
embed = outputs.last_hidden_state[0][0]
result = index.query(vector=embed.tolist(), top_k=4)
return [dataset["train"][int(vector.id)]["image"] for vector in result]
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")
with gr.Column(scale=3):
adv_output_image = gr.Gallery(height=1000)
@adv_input_image.upload(
inputs=[adv_input_image, adv_image_count], outputs=[adv_output_image]
)
def find_similar_faces(image, count):
inputs = extractor(images=image, return_tensors="pt")
outputs = model(**inputs)
embed = outputs.last_hidden_state[0][0]
result = index.query(vector=embed.tolist(), top_k=max(1, min(19, count)))
return [dataset["train"][int(vector.id)]["image"] for vector in result]
if __name__ == "__main__":
demo.launch(debug=True)
|