Update app.py

app.py

@@ -93,7 +93,7 @@ def get_gpu_info():
 
 
 def select_model(model_name):
-    if model_name == "Model 1":
+    if model_name == "YOLOv11_my_v1":
         return model1
     # elif model_name == "Model 2":
     #     return model2
@@ -101,7 +101,7 @@ def select_model(model_name):
     #     return model3
     elif model_name == "Model 4":
         return model4
-    elif model_name == "Model 5":
+    elif model_name == "YOLOv11_my_v5":
         return model5
     else:
         return model1
@@ -113,7 +113,7 @@ def predict_image(image, model_choice):
         return None, "Please upload an image"
 
     try:
-        if model_choice == "Model 6":
+        if model_choice == "ppe_rt_det":
             return predict_with_model6(image)
 
         # Otherwise, use YOLO models
@@ -311,8 +311,8 @@ with gr.Blocks(
     with gr.Tabs():
         # Model selector common for image and video
         model_selector = gr.Dropdown(
-            choices=["Model 1", "Model 5", "Model 6"],
-            value="Model 1",
+            choices=["YOLOv11_my_v1", "YOLOv11_my_v5", "ppe_rt_det"],
+            value="YOLOv11_my_v1",
             label="Choose model"
         )
         # Image Tab