app.py

import gradio as gr
import spaces

import sahi.utils
from sahi import AutoDetectionModel
import sahi.predict
import sahi.slicing
from PIL import Image
import numpy
from ultralytics import YOLO
import sys
import types

if 'huggingface_hub.utils._errors' not in sys.modules:
    mock_errors = types.ModuleType('_errors')
    mock_errors.RepositoryNotFoundError = Exception
    sys.modules['huggingface_hub.utils._errors'] = mock_errors

IMAGE_SIZE = 640

# Images
sahi.utils.file.download_from_url(
    "https://user-images.githubusercontent.com/34196005/142730935-2ace3999-a47b-49bb-83e0-2bdd509f1c90.jpg",
    "apple_tree.jpg",
)
sahi.utils.file.download_from_url(
    "https://user-images.githubusercontent.com/34196005/142730936-1b397756-52e5-43be-a949-42ec0134d5d8.jpg",
    "highway.jpg",
)

sahi.utils.file.download_from_url(
    "https://user-images.githubusercontent.com/34196005/142742871-bf485f84-0355-43a3-be86-96b44e63c3a2.jpg",
    "highway2.jpg",
)

sahi.utils.file.download_from_url(
    "https://user-images.githubusercontent.com/34196005/142742872-1fefcc4d-d7e6-4c43-bbb7-6b5982f7e4ba.jpg",
    "highway3.jpg",
)

# Global model variable
model = None

@spaces.GPU(duration=60)
def load_yolo_model(model_name, confidence_threshold=0.5):
    """
    Loads a YOLOv11 detection model.

    Args:
        model_name (str): The name of the YOLOv11 model to load (e.g., "yolo11n.pt").
        confidence_threshold (float): The confidence threshold for object detection.

    Returns:
        AutoDetectionModel: The loaded SAHI AutoDetectionModel.
    """
    global model
    model_path = model_name
    model = AutoDetectionModel.from_pretrained(
        model_type="ultralytics", model_path=model_path, device=None, # auto device selection
        confidence_threshold=confidence_threshold, image_size=IMAGE_SIZE
    )
    return model

@spaces.GPU(duration=60)
def sahi_yolo_inference(
    image,
    yolo_model_name,
    confidence_threshold,
    max_detections,
    slice_height=512,
    slice_width=512,
    overlap_height_ratio=0.2,
    overlap_width_ratio=0.2,
    postprocess_type="NMS",
    postprocess_match_metric="IOU",
    postprocess_match_threshold=0.5,
    postprocess_class_agnostic=False,
):
    """
    Performs object detection using SAHI with a specified YOLOv11 model.

    Args:
        image (PIL.Image.Image): The input image for detection.
        yolo_model_name (str): The name of the YOLOv11 model to use for inference.
        confidence_threshold (float): The confidence threshold for object detection.
        max_detections (int): The maximum number of detections to return.
        slice_height (int): The height of each slice for sliced inference.
        slice_width (int): The width of each slice for sliced inference.
        overlap_height_ratio (float): The overlap ratio for slice height.
        overlap_width_ratio (float): The overlap ratio for slice width.
        postprocess_type (str): The type of postprocessing to apply ("NMS" or "GREEDYNMM").
        postprocess_match_metric (str): The metric for postprocessing matching ("IOU" or "IOS").
        postprocess_match_threshold (float): The threshold for postprocessing matching.
        postprocess_class_agnostic (bool): Whether postprocessing should be class agnostic.

    Returns:
        tuple: A tuple containing two PIL.Image.Image objects:
               - The image with standard YOLO inference results.
               - The image with SAHI sliced YOLO inference results.
    """
    load_yolo_model(yolo_model_name, confidence_threshold)

    image_width, image_height = image.size
    sliced_bboxes = sahi.slicing.get_slice_bboxes(
        image_height,
        image_width,
        slice_height,
        slice_width,
        False,
        overlap_height_ratio,
        overlap_width_ratio,
    )
    if len(sliced_bboxes) > 60:
        raise ValueError(
            f"{len(sliced_bboxes)} slices are too much for huggingface spaces, try smaller slice size."
        )

    # Standard inference
    prediction_result_1 = sahi.predict.get_prediction(
        image=image, detection_model=model,
    )
    
    # Filter by max_detections for standard inference
    if max_detections is not None and len(prediction_result_1.object_prediction_list) > max_detections:
        prediction_result_1.object_prediction_list = sorted(
            prediction_result_1.object_prediction_list, key=lambda x: x.score.value, reverse=True
        )[:max_detections]

    visual_result_1 = sahi.utils.cv.visualize_object_predictions(
        image=numpy.array(image),
        object_prediction_list=prediction_result_1.object_prediction_list,
    )
    output_1 = Image.fromarray(visual_result_1["image"])

    # Sliced inference
    prediction_result_2 = sahi.predict.get_sliced_prediction(
        image=image,
        detection_model=model,
        slice_height=int(slice_height),
        slice_width=int(slice_width),
        overlap_height_ratio=overlap_height_ratio,
        overlap_width_ratio=overlap_width_ratio,
        postprocess_type=postprocess_type,
        postprocess_match_metric=postprocess_match_metric,
        postprocess_match_threshold=postprocess_match_threshold,
        postprocess_class_agnostic=postprocess_class_agnostic,
    )

    # Filter by max_detections for sliced inference
    if max_detections is not None and len(prediction_result_2.object_prediction_list) > max_detections:
        prediction_result_2.object_prediction_list = sorted(
            prediction_result_2.object_prediction_list, key=lambda x: x.score.value, reverse=True
        )[:max_detections]

    visual_result_2 = sahi.utils.cv.visualize_object_predictions(
        image=numpy.array(image),
        object_prediction_list=prediction_result_2.object_prediction_list,
    )

    output_2 = Image.fromarray(visual_result_2["image"])

    return output_1, output_2


with gr.Blocks() as app:
    gr.Markdown("# Small Object Detection with SAHI + YOLOv11")
    gr.Markdown(
        "SAHI + YOLOv11 demo for small object detection. "
        "Upload your own image or click an example image to use."
    )

    with gr.Row():
        with gr.Column():
            original_image_input = gr.Image(type="pil", label="Original Image")
            yolo_model_dropdown = gr.Dropdown(
                choices=["yolo11n.pt", "yolo11s.pt", "yolo11m.pt", "yolo11l.pt", "yolo11x.pt"],
                value="yolo11s.pt",
                label="YOLOv11 Model",
            )
            confidence_threshold_slider = gr.Slider(
                minimum=0.0,
                maximum=1.0,
                step=0.01,
                value=0.5,
                label="Confidence Threshold",
            )
            max_detections_slider = gr.Slider(
                minimum=1,
                maximum=500,
                step=1,
                value=300,
                label="Max Detections",
            )
            slice_height_input = gr.Number(value=512, label="Slice Height")
            slice_width_input = gr.Number(value=512, label="Slice Width")
            overlap_height_ratio_slider = gr.Slider(
                minimum=0.0,
                maximum=1.0,
                step=0.01,
                value=0.2,
                label="Overlap Height Ratio",
            )
            overlap_width_ratio_slider = gr.Slider(
                minimum=0.0,
                maximum=1.0,
                step=0.01,
                value=0.2,
                label="Overlap Width Ratio",
            )
            postprocess_type_dropdown = gr.Dropdown(
                ["NMS", "GREEDYNMM"],
                type="value",
                value="NMS",
                label="Postprocess Type",
            )
            postprocess_match_metric_dropdown = gr.Dropdown(
                ["IOU", "IOS"], type="value", value="IOU", label="Postprocess Match Metric"
            )
            postprocess_match_threshold_slider = gr.Slider(
                minimum=0.0,
                maximum=1.0,
                step=0.01,
                value=0.5,
                label="Postprocess Match Threshold",
            )
            postprocess_class_agnostic_checkbox = gr.Checkbox(value=True, label="Postprocess Class Agnostic")

            submit_button = gr.Button("Run Inference")

        with gr.Column():
            output_standard = gr.Image(type="pil", label="YOLOv11 Standard")
            output_sahi_sliced = gr.Image(type="pil", label="YOLOv11 + SAHI Sliced")

    gr.Examples(
        examples=[
            ["apple_tree.jpg", "yolo11s.pt", 0.5, 300, 256, 256, 0.2, 0.2, "NMS", "IOU", 0.4, True],
            ["highway.jpg", "yolo11s.pt", 0.5, 300, 256, 256, 0.2, 0.2, "NMS", "IOU", 0.4, True],
            ["highway2.jpg", "yolo11s.pt", 0.5, 300, 512, 512, 0.2, 0.2, "NMS", "IOU", 0.4, True],
            ["highway3.jpg", "yolo11s.pt", 0.5, 300, 512, 512, 0.2, 0.2, "NMS", "IOU", 0.4, True],
        ],
        inputs=[
            original_image_input,
            yolo_model_dropdown,
            confidence_threshold_slider,
            max_detections_slider,
            slice_height_input,
            slice_width_input,
            overlap_height_ratio_slider,
            overlap_width_ratio_slider,
            postprocess_type_dropdown,
            postprocess_match_metric_dropdown,
            postprocess_match_threshold_slider,
            postprocess_class_agnostic_checkbox,
        ],
        outputs=[output_standard, output_sahi_sliced],
        fn=sahi_yolo_inference,
        cache_examples=True,
    )

    submit_button.click(
        fn=sahi_yolo_inference,
        inputs=[
            original_image_input,
            yolo_model_dropdown,
            confidence_threshold_slider,
            max_detections_slider,
            slice_height_input,
            slice_width_input,
            overlap_height_ratio_slider,
            overlap_width_ratio_slider,
            postprocess_type_dropdown,
            postprocess_match_metric_dropdown,
            postprocess_match_threshold_slider,
            postprocess_class_agnostic_checkbox,
        ],
        outputs=[output_standard, output_sahi_sliced],
    )

app.launch(mcp_server=True)