import spaces
import torch
from transformers import AutoImageProcessor, LwDetrForObjectDetection
import supervision as sv
import gradio as gr


@spaces.GPU
def infer(model_name, image, confidence_threshold):
    # Dynamically scale text and boxes based on image size
    width, height = image.size
    text_scale = (width / 1000) * 0.5
    text_thickness = max(0.7, width / 500)

    label_annotator = sv.LabelAnnotator(
        text_padding=4,
        text_scale=text_scale,
        text_thickness=text_thickness,
        smart_position=True,
    )
    box_annotator = sv.BoxAnnotator()

    model_name = f"AnnaZhang/{model_name}"
    processor = AutoImageProcessor.from_pretrained(model_name)
    model = LwDetrForObjectDetection.from_pretrained(model_name)
    inputs = processor(images=image, return_tensors="pt")
    outputs = model(**inputs)

    target_sizes = torch.tensor([image.size[::-1]])
    results = processor.post_process_object_detection(
        outputs, target_sizes=target_sizes, threshold=confidence_threshold
    )[0]
    detections = sv.Detections.from_transformers(
        transformers_results=results, id2label=model.config.id2label
    )
    labels = [
        f"{class_name} {confidence:.2f}"
        for class_name, confidence in zip(
            detections["class_name"], detections.confidence
        )
    ]
    image = box_annotator.annotate(image, detections)
    image = label_annotator.annotate(image, detections, labels)
    return image


with gr.Blocks() as demo:
    gr.Markdown("# LWDETR Object Detection")
    gr.Markdown(
        "LWDETR is a transformer-based object detection model that is trained on the Objects365 and COCO datasets."
    )
    gr.Markdown(
        "This space is a demo of the LWDETR model. You can select a model and an image and see the results."
    )
    with gr.Row():
        with gr.Column():
            model = gr.Radio(
                [
                    "lwdetr_tiny_30e_objects365",
                    "lwdetr_small_30e_objects365",
                    "lwdetr_medium_30e_objects365",
                    "lwdetr_large_30e_objects365",
                    "lwdetr_xlarge_30e_objects365",
                    "lwdetr_tiny_60e_coco",
                    "lwdetr_small_60e_coco",
                    "lwdetr_medium_60e_coco",
                    "lwdetr_large_60e_coco",
                    "lwdetr_xlarge_60e_coco",
                ],
                value="lwdetr_xlarge_60e_coco",
                label="Model",
            )
            confidence_threshold = gr.Slider(
                minimum=0.0,
                maximum=1.0,
                value=0.3,
                step=0.1,
                label="Confidence Threshold",
            )
            input_image = gr.Image(label="Input Image", type="pil")
            send_btn = gr.Button("Infer", variant="primary")
        with gr.Column():
            output_image = gr.Image(label="Output Image", type="pil")

    gr.Examples(
        examples=[
            "samples/cats.jpg",
            "samples/detectron2.png",
            "samples/cat.jpg",
            "samples/hotdog.jpg",
        ],
        inputs=input_image,
    )
    send_btn.click(
        fn=infer,
        inputs=[model, input_image, confidence_threshold],
        outputs=[output_image],
    )


demo.launch(debug=True)