switch to extractor

app.py

@@ -1,35 +1,22 @@
 import gradio as gr
-import os
-import torchvision.transforms as T
 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")
 
-# Data transformation chain.
-transformation_chain = T.Compose(
-    [
-        T.Resize(extractor.size["height"]),
-        T.CenterCrop(extractor.size["height"]),
-        T.ToTensor(),
-        T.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
-    ]
-)
-
 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) 🚀
+        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 🚀 
         """
     )
 
@@ -43,11 +30,10 @@ with gr.Blocks() as demo:
 
         @input_image.upload(inputs=input_image, outputs=output_image)
         def find_similar_faces(image):
-            t_image = transformation_chain(image)
-            inputs = extractor(images=t_image, return_tensors="pt")
+            inputs = extractor(images=image, return_tensors="pt")
             outputs = model(**inputs)
             embed = outputs.last_hidden_state[0][0]
-            result = index.query(vector=embed, top_k=4)
+            result = index.query(vector=embed.tolist(), top_k=4)
             return [dataset["train"][int(vector.id)]["image"] for vector in result]
 
     with gr.Tab("Advanced"):
@@ -60,13 +46,14 @@ with gr.Blocks() as demo:
                 adv_output_image = gr.Gallery(height=1000)
 
 
-        @adv_input_image.upload(inputs=[adv_input_image, adv_image_count], outputs=[adv_output_image])
+        @adv_input_image.upload(
+            inputs=[adv_input_image, adv_image_count], outputs=[adv_output_image]
+        )
         def find_similar_faces(image, count):
-            t_image = transformation_chain(image)
-            inputs = extractor(images=t_image, return_tensors="pt")
+            inputs = extractor(images=image, return_tensors="pt")
             outputs = model(**inputs)
             embed = outputs.last_hidden_state[0][0]
-            result = index.query(vector=embed, top_k=max(1, min(19, count)))
+            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__":