Update app.py

app.py

@@ -25,7 +25,7 @@ def process_class_list(classes_string: str):
     return classes_list
 
 
-def model_inference(img, model_name: str, prob_threshold: int, classes_to_show=str):
+def model_inference(img, model_name: str = "yolos-small-dwr", prob_threshold: int, classes_to_show=str):
     feature_extractor = AutoFeatureExtractor.from_pretrained(f"hustvl/{model_name}")
     model = YolosForObjectDetection.from_pretrained(f"hustvl/{model_name}")
 
@@ -85,9 +85,8 @@ description = """Object Detection"""
 
 image_in = gr.components.Image()
 image_out = gr.components.Image()
-model_choice = "yolos-small"
 prob_threshold_slider = gr.components.Slider(
-    minimum=0, maximum=1.0, step=0.01, value=0.9, label="Probability Threshold"
+    minimum=0, maximum=1.0, step=0.01, value=0.7, label="Probability Threshold"
 )
 classes_to_show = gr.components.Textbox(
     placeholder="e.g. car, dog",
@@ -96,7 +95,7 @@ classes_to_show = gr.components.Textbox(
 
 Iface = gr.Interface(
     fn=model_inference,
-    inputs=[image_in, model_choice, prob_threshold_slider, classes_to_show],
+    inputs=[image_in, prob_threshold_slider, classes_to_show],
     outputs=image_out,
     title="Object Detection",
     description=description,