Update app.py

app.py

@@ -1,5 +1,11 @@
 """
-Inspired by HugoSchtr/DataCat_Yolov5
+Gradio APP for Yolo segmentation model of medival registers.
+
+Date: 20/12/2022
+Author: lterriel
+
+Test segmentation model with object-detection Yolo model
+App inspired by HugoSchtr/DataCat_Yolov5
 """
 
 import gradio as gr
@@ -8,12 +14,13 @@ from PIL import Image
 import torch
 from huggingface_hub import hf_hub_download
 
+# Add new models here
 model_names = [
-    "endp-yolov5x-35e-bs4",
+    "lterriel/endp-yolov5x-35e-bs4",
 ]
 
 def load_model(model_name):
-    model_path = hf_hub_download(repo_id=f"lterriel/{model_name}", filename="best.pt")
+    model_path = hf_hub_download(repo_id=model_name, filename="best.pt")
     return model_path
 
 def yolo_inference(im, model_path, threshold=0.50):
@@ -32,7 +39,7 @@ title = "YOLOv5 - Medieval Register Segmentation"
 #  ['', 0.50]]
 
 inputs = [gr.Image(type="pil", label="document image"),
-          gr.inputs.Dropdown(model_names, label="Model"),
+          gr.inputs.Dropdown(model_names, label="Model", value=model_names[0]),
           gr.Slider(maximum=1, step=0.01, value=0.50)]
 
 demo=gr.Interface(fn=yolo_inference,