Update app.py

app.py

@@ -16,13 +16,13 @@ def normalize_image(img):
 
 def load_model(model_name):
     # Load the model based on the model_name input
-    if model_name == "model1":
+    if model_name == 1:
         return tf.keras.models.load_model("model/Xception.h5")
-    elif model_name == "model2":
+    elif model_name == 2:
         return tf.keras.models.load_model("model/InceptionV3.h5")
-    elif model_name == "model3":
+    elif model_name == 3:
         return tf.keras.models.load_model("model/InceptionResNetV2.h5")
-    elif model_name == "model4":
+    elif model_name == 4:
         return tf.keras.models.load_model("model/DenseNet201.h5")
     else:
         raise ValueError("Invalid model_name")
@@ -41,14 +41,22 @@ def predict_top_classes(img, model_name):
 
     return {CLASS_LABEL[i]: str(v) for i, v in zip(top5_idx, top5_value)}
 
+models = {
+    "Xception": 1,
+    "InceptionV3": 2,
+    "InceptionResNetV2": 3,
+    "DenseNet201": 4
+}
+
 interface = gr.Interface(
     predict_top_classes,
     [
         gr.inputs.Image(type='pil'),
-        gr.inputs.Button(label="Model 1 (Xception)", value="model1"),
-        gr.inputs.Button(label="Model 2 (InceptionV3)", value="model2"),
-        gr.inputs.Button(label="Model 3 (InceptionResNetV2)", value="model3"),
-        gr.inputs.Button(label="Model 4 (DenseNet201)", value="model4")
+        gr.Dropdown(
+            choices=list(models.keys()),
+            value=list(models.values()),
+            label="Select a model"
+        )
     ],
     outputs='label'
 )