Spaces:
Running
Running
| 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") | |
| TOP_K = 1000 | |
| BASE_COUNT=4 | |
| MAX_COUNT = 10 | |
| with gr.Blocks() as demo: | |
| gr.Markdown( | |
| """ | |
| # Find Your Twins | |
| Upload your face and find the most similar faces 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. For best results please use 1x1 ratio face images, take a look at examples. Also increasing count in the advanced section results with more accurate searches. Disclaimer, this demo doesn't find your twins :), it finds similar face parts, shapes, features(nose, cheek, face, forehead shapes) that are encoded in the model. The Vector similarity search is powered by [Upstash Vector](https://upstash.com) 🚀. You can check our blog [post](https://huggingface.co/blog/omerXfaruq/serverless-image-similarity-with-upstash-vector) to learn more. | |
| """ | |
| ) | |
| with gr.Tab("Basic"): | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| input_image = gr.Image(type="pil") | |
| with gr.Column(scale=2): | |
| output_images = gr.Gallery() | |
| 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=TOP_K) | |
| return [dataset["train"][int(vector.id)]["image"] for vector in result[:BASE_COUNT]] | |
| gr.Examples( | |
| examples=[ | |
| dataset["train"][6]["image"], | |
| dataset["train"][7]["image"], | |
| dataset["train"][8]["image"], | |
| ], | |
| inputs=input_image, | |
| outputs=output_images, | |
| 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.Slider(1, MAX_COUNT, BASE_COUNT, label="Image Count") | |
| adv_button = gr.Button("Submit") | |
| with gr.Column(scale=2): | |
| adv_output_images = gr.Gallery() | |
| async def find_similar_faces(image, count): | |
| 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=TOP_K | |
| ) | |
| return [dataset["train"][int(vector.id)]["image"] for vector in result[:int(count)]] | |
| adv_button.click( | |
| fn=find_similar_faces, | |
| inputs=[adv_input_image, adv_image_count], | |
| outputs=[adv_output_images], | |
| ) | |
| adv_input_image.change( | |
| fn=find_similar_faces, | |
| inputs=[adv_input_image, adv_image_count], | |
| outputs=[adv_output_images], | |
| ) | |
| gr.Examples( | |
| examples=[ | |
| [dataset["train"][6]["image"], BASE_COUNT], | |
| [dataset["train"][7]["image"], BASE_COUNT], | |
| [dataset["train"][8]["image"], BASE_COUNT], | |
| ], | |
| inputs=[adv_input_image, adv_image_count], | |
| outputs=adv_output_images, | |
| fn=find_similar_faces, | |
| cache_examples=False, | |
| ) | |
| if __name__ == "__main__": | |
| demo.queue(default_concurrency_limit=40) | |
| demo.launch() | |