"""Gradio demo: Hand Visibility Detector on ZeroGPU.

Detects hands, estimates 3D hand pose (WiLoR-mini), and predicts
per-keypoint visibility using the model from ryhara/hand-visibility-detector.
Green keypoints = visible, Red = occluded.
"""

import spaces  # MUST be before any torch / CUDA-touching import
import torch

# torch>=2.6 defaults to weights_only=True which breaks loading ultralytics
# YOLO checkpoints (they contain custom nn.Module subclasses). Override back.
_orig_torch_load = torch.load
def _patched_torch_load(*args, **kwargs):
    kwargs["weights_only"] = False
    return _orig_torch_load(*args, **kwargs)
torch.load = _patched_torch_load

import gradio as gr
import numpy as np

from hand_visibility_detector import HandVisibilityPipeline, draw_detections

pipe = HandVisibilityPipeline(
    device="cuda",
    dtype=torch.float32,
)


@spaces.GPU(duration=15)
def detect(
    image: np.ndarray,
    hand_conf: float = 0.3,
    show_bones: bool = True,
) -> tuple[np.ndarray, str]:
    """Detect hands and estimate per-keypoint visibility.

    Args:
        image: Input RGB image.
        hand_conf: Hand detection confidence threshold (0.1–0.9).
        show_bones: Whether to draw the skeleton bones.
    """
    if image is None:
        return np.zeros((256, 256, 3), dtype=np.uint8), "No image provided"
    pipe.hand_conf = hand_conf
    results = pipe.predict(image)
    annotated = draw_detections(image, results, show_bones=show_bones)
    info_lines = [f"Detected {len(results)} hand(s)"]
    for i, r in enumerate(results):
        side = "R" if r.is_right else "L"
        vis_str = np.array2string(r.visibility, precision=2, separator=", ")
        info_lines.append(
            f"  [{i}] {side}  conf={r.bbox_conf:.2f}  vis={vis_str}"
        )
    return annotated, "\n".join(info_lines)


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

with gr.Blocks(title="Hand Visibility Detector") as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown("# 🤚 Hand Visibility Detector")
        gr.Markdown(
            "Detect hands, estimate 3D pose (WiLoR-mini), and predict "
            "per-keypoint visibility. **Green** = visible, **Red** = occluded.\n\n"
            "Based on [Hand Visibility Detector](https://arxiv.org/abs/2608.11574) "
            "by Hara et al. · [GitHub](https://github.com/ryhara/hand_visibility_detector)"
        )

        with gr.Row():
            img_input = gr.Image(label="Input image", type="numpy")
            img_output = gr.Image(label="Result", type="numpy")

        with gr.Accordion("Settings", open=True):
            with gr.Row():
                hand_conf_slider = gr.Slider(
                    minimum=0.1, maximum=0.9, value=0.3, step=0.05,
                    label="Hand detection confidence",
                )
                show_bones_cb = gr.Checkbox(value=True, label="Show bones")

        img_info = gr.Textbox(label="Info", interactive=False)
        img_btn = gr.Button("Detect", variant="primary")
        img_btn.click(
            fn=detect,
            inputs=[img_input, hand_conf_slider, show_bones_cb],
            outputs=[img_output, img_info],
            api_name="detect",
        )

        gr.Examples(
            examples=[
                ["sample.png", 0.1, True],
            ],
            inputs=[img_input, hand_conf_slider, show_bones_cb],
            outputs=[img_output, img_info],
            fn=detect,
            cache_examples=True,
            cache_mode="lazy",
        )

demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)