Update app.py

app.py

@@ -2,22 +2,27 @@ import spaces
 import supervision as sv
 import PIL.Image as Image
 from ultralytics import YOLO
-from huggingface_hub import hf_hub_download, list_repo_files
+from huggingface_hub import hf_hub_download
 import gradio as gr
 
 global repo_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)]
-    
+# Model filenames directly provided, since they are known
+model_filenames = [
+    "yolov10s.pt",
+    "yolov10x.pt",
+    "yolov8s.pt",
+    "yolov8x.pt",
+    "yolov9e.pt",
+    "yolov9s.pt"
+]
+
+def download_models(repo_id, model_id):
     # Download the selected model
     hf_hub_download(repo_id, filename=model_id, local_dir=f"./")
-    return model_filenames, f"./{model_id}"
+    return 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',
@@ -26,8 +31,8 @@ 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):
-    # Download models and get filenames within the same function
-    model_filenames, model_path = download_models(repo_id, model_id)
+    # Download models
+    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]
@@ -42,9 +47,6 @@ 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():
@@ -142,4 +144,4 @@ with gradio_app:
         with gr.Column():
             app()
 
-gradio_app.launch(debug=True)
+gradio_app.launch()