Spaces:
Running
Running
Commit
•
b33e6dd
1
Parent(s):
3bc904e
switch to extractor
Browse files
app.py
CHANGED
@@ -1,35 +1,22 @@
|
|
1 |
import gradio as gr
|
2 |
-
import os
|
3 |
-
import torchvision.transforms as T
|
4 |
from upstash_vector import Index
|
5 |
from datasets import load_dataset
|
6 |
from transformers import AutoFeatureExtractor, AutoModel
|
7 |
|
8 |
index = Index.from_env()
|
9 |
|
10 |
-
|
11 |
model_ckpt = "google/vit-base-patch16-224-in21k"
|
12 |
extractor = AutoFeatureExtractor.from_pretrained(model_ckpt)
|
13 |
model = AutoModel.from_pretrained(model_ckpt)
|
14 |
hidden_dim = model.config.hidden_size
|
15 |
dataset = load_dataset("BounharAbdelaziz/Face-Aging-Dataset")
|
16 |
|
17 |
-
# Data transformation chain.
|
18 |
-
transformation_chain = T.Compose(
|
19 |
-
[
|
20 |
-
T.Resize(extractor.size["height"]),
|
21 |
-
T.CenterCrop(extractor.size["height"]),
|
22 |
-
T.ToTensor(),
|
23 |
-
T.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
|
24 |
-
]
|
25 |
-
)
|
26 |
-
|
27 |
with gr.Blocks() as demo:
|
28 |
gr.Markdown(
|
29 |
"""
|
30 |
# Find Your Twins
|
31 |
|
32 |
-
Upload your face and find the most similar people from
|
33 |
"""
|
34 |
)
|
35 |
|
@@ -43,11 +30,10 @@ with gr.Blocks() as demo:
|
|
43 |
|
44 |
@input_image.upload(inputs=input_image, outputs=output_image)
|
45 |
def find_similar_faces(image):
|
46 |
-
|
47 |
-
inputs = extractor(images=t_image, return_tensors="pt")
|
48 |
outputs = model(**inputs)
|
49 |
embed = outputs.last_hidden_state[0][0]
|
50 |
-
result = index.query(vector=embed, top_k=4)
|
51 |
return [dataset["train"][int(vector.id)]["image"] for vector in result]
|
52 |
|
53 |
with gr.Tab("Advanced"):
|
@@ -60,13 +46,14 @@ with gr.Blocks() as demo:
|
|
60 |
adv_output_image = gr.Gallery(height=1000)
|
61 |
|
62 |
|
63 |
-
@adv_input_image.upload(
|
|
|
|
|
64 |
def find_similar_faces(image, count):
|
65 |
-
|
66 |
-
inputs = extractor(images=t_image, return_tensors="pt")
|
67 |
outputs = model(**inputs)
|
68 |
embed = outputs.last_hidden_state[0][0]
|
69 |
-
result = index.query(vector=embed, top_k=max(1, min(19, count)))
|
70 |
return [dataset["train"][int(vector.id)]["image"] for vector in result]
|
71 |
|
72 |
if __name__ == "__main__":
|
|
|
1 |
import gradio as gr
|
|
|
|
|
2 |
from upstash_vector import Index
|
3 |
from datasets import load_dataset
|
4 |
from transformers import AutoFeatureExtractor, AutoModel
|
5 |
|
6 |
index = Index.from_env()
|
7 |
|
|
|
8 |
model_ckpt = "google/vit-base-patch16-224-in21k"
|
9 |
extractor = AutoFeatureExtractor.from_pretrained(model_ckpt)
|
10 |
model = AutoModel.from_pretrained(model_ckpt)
|
11 |
hidden_dim = model.config.hidden_size
|
12 |
dataset = load_dataset("BounharAbdelaziz/Face-Aging-Dataset")
|
13 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
14 |
with gr.Blocks() as demo:
|
15 |
gr.Markdown(
|
16 |
"""
|
17 |
# Find Your Twins
|
18 |
|
19 |
+
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 🚀
|
20 |
"""
|
21 |
)
|
22 |
|
|
|
30 |
|
31 |
@input_image.upload(inputs=input_image, outputs=output_image)
|
32 |
def find_similar_faces(image):
|
33 |
+
inputs = extractor(images=image, return_tensors="pt")
|
|
|
34 |
outputs = model(**inputs)
|
35 |
embed = outputs.last_hidden_state[0][0]
|
36 |
+
result = index.query(vector=embed.tolist(), top_k=4)
|
37 |
return [dataset["train"][int(vector.id)]["image"] for vector in result]
|
38 |
|
39 |
with gr.Tab("Advanced"):
|
|
|
46 |
adv_output_image = gr.Gallery(height=1000)
|
47 |
|
48 |
|
49 |
+
@adv_input_image.upload(
|
50 |
+
inputs=[adv_input_image, adv_image_count], outputs=[adv_output_image]
|
51 |
+
)
|
52 |
def find_similar_faces(image, count):
|
53 |
+
inputs = extractor(images=image, return_tensors="pt")
|
|
|
54 |
outputs = model(**inputs)
|
55 |
embed = outputs.last_hidden_state[0][0]
|
56 |
+
result = index.query(vector=embed.tolist(), top_k=max(1, min(19, count)))
|
57 |
return [dataset["train"][int(vector.id)]["image"] for vector in result]
|
58 |
|
59 |
if __name__ == "__main__":
|