Update app.py

app.py

@@ -11,14 +11,14 @@ COLORS = [[0.000, 0.447, 0.741], [0.850, 0.325, 0.098], [0.929, 0.694, 0.125],
           [0.494, 0.184, 0.556], [0.466, 0.674, 0.188], [0.301, 0.745, 0.933]]
 
 
-def get_class_list_from_input(classes_string: str):
+def process_class_list(classes_string: str):
     if classes_string == "":
         return []
     classes_list = classes_string.split(",")
     classes_list = [x.strip() for x in classes_list]
     return classes_list
 
-def infer(img, model_name: str, prob_threshold: int, classes_to_show = str):
+def model_inference(img, model_name: str, prob_threshold: int, classes_to_show = str):
     feature_extractor = AutoFeatureExtractor.from_pretrained(f"hustvl/{model_name}")
     model = YolosForObjectDetection.from_pretrained(f"hustvl/{model_name}")
 
@@ -36,7 +36,7 @@ def infer(img, model_name: str, prob_threshold: int, classes_to_show = str):
     postprocessed_outputs = feature_extractor.post_process(outputs, target_sizes)
     bboxes_scaled = postprocessed_outputs[0]['boxes']
     
-    classes_list = get_class_list_from_input(classes_to_show)
+    classes_list = process_class_list(classes_to_show)
     res_img = plot_results(img, probas[keep], bboxes_scaled[keep], model, classes_list)
     
     return res_img
@@ -75,10 +75,10 @@ image_in = gr.components.Image()
 image_out = gr.components.Image()
 model_choice = gr.components.Dropdown(["yolos-tiny", "yolos-small", "yolos-base", "yolos-small-300", "yolos-small-dwr"], value="yolos-small", label="YOLOS Model")
 prob_threshold_slider = gr.components.Slider(minimum=0, maximum=1.0, step=0.01, value=0.9, label="Probability Threshold")
-classes_to_show = gr.components.Textbox(placeholder="e.g. person, boat", label="Classes to use (empty means all classes)")
+classes_to_show = gr.components.Textbox(placeholder="e.g. person, truck", label="Classes to use (defaulted to detect all classes)")
 
 Iface = gr.Interface(
-    fn=infer,
+    fn=model_inference,
     inputs=[image_in,model_choice, prob_threshold_slider, classes_to_show],
     outputs=image_out,
     title="YOLOS - Object Detection",