Spaces:
Runtime error
Runtime error
File size: 2,434 Bytes
dbb1308 d3d5d4c 77403d5 cf029f9 77403d5 cf029f9 77403d5 cf029f9 33de0ba dbb1308 d21fc46 77403d5 153cbe2 dbb1308 77403d5 a1a1e17 dbb1308 cf029f9 77403d5 cf029f9 a1a1e17 cf029f9 dbb1308 77403d5 dbb1308 77403d5 dbb1308 cf029f9 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 62 63 64 65 66 67 68 |
import gradio as gr
import os
from torchvision.transforms import Resize
from upstash_vector import Index
from datasets import load_dataset
from transformers import AutoFeatureExtractor, AutoModel
index = Index.from_env()
print(os.environ["UPSTASH_VECTOR_REST_URL"])
print(os.environ["UPSTASH_VECTOR_REST_TOKEN"])
resize_transform = Resize((250, 250))
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("HengJi/human_faces")
with gr.Blocks() as demo:
gr.Markdown(
"""
# Find Your Twins
Upload your face and find the most similar people from the X dataset. Powered by [Upstash Vector](https://upstash.com) 🚀
"""
)
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=600)
@input_image.upload(inputs=input_image, outputs=output_image)
def find_similar_faces(image):
resized_image = resize_transform(image)
inputs = extractor(images=resized_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[3:])]["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):
resized_image = resize_transform(image)
inputs = extractor(images=resized_image, return_tensors="pt")
outputs = model(**inputs)
embed = outputs.last_hidden_state[0][0]
result = index.query(vector=embed.tolist(), top_k=min(count, 9))
return [dataset["train"][int(vector.id[3:])]["image"] for vector in result]
if __name__ == "__main__":
demo.launch(debug=True)
|