Update app.py

app.py

@@ -27,13 +27,25 @@ models = {
 
 }
 
+model_names = [
+    "0-ma/swin-geometric-shapes-tiny",
+    "0-ma/mobilenet-v2-geometric-shapes",
+    "0-ma/focalnet-geometric-shapes-tiny",
+    "0-ma/efficientnet-b2-geometric-shapes",
+    "0-ma/beit-geometric-shapes-base",
+    "0-ma/mit-b0-geometric-shapes",
+    "0-ma/vit-geometric-shapes-base",
+    "0-ma/resnet-geometric-shapes",
+    "0-ma/vit-geometric-shapes-tiny",
+]
+
 # 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(models[model_name]) for model_name in models}
-classification_models = { model_name :  AutoModelForImageClassification.from_pretrained(models[model_name]) for model_name in models}
+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, selected_model):
@@ -73,7 +85,7 @@ examples = [
 # Adding a dropdown for model selection
 gr.Interface(
     fn=predict,
-    inputs=[gr.Image(type="pil"), gr.Dropdown(list(models.keys()), label="Select Model")],
+    inputs=[gr.Image(type="pil"), gr.Dropdown(model_names), label="Select Model")],
     outputs=gr.Label(),
     title=title,
     description=description,