import gradio as gr
import spaces
import supervision as sv
from rfdetr import RFDETRBase
from rfdetr.util.coco_classes import COCO_CLASSES

MARKDOWN = """
# RF-DETR 🔥

<div style="display: flex; align-items: center; gap: 8px;">
  <a href="https://colab.research.google.com/github/roboflow-ai/notebooks/blob/main/notebooks/how-to-finetune-rf-detr-on-detection-dataset.ipynb">
    <img src="https://colab.research.google.com/assets/colab-badge.svg" alt="colab" />
  </a>
  <a href="https://blog.roboflow.com/rf-detr">
    <img src="https://raw.githubusercontent.com/roboflow-ai/notebooks/main/assets/badges/roboflow-blogpost.svg" alt="roboflow" />
  </a>
  <a href="https://github.com/roboflow/rf-detr">
    <img src="https://badges.aleen42.com/src/github.svg" alt="roboflow" />
  </a>
</div>

RF-DETR is a real-time, transformer-based object detection model architecture developed 
by [Roboflow](https://roboflow.com/) and released under the Apache 2.0 license.
"""

COLOR = sv.ColorPalette.from_hex([
    "#ffff00", "#ff9b00", "#ff8080", "#ff66b2", "#ff66ff", "#b266ff",
    "#9999ff", "#3399ff", "#66ffff", "#33ff99", "#66ff66", "#99ff00"
])

MODEL = RFDETRBase()


@spaces.GPU()
def inference(image, confidence):
    detections = MODEL.predict(image, threshold=confidence)

    text_scale = sv.calculate_optimal_text_scale(resolution_wh=image.size)
    thickness = sv.calculate_optimal_line_thickness(resolution_wh=image.size)

    bbox_annotator = sv.BoxAnnotator(color=COLOR, thickness=thickness)
    label_annotator = sv.LabelAnnotator(
        color=COLOR,
        text_color=sv.Color.BLACK,
        text_scale=text_scale,
        smart_position=True
    )

    labels = [
        f"{COCO_CLASSES[class_id]} {confidence:.2f}"
        for class_id, confidence
        in zip(detections.class_id, detections.confidence)
    ]

    annotated_image = image.copy()
    annotated_image = bbox_annotator.annotate(annotated_image, detections)
    annotated_image = label_annotator.annotate(annotated_image, detections, labels)
    return annotated_image

with gr.Blocks() as demo:
    gr.Markdown(MARKDOWN)
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(
                label="Input Image",
                image_mode='RGB',
                type='pil',
                height=600
            )
            confidence_slider = gr.Slider(
                label="Confidence",
                minimum=0.0,
                maximum=1.0,
                step=0.05,
                value=0.5,
            )
            submit_button = gr.Button("Submit")
        with gr.Column():
            output_image = gr.Image(
                label="Input Image",
                image_mode='RGB',
                type='pil',
                height=600
            )

        submit_button.click(
            inference,
            inputs=[input_image, confidence_slider],
            outputs=output_image
        )

demo.launch(debug=False, show_error=True)