app.py

import spaces
import supervision as sv
import PIL.Image as Image
from PIL import ImageDraw, ImageFont # Added ImageDraw and ImageFont import
from ultralytics import YOLO
from huggingface_hub import hf_hub_download, HfApi
import gradio as gr
import torch
import cv2
import numpy as np
import tempfile

global repo_id

repo_id = "atalaydenknalbant/asl-yolo-models"

def get_model_filenames(repo_id):
    """
    Retrieves a list of YOLO model filenames from a specified Hugging Face repository.

    This function connects to the Hugging Face Hub API, lists all files
    within the given repository, and filters for files ending with '.pt'
    to identify potential model weight files.

    Args:
        repo_id (str): The repository ID on Hugging Face Hub (e.g., "user/repo_name").

    Returns:
        list: A list of strings, where each string is the filename of a
              '.pt' model found in the repository.
    """
    api = HfApi()
    files = api.list_repo_files(repo_id)
    model_filenames = [file for file in files if file.endswith('.pt')]
    return model_filenames


model_filenames = get_model_filenames(repo_id)

def download_models(repo_id, model_id):
    """
    Downloads a specific model file from a Hugging Face repository to a local directory.

    This function uses `hf_hub_download` to fetch the model identified by `model_id`
    from the `repo_id` and saves it in the current working directory.

    Args:
        repo_id (str): The repository ID on Hugging Face Hub where the model is stored.
        model_id (str): The filename of the specific model to download (e.g., 'yolo11n.pt').

    Returns:
        str: The local file path to the downloaded model.
    """
    hf_hub_download(repo_id, filename=model_id, local_dir=f"./")
    return f"./{model_id}"

box_annotator = sv.BoxAnnotator()
category_dict = {0: 'A', 1: 'B', 2: 'C', 3: 'D', 4: 'E', 5: 'F', 6: 'G', 7: 'H', 8: 'I',
                 9: 'J', 10: 'K', 11: 'L', 12: 'M', 13: 'N', 14: 'O', 15: 'P', 16: 'Q',
                 17: 'R', 18: 'S', 19: 'T', 20: 'U', 21: 'V', 22: 'W', 23: 'X', 24: 'Y', 25: 'Z'}


@spaces.GPU
def yolo_inference(input_type, image, video, model_id, conf_threshold, iou_threshold, max_detection):
    """
    Performs ASL letter detection inference on an image or video using a YOLO model.

    This function first downloads the specified YOLO model. It then applies the model
    to the input, which can be either an image or a video. For images, it returns an
    annotated image. For videos, it processes each frame and reconstructs an annotated video.
    Error handling for missing inputs is included, returning blank outputs with 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 filename of the YOLO model to use (e.g., 'yolo11n.pt').
        conf_threshold (float): The confidence threshold for filtering detections.
                                Detections with confidence below this value are discarded.
        iou_threshold (float): The Intersection over Union (IoU) threshold for
                               Non-Maximum Suppression (NMS) to remove duplicate detections.
        max_detection (int): The maximum number of detections to display.

    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.
    """
    model_path = download_models(repo_id, model_id)
    model = YOLO(model_path)

    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(source=image, imgsz=640, iou=iou_threshold, conf=conf_threshold, verbose=False, max_det=max_detection)[0]
        detections = sv.Detections.from_ultralytics(results)
        
        labels = [
            f"{category_dict[class_id]} {confidence:.2f}"
            for class_id, confidence in zip(detections.class_id, detections.confidence)
        ]
        annotated_image = box_annotator.annotate(image, detections=detections, labels=labels)
        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(source=pil_frame, imgsz=640, iou=iou_threshold, conf=conf_threshold, verbose=False, max_det=max_detection)[0]
            detections = sv.Detections.from_ultralytics(results)

            labels = [
                f"{category_dict[class_id]} {confidence:.2f}"
                for class_id, confidence in zip(detections.class_id, detections.confidence)
            ]
            annotated_frame_array = box_annotator.annotate(np.array(pil_frame), detections=detections, labels=labels)
            annotated_frame = cv2.cvtColor(annotated_frame_array, cv2.COLOR_RGB2BGR)
            frames.append(annotated_frame)
        cap.release()
        if not frames:
            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:
            out.write(f)
        out.release()
        return None, temp_video_file

    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":
        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(title="ASL Letter Detector") as app:
    gr.HTML(
        """
    <h1 style='text-align: center'>
    YOLO Powered ASL(American Sign Language) Letter Detector PSA: It can't detect J or Z
    </h1>
        """)
    gr.Markdown("Upload an image or video for ASL letter detection using a YOLO model.")

    with gr.Row():
        with gr.Column():
            image = gr.Image(type="pil", label="Image Input", interactive=True, visible=True)
            video = gr.Video(label="Video Input", interactive=True, visible=False)
            input_type = gr.Radio(
                choices=["Image", "Video"],
                value="Image",
                label="Input Type",
            )
            
            model_id = gr.Dropdown(
                label="Model",
                choices=model_filenames,
                value=model_filenames[0] if model_filenames else "",
            )
            conf_threshold = gr.Slider(
                label="Confidence Threshold",
                minimum=0.1,
                maximum=1.0,
                step=0.1,
                value=0.45,
            )
            iou_threshold = gr.Slider(
                label="IoU Threshold",
                minimum=0.1,
                maximum=1.0,
                step=0.1,
                value=0.7,
            )
            
            max_detection = gr.Slider(
                label="Max Detection",
                minimum=1,
                step=1,
                value=1,
            )
            yolov_infer = gr.Button(value="Detect Objects")

        with gr.Column():
            output_image = gr.Image(type="pil", label="Annotated Image", interactive=False, visible=True)
            output_video = gr.Video(label="Annotated Video", interactive=False, visible=False)
            gr.DeepLinkButton()

    input_type.change(
        fn=update_visibility,
        inputs=input_type,
        outputs=[image, video, output_image, output_video],
    )

    yolov_infer.click(
        fn=yolo_inference,
        inputs=[
            input_type,
            image,
            video,
            model_id,
            conf_threshold,
            iou_threshold,
            max_detection,
        ],
        outputs=[output_image, output_video],
    )

    gr.Examples(
        examples=[
            ["b.jpg", "yolo11x.pt", 0.45, 0.7, 1],
            ["a.jpg", "yolo11s.pt", 0.45, 0.7, 1],
            ["y.jpg", "yolo11m.pt", 0.45, 0.7, 1],
        ],
        fn=yolo_inference_for_examples,
        inputs=[
            image,
            model_id,
            conf_threshold,
            iou_threshold,
            max_detection,
        ],
        outputs=[output_image],
        cache_examples=True,
        label="Examples (Images)",
    )

app.launch(mcp_server=True)