Update app.py

app.py

@@ -3,7 +3,6 @@ import numpy as np
 from PIL import Image
 from transformers import AutoImageProcessor, AutoModelForImageClassification
 import requests
-selected_model = "0-ma/vit-geometric-shapes-tiny"
 
 model_names = [
     "0-ma/swin-geometric-shapes-tiny",
@@ -27,23 +26,13 @@ examples = [
 ]
 
 labels = [example.split("_")[1].split(".")[0] for example in examples]
-# Load the default model
-#feature_extractor = AutoImageProcessor.from_pretrained(models["Tiny Model"])
-#model = AutoModelForImageClassification.from_pretrained(models["Tiny Model"])
 
+feature_extractors = {model_name: AutoImageProcessor.from_pretrained(model_name) for model_name in model_names}
+classification_models = {model_name: AutoModelForImageClassification.from_pretrained(model_name) for model_name in model_names}
 
-feature_extractors = { model_name :  AutoImageProcessor.from_pretrained(model_name) for model_name in model_names}
-classification_models = { model_name :  AutoModelForImageClassification.from_pretrained(model_name) for model_name in model_names}
-
-
-def predict(image):
-    # Load the selected model
-    # feature_extractor = AutoImageProcessor.from_pretrained(models[selected_model])
-    # model = AutoModelForImageClassification.from_pretrained(models[selected_model])
-
+def predict(image, selected_model):
     feature_extractor = feature_extractors[selected_model]
     model = classification_models[selected_model]
-    
 
     inputs = feature_extractor(images=[image], return_tensors="pt")
     logits = model(**inputs)['logits'].cpu().detach().numpy()[0]
@@ -57,17 +46,20 @@ def predict(image):
             confidences[labels[i]] = float(logits_positive[i])
     return confidences
 
-
 title = "Geometric Shape Classifier"
 description = "Select a model to classify geometric shapes."
 
-
-# Adding a dropdown for model selection
-gr.Interface(
+# Create the Gradio interface
+iface = gr.Interface(
     fn=predict,
-    inputs=gr.Image(type="pil"),
+    inputs=[
+        gr.Image(type="pil"),
+        gr.Dropdown(choices=model_names, label="Select Model", value=model_names[0])
+    ],
     outputs=gr.Label(),
     title=title,
     description=description,
     examples=examples
-).launch()
+)
+
+# Launch the interf