app.py

import spaces
import gradio as gr
from PIL import Image, ImageDraw, ImageFont
from ultralytics import YOLO
import cv2
import numpy as np
import tempfile

@spaces.GPU
def yolo_inference(input_type, image, video, model_id, conf_threshold, iou_threshold, max_detection):
    """
    Performs object detection, instance segmentation, pose estimation,
    oriented object detection, or classification using a YOLOv11 model
    on either an image or a video.

    This function loads the specified YOLOv11 model and applies it to the
    provided input. For images, it returns an annotated image. For videos, it
    processes each frame and returns an annotated video. It includes error
    handling for missing inputs, returning blank outputs with informative messages.

    Args:
        input_type (str): Specifies the input type, either "Image" or "Video".
        image (PIL.Image.Image or None): The input image if `input_type` is "Image".
                                         None otherwise.
        video (str or None): The path to the input video file if `input_type` is "Video".
                             None otherwise.
        model_id (str): The identifier of the YOLOv11 model to use
                        (e.g., 'yolo11n.pt', 'yolo11s-seg.pt').
        conf_threshold (float): The confidence threshold for object detection.
                                Detections with lower confidence are discarded.
        iou_threshold (float): The Intersection over Union (IoU) threshold for
                               Non-Maximum Suppression (NMS). This is relevant
                               for detection and segmentation tasks.
        max_detection (int): The maximum number of detections to return per image or frame.

    Returns:
        tuple: A tuple containing two elements:
            - PIL.Image.Image or None: The annotated image if `input_type` was "Image",
                                       otherwise None.
            - str or None: The path to the annotated video file if `input_type` was "Video",
                           otherwise None.
    """
    # For YOLOv11, the model_id can directly be used by YOLO() as they are often
    # pre-trained weights included with the Ultralytics package.
    model = YOLO(model_id)

    if input_type == "Image":
        if image is None:
            width, height = 640, 480
            blank_image = Image.new("RGB", (width, height), color="white")
            draw = ImageDraw.Draw(blank_image)
            message = "No image provided"
            font = ImageFont.load_default(size=40)
            bbox = draw.textbbox((0, 0), message, font=font)
            text_width = bbox[2] - bbox[0]
            text_height = bbox[3] - bbox[1]
            text_x = (width - text_width) / 2
            text_y = (height - text_height) / 2
            draw.text((text_x, text_y), message, fill="black", font=font)
            return blank_image, None

        results = model.predict(
            source=image,
            conf=conf_threshold,
            iou=iou_threshold,
            imgsz=640,
            max_det=max_detection,
            show_labels=True,
            show_conf=True,
        )
        for r in results:
            image_array = r.plot()
            annotated_image = Image.fromarray(image_array[..., ::-1])
        return annotated_image, None

    elif input_type == "Video":
        if video is None:
            width, height = 640, 480
            blank_image = Image.new("RGB", (width, height), color="white")
            draw = ImageDraw.Draw(blank_image)
            message = "No video provided"
            font = ImageFont.load_default(size=40)
            bbox = draw.textbbox((0, 0), message, font=font)
            text_width = bbox[2] - bbox[0]
            text_height = bbox[3] - bbox[1]
            text_x = (width - text_width) / 2
            text_y = (height - text_height) / 2
            draw.text((text_x, text_y), message, fill="black", font=font)
            temp_video_file = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
            fourcc = cv2.VideoWriter_fourcc(*"mp4v")
            out = cv2.VideoWriter(temp_video_file, fourcc, 1, (width, height))
            frame = cv2.cvtColor(np.array(blank_image), cv2.COLOR_RGB2BGR)
            out.write(frame)
            out.release()
            return None, temp_video_file

        cap = cv2.VideoCapture(video)
        fps = cap.get(cv2.CAP_PROP_FPS) if cap.get(cv2.CAP_PROP_FPS) > 0 else 25
        frames = []
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            pil_frame = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
            results = model.predict(
                source=pil_frame,
                conf=conf_threshold,
                iou=iou_threshold,
                imgsz=640,
                max_det=max_detection,
                show_labels=True,
                show_conf=True,
            )
            for r in results:
                annotated_frame_array = r.plot()
                annotated_frame = cv2.cvtColor(annotated_frame_array, cv2.COLOR_BGR2RGB)
            frames.append(annotated_frame)
        cap.release()
        if len(frames) == 0:
            return None, None

        height_out, width_out, _ = frames[0].shape
        temp_video_file = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
        fourcc = cv2.VideoWriter_fourcc(*"mp4v")
        out = cv2.VideoWriter(temp_video_file, fourcc, fps, (width_out, height_out))
        for f in frames:
            f_bgr = cv2.cvtColor(f, cv2.COLOR_RGB2BGR)
            out.write(f_bgr)
        out.release()
        return None, temp_video_file

    else:
        return None, None

def update_visibility(input_type):
    """
    Adjusts the visibility of Gradio components based on the selected input type.

    This function dynamically shows or hides the image and video input/output
    components in the Gradio interface to ensure only relevant fields are visible.

    Args:
        input_type (str): The selected input type, either "Image" or "Video".

    Returns:
        tuple: A tuple of `gr.update` objects for the visibility of:
               (image input, video input, image output, video output).
    """
    if input_type == "Image":
        # image, video, output_image, output_video
        return gr.update(visible=True), gr.update(visible=False), gr.update(visible=True), gr.update(visible=False)
    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):
    """
    Wrapper function for `yolo_inference` specifically for Gradio examples that use images.

    This function simplifies the `yolo_inference` call for the `gr.Examples` component,
    ensuring only image-based inference is performed for predefined examples.

    Args:
        image (PIL.Image.Image): The input image for the example.
        model_id (str): The identifier of the YOLO model to use.
        conf_threshold (float): The confidence threshold.
        iou_threshold (float): The IoU threshold.
        max_detection (int): The maximum number of detections.

    Returns:
        PIL.Image.Image or None: The annotated image. Returns None if no image is processed.
    """
    annotated_image, _ = yolo_inference(
        input_type="Image",
        image=image,
        video=None,
        model_id=model_id,
        conf_threshold=conf_threshold,
        iou_threshold=iou_threshold,
        max_detection=max_detection
    )
    return annotated_image

with gr.Blocks() as app:
    gr.Markdown("# Yolo11: Object Detection, Instance Segmentation, Pose/Keypoints, Oriented Detection, Classification")
    gr.Markdown("Upload image(s) or video(s) for inference using the latest Ultralytics YOLO11 models.")
    
    with gr.Row():
        with gr.Column():
            image = gr.Image(type="pil", label="Image", visible=True)
            video = gr.Video(label="Video", visible=False)
            input_type = gr.Radio(
                choices=["Image", "Video"],
                value="Image",
                label="Input Type",
            )
            model_id = gr.Dropdown(
                label="Model Name",
                choices=[
                    'yolo11n.pt', 'yolo11s.pt', 'yolo11m.pt', 'yolo11l.pt', 'yolo11x.pt', 
                    'yolo11n-seg.pt', 'yolo11s-seg.pt', 'yolo11m-seg.pt', 'yolo11l-seg.pt', 'yolo11x-seg.pt',
                    'yolo11n-pose.pt', 'yolo11s-pose.pt', 'yolo11m-pose.pt', 'yolo11l-pose.pt', 'yolo11x-pose.pt',
                    'yolo11n-obb.pt', 'yolo11s-obb.pt', 'yolo11m-obb.pt', 'yolo11l-obb.pt', 'yolo11x-obb.pt',
                    'yolo11n-cls.pt', 'yolo11s-cls.pt', 'yolo11m-cls.pt', 'yolo11l-cls.pt', 'yolo11x-cls.pt'
                ],
                value="yolo11n.pt",
            )
            conf_threshold = gr.Slider(minimum=0, maximum=1, value=0.25, 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")
            infer_button = gr.Button("Detect Objects")
        with gr.Column():
            output_image = gr.Image(type="pil", show_label=False, visible=True)
            output_video = gr.Video(show_label=False, visible=False)
            gr.DeepLinkButton()
    
    # Toggle input/output visibility
    input_type.change(
        fn=update_visibility,
        inputs=input_type,
        outputs=[image, video, output_image, output_video],
    )
    
    # Main inference for button click
    infer_button.click(
        fn=yolo_inference,
        inputs=[input_type, image, video, model_id, conf_threshold, iou_threshold, max_detection],
        outputs=[output_image, output_video],
    )
    
    # Examples for images only
    gr.Examples(
        examples=[
            ["zidane.jpg", "yolo11s.pt", 0.25, 0.45, 300],
            ["bus.jpg", "yolo11m.pt", 0.25, 0.45, 300],
            ["yolo_vision.jpg", "yolo11x.pt", 0.25, 0.45, 300],
            ["Tricycle.jpg", "yolo11x-cls.pt", 0.25, 0.45, 300],
            ["tcganadolu.jpg", "yolo11m-obb.pt", 0.25, 0.45, 300],
            ["San Diego Airport.jpg", "yolo11x-seg.pt", 0.25, 0.45, 300],
            ["Theodore_Roosevelt.png", "yolo11l-pose.pt", 0.25, 0.45, 300],
        ],
        fn=yolo_inference_for_examples,
        inputs=[image, model_id, conf_threshold, iou_threshold, max_detection],
        outputs=[output_image],
        label="Examples (Images)",
    )

if __name__ == '__main__':
    app.launch(mcp_server=True)