import spaces  # MUST come before torch / any CUDA-touching import
import torch
import torchvision.transforms as T
import gradio as gr
import numpy as np
import cv2
from PIL import Image
from copy import deepcopy
from detrpose import DETR

# Load the model at module scope, .to("cuda") eagerly.
model = DETR(model="detrpose_hgnetv2_l", device="cuda")

transforms = T.Compose(
    [
        T.Resize((640, 640)),
        T.ToTensor(),
    ]
)


@spaces.GPU(duration=30)
def predict(image: "Image.Image", threshold: float = 0.5) -> "Image.Image":
    """Detect multi-person poses in an image using DETRPose-L.

    Args:
        image: Input PIL image.
        threshold: Confidence threshold for keeping detected poses.
    """
    if image is None:
        return None

    im_pil = image.convert("RGB")

    w, h = im_pil.size
    orig_size = torch.tensor([[w, h]]).to("cuda")

    im_data = transforms(im_pil).unsqueeze(0).to("cuda")

    with torch.no_grad():
        outputs = model.model.model(im_data)
        results = model.model.postprocessor(outputs, orig_size)

    scores, labels, keypoints = results
    scores = scores[0].detach().cpu().numpy()
    keypoints = keypoints[0].detach().cpu().numpy()

    idx = scores > threshold
    valid_keypoints = keypoints[idx]

    im_cv2 = cv2.cvtColor(np.array(im_pil), cv2.COLOR_RGB2BGR)
    model.annotator.draw_on(im_cv2, valid_keypoints)

    result_rgb = cv2.cvtColor(im_cv2, cv2.COLOR_BGR2RGB)
    return Image.fromarray(result_rgb)


CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks() as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            "# DETRPose: Real-Time Multi-Person Pose Estimation\n"
            "Upload an image to detect human poses with skeleton keypoints overlaid. "
            "Powered by [DETRPose-L](https://huggingface.co/SebasJanampa/DETRPose_L_COCO) — "
            "the first real-time end-to-end transformer for multi-person pose estimation."
        )

        with gr.Row():
            with gr.Column(scale=1):
                input_image = gr.Image(
                    type="pil", label="Input Image", sources=["upload", "clipboard"]
                )
                threshold_slider = gr.Slider(
                    minimum=0.1,
                    maximum=0.9,
                    value=0.5,
                    step=0.05,
                    label="Confidence Threshold",
                )
                run_btn = gr.Button("Run", variant="primary")

            with gr.Column(scale=1):
                output_image = gr.Image(
                    type="pil", label="Pose Estimation Result"
                )

        with gr.Accordion("Advanced settings", open=False):
            gr.Markdown(
                "**Confidence Threshold**: Lower this to detect more poses (may include "
                "false positives). Raise it for stricter, more confident detections."
            )

        gr.Examples(
            examples=[
                ["example1.jpg", 0.5],
                ["example2.jpg", 0.5],
            ],
            inputs=[input_image, threshold_slider],
            outputs=output_image,
            fn=predict,
            cache_examples=True,
            cache_mode="lazy",
        )

    run_btn.click(
        fn=predict,
        inputs=[input_image, threshold_slider],
        outputs=output_image,
        api_name="predict",
    )

if __name__ == "__main__":
    demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)