app.py

import os  # added for cache_examples
from pathlib import Path

import gradio as gr
import numpy as np
import supervision as sv
from gradio import ColorPicker
from PIL import Image
from torch import cuda, device
from ultralytics import YOLO

# Use GPU if available
if cuda.is_available():
    device = device("cuda")
else:
    device = device("cpu")


TITLE = """<h1 align="center">Supervision Annotator Playground 🚀</h1>"""
SUBTITLE = """<h2 align="center">Experiment with Supervision Annotators</h2>"""
BANNER = """
<div align="center">
    <p>
        <a align="center" href="https://supervision.roboflow.com/" target="_blank">
            <img style="max-width: 50%; height: auto; margin: 0 auto; display: block; padding: 20"
                src="https://media.roboflow.com/open-source/supervision/rf-supervision-banner.png?updatedAt=1678995927529">
        </a>
    </p>
</div>
"""  # noqa: E501 title/docs
DESC = """
<div style="text-align: center; display: flex; justify-content: center; align-items: center;">
    <a href="https://huggingface.co/spaces/Roboflow/Annotators?duplicate=true">
        <img src="https://bit.ly/3gLdBN6" alt="Duplicate Space" style="margin-right: 10px;">
    </a>
    <a href="https://github.com/roboflow/supervision">
        <img alt="GitHub Repo stars" src="https://img.shields.io/github/stars/roboflow/supervision"
            style="margin-right: 10px;">
    </a>
    <a href="https://colab.research.google.com/github/roboflow/supervision/blob/main/demo.ipynb">
        <img alt="Open In Colab" src="https://colab.research.google.com/assets/colab-badge.svg"
            style="margin-right: 10px;">
    </a>
</div>
"""  # noqa: E501 title/docs

last_detections = sv.Detections.empty()
last_labels: list[str] = []


def load_model(img, model: str | Path = "yolov8s-seg.pt"):
    # Load model, get results and return detections/labels
    model = YOLO(model=model)
    result = model(img, verbose=False, imgsz=1280)[0]
    detections = sv.Detections.from_ultralytics(result)
    labels = [
        f"{model.model.names[class_id]} {confidence:.2f}"
        for class_id, confidence in zip(detections.class_id, detections.confidence)
    ]

    return detections, labels


def calculate_crop_dim(a, b):
    # Calculates the crop dimensions of the image resultant
    if a > b:
        width = a
        height = a
    else:
        width = b
        height = b

    return width, height


def annotators(
    img,
    last_detections,
    annotators_list,
    last_labels,
    colorbb,
    colormask,
    colorellipse,
    colorbc,
    colorcir,
    colorlabel,
    colorhalo,
    colortri,
    colordot,
) -> np.ndarray:
    if last_detections == sv.Detections.empty():
        gr.Warning("Detection is empty please add image and annotate first")
        return np.zeros()

    if "Blur" in annotators_list:
        # Apply Blur
        blur_annotator = sv.BlurAnnotator()
        img = blur_annotator.annotate(img, detections=last_detections)

    if "BoundingBox" in annotators_list:
        # Draw Boundingbox
        box_annotator = sv.BoundingBoxAnnotator(sv.Color.from_hex(str(colorbb)))
        img = box_annotator.annotate(img, detections=last_detections)

    if "Mask" in annotators_list:
        # Draw Mask
        mask_annotator = sv.MaskAnnotator(sv.Color.from_hex(str(colormask)))
        img = mask_annotator.annotate(img, detections=last_detections)

    if "Ellipse" in annotators_list:
        # Draw Ellipse
        ellipse_annotator = sv.EllipseAnnotator(sv.Color.from_hex(str(colorellipse)))
        img = ellipse_annotator.annotate(img, detections=last_detections)

    if "BoxCorner" in annotators_list:
        # Draw Box corner
        corner_annotator = sv.BoxCornerAnnotator(sv.Color.from_hex(str(colorbc)))
        img = corner_annotator.annotate(img, detections=last_detections)

    if "Circle" in annotators_list:
        # Draw Circle
        circle_annotator = sv.CircleAnnotator(sv.Color.from_hex(str(colorcir)))
        img = circle_annotator.annotate(img, detections=last_detections)

    if "Label" in annotators_list:
        # Draw Label
        label_annotator = sv.LabelAnnotator(text_position=sv.Position.CENTER)
        label_annotator = sv.LabelAnnotator(sv.Color.from_hex(str(colorlabel)))
        img = label_annotator.annotate(
            img, detections=last_detections, labels=last_labels
        )

    if "Pixelate" in annotators_list:
        # Apply PixelateAnnotator
        pixelate_annotator = sv.PixelateAnnotator()
        img = pixelate_annotator.annotate(img, detections=last_detections)

    if "Halo" in annotators_list:
        # Draw HaloAnnotator
        halo_annotator = sv.HaloAnnotator(sv.Color.from_hex(str(colorhalo)))
        img = halo_annotator.annotate(img, detections=last_detections)

    if "HeatMap" in annotators_list:
        # Draw HeatMapAnnotator
        heatmap_annotator = sv.HeatMapAnnotator()
        img = heatmap_annotator.annotate(img, detections=last_detections)

    if "Dot" in annotators_list:
        # Dot DotAnnotator
        dot_annotator = sv.DotAnnotator(sv.Color.from_hex(str(colordot)))
        img = dot_annotator.annotate(img, detections=last_detections)

    if "Triangle" in annotators_list:
        # Draw TriangleAnnotator
        tri_annotator = sv.TriangleAnnotator(sv.Color.from_hex(str(colortri)))
        img = tri_annotator.annotate(img, detections=last_detections)

    # crop image for the largest possible square
    res_img = Image.fromarray(img)
    # print(type(res_img))
    x = 0
    y = 0

    # print("size of the pil im=", res_img.size)
    (v1, v2) = res_img.size
    width, height = calculate_crop_dim(v1, v2)
    # print(width, height)
    my_img = np.array(res_img)

    crop_img = my_img[y : y + height, x : x + width]
    # print(type(crop_img))

    return crop_img[..., ::-1].copy()  # BGR to RGB using numpy


def annotator(
    img,
    model,
    annotators_list,
    colorbb,
    colormask,
    colorellipse,
    colorbc,
    colorcir,
    colorlabel,
    colorhalo,
    colortri,
    colordot,
    progress=gr.Progress(track_tqdm=True),
) -> np.ndarray:
    """
    Function that changes the color of annotators
    Args:
        annotators: Icon whose color needs to be changed.
        color: Chosen color with which to edit the input icon in Hex.
        img: Input image is numpy matrix in BGR.
    Returns:
        annotators: annotated image
    """

    img = img[..., ::-1].copy()  # BGR to RGB using numpy

    detections, labels = load_model(img, model)
    last_detections = detections
    last_labels = labels

    return annotators(
        img,
        last_detections,
        annotators_list,
        last_labels,
        colorbb,
        colormask,
        colorellipse,
        colorbc,
        colorcir,
        colorlabel,
        colorhalo,
        colortri,
        colordot,
    )


purple_theme = theme = gr.themes.Soft(primary_hue=gr.themes.colors.purple).set(
    button_primary_background_fill="*primary_600",
    button_primary_background_fill_hover="*primary_700",
    checkbox_label_background_fill_selected="*primary_600",
    checkbox_background_color_selected="*primary_400",
)

with gr.Blocks(theme=purple_theme) as app:
    gr.HTML(TITLE)
    gr.HTML(SUBTITLE)
    gr.HTML(BANNER)
    gr.HTML(DESC)

    models = gr.Dropdown(
        [
            "yolov8n-seg.pt",
            "yolov8s-seg.pt",
            "yolov8m-seg.pt",
            "yolov8l-seg.pt",
            "yolov8x-seg.pt",
        ],
        type="value",
        value="yolov8s-seg.pt",
        label="Select Model:",
    )

    annotators_list = gr.CheckboxGroup(
        choices=[
            "BoundingBox",
            "Mask",
            "Halo",
            "Ellipse",
            "BoxCorner",
            "Circle",
            "Label",
            "Blur",
            "Pixelate",
            "HeatMap",
            "Dot",
            "Triangle",
        ],
        value=["BoundingBox", "Mask"],
        label="Select Annotators:",
    )

    gr.Markdown("## Color Picker 🎨")
    with gr.Row(variant="panel"):
        with gr.Column():
            colorbb = gr.ColorPicker(value="#A351FB", label="BoundingBox")
            colormask = gr.ColorPicker(value="#A351FB", label="Mask")
            colorellipse = gr.ColorPicker(value="#A351FB", label="Ellipse")
        with gr.Column():
            colorbc = gr.ColorPicker(value="#A351FB", label="BoxCorner")
            colorcir = gr.ColorPicker(value="#A351FB", label="Circle")
            colorlabel = gr.ColorPicker(value="#A351FB", label="Label")
        with gr.Column():
            colorhalo = gr.ColorPicker(value="#A351FB", label="Halo")
            colordot = gr.ColorPicker(value="#A351FB", label="Dot")
            colortri = gr.ColorPicker(value="#A351FB", label="Triangle")

    with gr.Row():
        with gr.Column():
            with gr.Tab("Input image"):
                image_input = gr.Image(type="numpy", show_label=False)
        with gr.Column():
            with gr.Tab("Result image"):
                image_output = gr.Image(type="numpy", show_label=False)
    image_button = gr.Button(value="Annotate it!", variant="primary")

    image_button.click(
        annotator,
        inputs=[
            image_input,
            models,
            annotators_list,
            colorbb,
            colormask,
            colorellipse,
            colorbc,
            colorcir,
            colorlabel,
            colorhalo,
            colortri,
            colordot,
        ],
        outputs=image_output,
    )

    gr.Markdown("## Image Examples 🖼️")
    gr.Examples(
        examples=[
            os.path.join(os.path.abspath(""), "./assets/city.jpg"),
            os.path.join(os.path.abspath(""), "./assets/household.jpg"),
            os.path.join(os.path.abspath(""), "./assets/industry.jpg"),
            os.path.join(os.path.abspath(""), "./assets/retail.jpg"),
            os.path.join(os.path.abspath(""), "./assets/aerodefence.jpg"),
        ],
        inputs=image_input,
        outputs=image_output,
        fn=annotator,
        cache_examples=False,
    )

    annotators_list.change(
        fn=annotator,
        inputs=[
            image_input,
            models,
            annotators_list,
            colorbb,
            colormask,
            colorellipse,
            colorbc,
            colorcir,
            colorlabel,
            colorhalo,
            colortri,
            colordot,
        ],
        outputs=image_output,
    )

    def change_color(color: ColorPicker):
        color.change(
            fn=annotator,
            inputs=[
                image_input,
                models,
                annotators_list,
                colorbb,
                colormask,
                colorellipse,
                colorbc,
                colorcir,
                colorlabel,
                colorhalo,
                colortri,
                colordot,
            ],
            outputs=image_output,
        )

    colors = [
        colorbb,
        colormask,
        colorellipse,
        colorbc,
        colorcir,
        colorlabel,
        colorhalo,
        colortri,
        colordot,
    ]

    for color in colors:
        change_color(color)


if __name__ == "__main__":
    print("Starting app...")
    print("Dark theme is available at: http://localhost:7860/?__theme=dark")
    # app.launch(debug=False, server_name="0.0.0.0")  # for local network
    app.launch(debug=False)