Pull the existing models from the model hub (#6)

- Pull the existing models from the model hub (b971dc429b86a8ebc93f491222263f7e6c756778)


Co-authored-by: Nima Boscarino <NimaBoscarino@users.noreply.huggingface.co>

app.py

@@ -1,10 +1,6 @@
 import gradio as gr
 import os
 
-os.system("wget https://github.com/WongKinYiu/yolov7/releases/download/v0.1/yolov7.pt")
-os.system("wget https://github.com/WongKinYiu/yolov7/releases/download/v0.1/yolov7-e6e.pt")
-os.system("wget https://github.com/WongKinYiu/yolov7/releases/download/v0.1/yolov7-e6.pt")
-
 import argparse
 import time
 from pathlib import Path
@@ -22,12 +18,24 @@ from utils.plots import plot_one_box
 from utils.torch_utils import select_device, load_classifier, time_synchronized, TracedModel
 from PIL import Image
  
+from huggingface_hub import hf_hub_download
+
+def load_model(model_name):
+    model_path = hf_hub_download(repo_id=f"Yolov7/{model_name}", filename=f"{model_name}.pt"), None)
+
+    return model_path
+
 
+model_names = [
+    "yolov7",
+    "yolov7-e6e",
+    "yolov7-e6",
+]
 
 
 def detect(img,model):
     parser = argparse.ArgumentParser()
-    parser.add_argument('--weights', nargs='+', type=str, default=model+".pt", help='model.pt path(s)')
+    parser.add_argument('--weights', nargs='+', type=str, default=load_model(model), help='model.pt path(s)')
     parser.add_argument('--source', type=str, default='Inference/', help='source')  # file/folder, 0 for webcam
     parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
     parser.add_argument('--conf-thres', type=float, default=0.25, help='object confidence threshold')
@@ -183,4 +191,4 @@ def detect(img,model):
     return Image.fromarray(im0[:,:,::-1])
 
    
-gr.Interface(detect,[gr.Image(type="pil"),gr.Dropdown(choices=["yolov7","yolov7-e6"])], gr.Image(type="pil"),title="Yolov7",examples=[["horses.jpeg", "yolov7"]],description="demo for <a href='https://github.com/WongKinYiu/yolov7' style='text-decoration: underline' target='_blank'>WongKinYiu/yolov7</a> Trainable bag-of-freebies sets new state-of-the-art for real-time object detectors").launch()
+gr.Interface(detect,[gr.Image(type="pil"),gr.Dropdown(choices=model_names)], gr.Image(type="pil"),title="Yolov7",examples=[["horses.jpeg", "yolov7"]],description="demo for <a href='https://github.com/WongKinYiu/yolov7' style='text-decoration: underline' target='_blank'>WongKinYiu/yolov7</a> Trainable bag-of-freebies sets new state-of-the-art for real-time object detectors").launch()