Update app.py

app.py

@@ -25,9 +25,9 @@ def process_class_list(classes_string: str):
     return classes_list
 
 
-def model_inference(img, model_name = "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}")
+def model_inference(img, prob_threshold, classes_to_show):
+    feature_extractor = AutoFeatureExtractor.from_pretrained(f"hustvl/yolos-small-dwr")
+    model = YolosForObjectDetection.from_pretrained(f"hustvl/yolos-small-dwr")
 
     img = Image.fromarray(img)