import json
import re
import subprocess
import sys
from pathlib import Path

import gradio as gr
import spaces
import torch
from PIL import Image, ImageDraw, ImageFont
from transformers import AutoModelForCausalLM, AutoProcessor

MODEL_ID = "google/gemma-4-26B-A4B-it"

processor = AutoProcessor.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(MODEL_ID, device_map="cuda:0", dtype=torch.bfloat16)

# Gemma 4 emits bounding boxes in a normalized 1000x1000 coordinate space.
COORD_SPACE = 1000
EXAMPLES_DIR = Path(__file__).parent / "examples"

def extract_json(text: str):
    """Pull the first JSON object/array out of a model response."""
    text = text.strip()
    text = re.sub(r"^```(?:json)?\s*", "", text)
    text = re.sub(r"\s*```$", "", text)

    try:
        return json.loads(text)
    except json.JSONDecodeError:
        pass

    match = re.search(r"(\{.*\}|\[.*\])", text, re.DOTALL)
    if match:
        return json.loads(match.group(1))
    raise ValueError("No valid JSON found in model output")


def draw_box(image: Image.Image, box, label: str) -> Image.Image:
    """Draw a Pascal-VOC style bounding box on a copy of the image.

    Gemma returns ``[ymin, xmin, ymax, xmax]`` in a 1000x1000 normalized space;
    we rescale to the image's pixel dimensions.
    """
    out = image.convert("RGB").copy()
    width, height = out.size
    ymin, xmin, ymax, xmax = box
    xmin = xmin / COORD_SPACE * width
    xmax = xmax / COORD_SPACE * width
    ymin = ymin / COORD_SPACE * height
    ymax = ymax / COORD_SPACE * height

    draw = ImageDraw.Draw(out)
    line_width = max(3, min(width, height) // 200)
    draw.rectangle([(xmin, ymin), (xmax, ymax)], outline="red", width=line_width)

    if label:
        font_size = max(14, min(width, height) // 40)
        try:
            font = ImageFont.truetype("DejaVuSans-Bold.ttf", font_size)
        except OSError:
            font = ImageFont.load_default()
        text_bbox = draw.textbbox((xmin, ymin), label, font=font)
        text_h = text_bbox[3] - text_bbox[1]
        text_w = text_bbox[2] - text_bbox[0]
        pad = 4
        text_y = max(0, ymin - text_h - 2 * pad)
        draw.rectangle(
            [(xmin, text_y), (xmin + text_w + 2 * pad, text_y + text_h + 2 * pad)],
            fill="yellow",
        )
        draw.text((xmin + pad, text_y + pad), label, fill="black", font=font)

    return out


@spaces.GPU(duration=60)
@torch.inference_mode()
def _detect_on_gpu(image: Image.Image, what_object: str) -> str:
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image", "image": image},
                {"type": "text", "text": f"What's the bounding box for the {what_object} in the image, in JSON format?"},
            ],
        }
    ]

    inputs = processor.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=True,
        return_dict=True,
        return_tensors="pt",
        enable_thinking=False,
    ).to(device=model.device)

    input_len = inputs["input_ids"].shape[-1]
    generated = model.generate(**inputs, max_new_tokens=60, do_sample=False)
    return processor.decode(generated[0, input_len:], skip_special_tokens=True)


def detect(image: Image.Image, what_object: str):
    if image is None:
        raise gr.Error("Please upload an image.")
    what_object = (what_object or "").strip()
    if not what_object:
        raise gr.Error("Please enter what to detect.")

    raw = _detect_on_gpu(image, what_object)

    if re.search(r"did not find", raw, re.IGNORECASE):
        gr.Info(f"No detections: the model could not find '{what_object}' in this image.")
        return image, raw

    try:
        parsed = extract_json(raw)
    except (ValueError, json.JSONDecodeError):
        gr.Warning(f"Could not parse model output as JSON. Raw response shown on the right.")
        return image, raw

    # Model usually returns a dict, but Claude insisted on testing for list
    detection = parsed[0] if isinstance(parsed, list) else parsed
    if "box_2d" not in detection:
        gr.Warning("Model output is missing 'box_2d'. Raw response shown on the right.")
        return image, json.dumps(detection, indent=2)

    box = detection["box_2d"]
    label = detection.get("label", what_object)
    annotated = draw_box(image, box, label)
    return annotated, json.dumps(detection, indent=2)


examples = [
    [str(EXAMPLES_DIR / "bike-48x48.jpg"), "bike"],
    [str(EXAMPLES_DIR / "boat-48x48.jpg"), "hat"],
    [str(EXAMPLES_DIR / "forbidden-48x48.jpg"), "person"],
    [str(EXAMPLES_DIR / "wheel-48x48.jpg"), "turquoise capsule"],
    [str(EXAMPLES_DIR / "recipe.png"), "view recipe button"],
]


with gr.Blocks(title="Gemma 4 Object Detection") as demo:
    gr.Markdown(
        """
        # Gemma 4 Object Detection

        This demo showcases the extraordinary out-of-the-box geometry awareness of Gemma 4.

        Just upload an image and ask the model ([Gemma 4 26B-A4B-it](https://huggingface.co/google/gemma-4-26B-A4B-it)) what to look for. It returns a
        bounding box for the requested object in JSON format, which we then draw on top of the image.

        To know more about Gemma 4, please visit [our blog post](https://hf.co/blog/gemma4).
        """
    )
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(label="Input image", type="pil", sources=["upload", "clipboard"])
            object_text = gr.Textbox(
                label="What to detect",
                placeholder="e.g. 'bike', 'person', 'turquoise capsule'",
            )
            run_btn = gr.Button("Detect", variant="primary")
        with gr.Column():
            output_image = gr.Image(label="Detection", type="pil")
            raw_json = gr.Code(label="Model output", language="json")

    gr.Examples(
        examples=examples,
        inputs=[input_image, object_text],
        outputs=[output_image, raw_json],
        fn=detect,
        cache_examples=False,
    )

    run_btn.click(fn=detect, inputs=[input_image, object_text], outputs=[output_image, raw_json])
    object_text.submit(fn=detect, inputs=[input_image, object_text], outputs=[output_image, raw_json])


if __name__ == "__main__":
    demo.launch()