Update app.py

app.py

@@ -1,22 +1,27 @@
 import gradio as gr
 import os
-from torchvision.transforms import Resize
+import torchvision.transforms as T
 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))
+# 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)),
+    ]
+)
 
 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")
+dataset = load_dataset("BounharAbdelaziz/Face-Aging-Dataset")
 
 with gr.Blocks() as demo:
     gr.Markdown(
@@ -31,18 +36,18 @@ with gr.Blocks() as demo:
         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)
+            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):
-            resized_image = resize_transform(image)
-            inputs = extractor(images=resized_image, return_tensors="pt")
+            t_image = transformation_chain(image)
+            inputs = extractor(images=t_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]
+            result = index.query(vector=embed, top_k=4)
+            return [dataset["train"][int(vector.id)]["image"] for vector in result]
 
     with gr.Tab("Advanced"):
         with gr.Row():
@@ -56,12 +61,12 @@ with gr.Blocks() as demo:
 
         @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")
+            t_image = transformation_chain(image)
+            inputs = extractor(images=t_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]
+            result = index.query(vector=embed, 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)