Update app.py

app.py

@@ -2,16 +2,23 @@ import spaces
 import supervision as sv
 import PIL.Image as Image
 from ultralytics import YOLO
-from huggingface_hub import hf_hub_download
+from huggingface_hub import hf_hub_download, list_repo_files
 import gradio as gr
 
 global repo_id
 
-def download_models(model_id):
-    hf_hub_download(repo_id, filename=f"{model_id}", local_dir=f"./")
-    return f"./{model_id}"
-
 repo_id = "atalaydenknalbant/asl-yolo-models"
+
+def download_models(repo_id, model_id, file_extension=".pt"):
+    # Get list of files in the repository without using HfApi
+    files = list_repo_files(repo_id)
+    # Filter for model filenames with the given extension
+    model_filenames = [file for file in files if file.endswith(file_extension)]
+    
+    # Download the selected model
+    hf_hub_download(repo_id, filename=model_id, local_dir=f"./")
+    return model_filenames, f"./{model_id}"
+
 box_annotator = sv.BoxAnnotator()
 category_dict = {0: 'A', 1: 'B', 2: 'C', 3: 'D', 4: 'E', 5: 'F', 6: 'G', 7: 'H', 8: 'I',
                  9: 'J', 10: 'K', 11: 'L', 12: 'M', 13: 'N', 14: 'O', 15: 'P', 16: 'Q',
@@ -19,7 +26,9 @@ category_dict = {0: 'A', 1: 'B', 2: 'C', 3: 'D', 4: 'E', 5: 'F', 6: 'G', 7: 'H',
 
 @spaces.GPU
 def yolo_inference(image, model_id, conf_threshold, iou_threshold, max_detection):
-    model_path = download_models(model_id)
+    # Download models and get filenames within the same function
+    model_filenames, model_path = download_models(repo_id, model_id)
+    
     model = YOLO(model_path)
     results = model(source=image, imgsz=640, iou=iou_threshold, conf=conf_threshold, verbose=False, max_det=max_detection)[0]
     detections = sv.Detections.from_ultralytics(results)
@@ -33,13 +42,19 @@ def yolo_inference(image, model_id, conf_threshold, iou_threshold, max_detection
     return annotated_image
 
 def app():
+    # Fetch the model filenames directly in the app
+    model_filenames, _ = download_models(repo_id, "")
+    
     with gr.Blocks():
         with gr.Row():
             with gr.Column():
                 image = gr.Image(type="pil", label="Image", interactive=True)
     
-                model_id = gr.Textbox(label="Model ID", placeholder="Enter model filename (.pt)")
-                
+                model_id = gr.Dropdown(
+                    label="Model",
+                    choices=model_filenames,
+                    value=model_filenames[0] if model_filenames else "",
+                )
                 conf_threshold = gr.Slider(
                     label="Confidence Threshold",
                     minimum=0.1,