Filter detections to sports ball by default

app.py

@@ -23,8 +23,36 @@ def download_model(model_filename):
     """
     return hf_hub_download(repo_id="atalaydenknalbant/Yolov13", filename=model_filename)
 
+TARGET_ALIASES = {
+    "ball": "sports ball",
+    "soccer ball": "sports ball",
+    "football": "sports ball",
+    "sports ball": "sports ball",
+}
+
+
+def resolve_target_class(model, target_label: str):
+    """
+    Resolve a human-provided class name to YOLO class indices.
+
+    Args:
+        model (YOLO): Loaded YOLO model instance.
+        target_label (str): Label entered by the user.
+
+    Returns:
+        list[int] | None: List of class indices to filter on, or None to keep all classes.
+    """
+    if not target_label:
+        return None
+
+    cleaned = target_label.strip().lower()
+    canonical = TARGET_ALIASES.get(cleaned, cleaned)
+    matching_ids = [idx for idx, name in model.names.items() if name.lower() == canonical]
+    return matching_ids or None
+
+
 @spaces.GPU
-def yolo_inference(input_type, image, video, model_id, conf_threshold, iou_threshold, max_detection):
+def yolo_inference(input_type, image, video, model_id, conf_threshold, iou_threshold, max_detection, target_class):
     """
     Performs object detection inference using a YOLOv13 model on either an image or a video.
 
@@ -71,12 +99,14 @@ def yolo_inference(input_type, image, video, model_id, conf_threshold, iou_thres
             return blank_image, None
 
         model = YOLO(model_path)
+        class_ids = resolve_target_class(model, target_class)
         results = model.predict(
             source=image,
             conf=conf_threshold,
             iou=iou_threshold,
             imgsz=640,
             max_det=max_detection,
+            classes=class_ids,
             show_labels=True,
             show_conf=True,
         )
@@ -107,6 +137,7 @@ def yolo_inference(input_type, image, video, model_id, conf_threshold, iou_thres
             return None, temp_video_file
 
         model = YOLO(model_path)
+        class_ids = resolve_target_class(model, target_class)
         cap = cv2.VideoCapture(video)
         fps = cap.get(cv2.CAP_PROP_FPS) if cap.get(cv2.CAP_PROP_FPS) > 0 else 25
         frames = []
@@ -121,6 +152,7 @@ def yolo_inference(input_type, image, video, model_id, conf_threshold, iou_thres
                 iou=iou_threshold,
                 imgsz=640,
                 max_det=max_detection,
+                classes=class_ids,
                 show_labels=True,
                 show_conf=True,
             )
@@ -163,7 +195,7 @@ def update_visibility(input_type):
     else:
         return gr.update(visible=False), gr.update(visible=True), gr.update(visible=False), gr.update(visible=True)
 
-def yolo_inference_for_examples(image, model_id, conf_threshold, iou_threshold, max_detection):
+def yolo_inference_for_examples(image, model_id, conf_threshold, iou_threshold, max_detection, target_class):
     """
     Wrapper function for `yolo_inference` specifically for Gradio examples that use images.
 
@@ -187,7 +219,8 @@ def yolo_inference_for_examples(image, model_id, conf_threshold, iou_threshold,
         model_id=model_id,
         conf_threshold=conf_threshold,
         iou_threshold=iou_threshold,
-        max_detection=max_detection
+        max_detection=max_detection,
+        target_class=target_class
     )
     return annotated_image
 
@@ -234,6 +267,7 @@ with gr.Blocks(theme=theme) as app:
             conf_threshold = gr.Slider(minimum=0, maximum=1, value=0.35, label="Confidence Threshold")
             iou_threshold = gr.Slider(minimum=0, maximum=1, value=0.45, label="IoU Threshold")
             max_detection = gr.Slider(minimum=1, maximum=300, step=1, value=300, label="Max Detection")
+            target_class = gr.Textbox(value="sports ball", label="Target class (default: sports ball)")
             infer_button = gr.Button("Detect Objects", variant="primary")
         with gr.Column():
             output_image = gr.Image(type="pil", show_label=False, show_share_button=False, visible=True)
@@ -248,18 +282,18 @@ with gr.Blocks(theme=theme) as app:
 
     infer_button.click(
         fn=yolo_inference,
-        inputs=[input_type, image, video, model_id, conf_threshold, iou_threshold, max_detection],
+        inputs=[input_type, image, video, model_id, conf_threshold, iou_threshold, max_detection, target_class],
         outputs=[output_image, output_video],
     )
 
     gr.Examples(
         examples=[
-            ["zidane.jpg", "yolov13s.pt", 0.35, 0.45, 300],
-            ["bus.jpg", "yolov13l.pt", 0.35, 0.45, 300],
-            ["yolo_vision.jpg", "yolov13x.pt", 0.35, 0.45, 300],
+            ["zidane.jpg", "yolov13s.pt", 0.35, 0.45, 300, "sports ball"],
+            ["bus.jpg", "yolov13l.pt", 0.35, 0.45, 300, "sports ball"],
+            ["yolo_vision.jpg", "yolov13x.pt", 0.35, 0.45, 300, "sports ball"],
         ],
         fn=yolo_inference_for_examples,
-        inputs=[image, model_id, conf_threshold, iou_threshold, max_detection],
+        inputs=[image, model_id, conf_threshold, iou_threshold, max_detection, target_class],
         outputs=[output_image],
         label="Examples (Images)",
     )