import gradio as gr
import spaces
import torch
from transformers import AutoProcessor, AutoModelForZeroShotObjectDetection
from PIL import Image
import time

def extract_model_short_name(model_id):
    return model_id.split("/")[-1].replace("-", " ").replace("_", " ")

model_llmdet_id = "iSEE-Laboratory/llmdet_tiny"
model_mm_grounding_id = "rziga/mm_grounding_dino_tiny_o365v1_goldg"
model_omdet_id = "omlab/omdet-turbo-swin-tiny-hf"
model_owlv2_id = "google/owlv2-large-patch14-ensemble"

model_llmdet_name = extract_model_short_name(model_llmdet_id)
model_mm_grounding_name = extract_model_short_name(model_mm_grounding_id)
model_omdet_name = extract_model_short_name(model_omdet_id)
model_owlv2_name = extract_model_short_name(model_owlv2_id)

@spaces.GPU
def detect_omdet(image: Image.Image, prompts: list, threshold: float):
    t0 = time.perf_counter()
    model_id = model_omdet_id
    device = "cuda" if torch.cuda.is_available() else "cpu"
    processor = AutoProcessor.from_pretrained(model_id)
    model = AutoModelForZeroShotObjectDetection.from_pretrained(model_id).to(device).eval()
    texts = [prompts]
    inputs = processor(images=image, text=texts, return_tensors="pt").to(device)
    with torch.no_grad():
        outputs = model(**inputs)
    results = processor.post_process_grounded_object_detection(
        outputs,
        threshold=threshold,
        target_sizes=[image.size[::-1]]
    )   
    result = results[0]
    annotations = []
    raw_results = []
    for box, score, label in zip(result["boxes"], result["scores"], result["labels"]):
        if score >= threshold:
            label_name = prompts[label]
            xmin, ymin, xmax, ymax = [int(x) for x in box.tolist()]
            annotations.append(((xmin, ymin, xmax, ymax), f"{label_name} {score:.2f}"))
            raw_results.append(f"Detected {label_name} with confidence {score:.2f} at location [{xmin}, {ymin}, {xmax}, {ymax}]")
    elapsed_ms = (time.perf_counter() - t0) * 1000
    time_taken = f"**Inference time ({model_omdet_name}):** {elapsed_ms:.0f} ms"
    raw_text = "\n".join(raw_results) if raw_results else "No detections"
    return annotations, raw_text, time_taken
@spaces.GPU
def detect_llmdet(image: Image.Image, prompts: list, threshold: float):
    t0 = time.perf_counter()
    model_id = model_llmdet_id
    device = "cuda" if torch.cuda.is_available() else "cpu"
    processor = AutoProcessor.from_pretrained(model_id)
    model = AutoModelForZeroShotObjectDetection.from_pretrained(model_id).to(device).eval()
    texts = [prompts]
    inputs = processor(images=image, text=texts, return_tensors="pt").to(device)
    with torch.no_grad():
        outputs = model(**inputs)
    results = processor.post_process_grounded_object_detection(
        outputs,
        threshold=threshold,
        target_sizes=[image.size[::-1]]
    )
    result = results[0]
    annotations = []
    raw_results = []
    for box, score, label in zip(result["boxes"], result["scores"], result["labels"]):
        if score >= threshold:
            xmin, ymin, xmax, ymax = [int(x) for x in box.tolist()]
            annotations.append(((xmin, ymin, xmax, ymax), f"{label} {score:.2f}"))
            raw_results.append(f"Detected {label} with confidence {score:.2f} at location [{xmin}, {ymin}, {xmax}, {ymax}]")
    elapsed_ms = (time.perf_counter() - t0) * 1000
    time_taken = f"**Inference time ({model_llmdet_name}):** {elapsed_ms:.0f} ms"
    raw_text = "\n".join(raw_results) if raw_results else "No detections"
    return annotations, raw_text, time_taken
@spaces.GPU
def detect_mm_grounding(image: Image.Image, prompts: list, threshold: float):
    t0 = time.perf_counter()
    model_id = model_mm_grounding_id
    device = "cuda" if torch.cuda.is_available() else "cpu"
    processor = AutoProcessor.from_pretrained(model_id)
    model = AutoModelForZeroShotObjectDetection.from_pretrained(model_id).to(device).eval()
    texts = [prompts]
    inputs = processor(images=image, text=texts, return_tensors="pt").to(device)
    with torch.no_grad():
        outputs = model(**inputs)
    results = processor.post_process_grounded_object_detection(
        outputs,
        threshold=threshold,
        target_sizes=[image.size[::-1]]
    )
    result = results[0]
    annotations = []
    raw_results = []
    for box, score, label in zip(result["boxes"], result["scores"], result["labels"]):
        if score >= threshold:
            xmin, ymin, xmax, ymax = [int(x) for x in box.tolist()]
            annotations.append(((xmin, ymin, xmax, ymax), f"{label} {score:.2f}"))
            raw_results.append(f"Detected {label} with confidence {score:.2f} at location [{xmin}, {ymin}, {xmax}, {ymax}]")
    elapsed_ms = (time.perf_counter() - t0) * 1000
    time_taken = f"**Inference time ({model_mm_grounding_name}):** {elapsed_ms:.0f} ms"
    raw_text = "\n".join(raw_results) if raw_results else "No detections"
    return annotations, raw_text, time_taken

@spaces.GPU
def detect_owlv2(image: Image.Image, prompts: list, threshold: float):
    t0 = time.perf_counter()
    model_id = model_owlv2_id
    device = "cuda" if torch.cuda.is_available() else "cpu"
    processor = AutoProcessor.from_pretrained(model_id)
    model = AutoModelForZeroShotObjectDetection.from_pretrained(model_id).to(device).eval()
    texts = [prompts]
    inputs = processor(images=image, text=texts, return_tensors="pt").to(device)
    with torch.no_grad():
        outputs = model(**inputs)
    results = processor.post_process_grounded_object_detection(
        outputs,
        threshold=threshold,
        target_sizes=[image.size[::-1]]
    )
    result = results[0]
    annotations = []
    raw_results = []
    for box, score, label in zip(result["boxes"], result["scores"], result["labels"]):
        if score >= threshold:
            label_name = prompts[label]
            xmin, ymin, xmax, ymax = [int(x) for x in box.tolist()]
            annotations.append(((xmin, ymin, xmax, ymax), f"{label_name} {score:.2f}"))
            raw_results.append(f"Detected {label_name} with confidence {score:.2f} at location [{xmin}, {ymin}, {xmax}, {ymax}]")
    elapsed_ms = (time.perf_counter() - t0) * 1000
    time_taken = f"**Inference time ({model_owlv2_name}):** {elapsed_ms:.0f} ms"
    raw_text = "\n".join(raw_results) if raw_results else "No detections"
    return annotations, raw_text, time_taken


def run_detection(image, prompts_str, threshold_llm, threshold_mm, threshold_owlv2, threshold_omdet):
    if image is None:
        return (None, []), "No detections", "", (None, []), "No detections", ""
    prompts = [p.strip() for p in prompts_str.split(",")]
    ann_llm, raw_llm, time_llm = detect_llmdet(image, prompts, threshold_llm)
    ann_mm, raw_mm, time_mm = detect_mm_grounding(image, prompts, threshold_mm)
    ann_owlv2, raw_owlv2, time_owlv2 = detect_owlv2(image, prompts, threshold_owlv2)
    ann_omdet, raw_omdet, time_omdet = detect_omdet(image, prompts, threshold_omdet)
    return (image, ann_llm), raw_llm, time_llm, (image, ann_mm), raw_mm, time_mm, (image, ann_owlv2), raw_owlv2, time_owlv2, (image, ann_omdet), raw_omdet, time_omdet

with gr.Blocks() as app:
    gr.Markdown("# Zero-Shot Object Detection Arena")
    gr.Markdown("### Compare different zero-shot object detection models on the same image and prompts.")
    with gr.Row():
        with gr.Column(scale=1):
            image = gr.Image(type="pil", label="Upload an image", height=400)
            prompts = gr.Textbox(label="Prompts (comma-separated)", value="a cat, a remote control")
            with gr.Accordion("Per-model confidence thresholds", open=True):
                threshold_llm = gr.Slider(label="Threshold for LLMDet", minimum=0.0, maximum=1.0, value=0.3)
                threshold_mm = gr.Slider(label="Threshold for MM GroundingDINO Tiny", minimum=0.0, maximum=1.0, value=0.3)
                threshold_owlv2 = gr.Slider(label="Threshold for OwlV2 Large", minimum=0.0, maximum=1.0, value=0.1)
                threshold_omdet = gr.Slider(label="Threshold for OMDet Turbo Swin Tiny", minimum=0.0, maximum=1.0, value=0.2)
            generate_btn = gr.Button(value="Detect")
        with gr.Row():
            with gr.Column(scale=2):
                output_image_llm = gr.AnnotatedImage(label=f"Annotated image for {model_llmdet_name}", height=400)
                output_text_llm = gr.Textbox(label=f"Model detections for {model_llmdet_name}", lines=5)
                output_time_llm = gr.Markdown()
            with gr.Column(scale=2):
                output_image_mm = gr.AnnotatedImage(label=f"Annotated image for {model_mm_grounding_name}", height=400)
                output_text_mm = gr.Textbox(label=f"Model detections for {model_mm_grounding_name}", lines=5)
                output_time_mm = gr.Markdown()
        with gr.Row():
            with gr.Column(scale=2):
                output_image_owlv2 = gr.AnnotatedImage(label=f"Annotated image for {model_owlv2_name}", height=400)
                output_text_owlv2 = gr.Textbox(label=f"Model detections for {model_owlv2_name}", lines=5)
                output_time_owlv2 = gr.Markdown()
            with gr.Column(scale=2):
                output_image_omdet = gr.AnnotatedImage(label=f"Annotated image for {model_omdet_name}", height=400)
                output_text_omdet = gr.Textbox(label=f"Model detections for {model_omdet_name}", lines=5)
                output_time_omdet = gr.Markdown()
    gr.Markdown("### Examples")
    example_data = [
        ["http://images.cocodataset.org/val2017/000000039769.jpg", "a cat, a remote control", 0.30, 0.30, 0.10, 0.30],
        ["http://images.cocodataset.org/val2017/000000000139.jpg", "a person, a tv, a remote", 0.35, 0.30, 0.12, 0.30],
    ]

    gr.Examples(
        examples=example_data,
        inputs=[image, prompts, threshold_llm, threshold_mm, threshold_owlv2, threshold_omdet],
        label="Click an example to populate the inputs",
    )
    inputs = [image, prompts, threshold_llm, threshold_mm, threshold_owlv2, threshold_omdet]
    outputs = [output_image_llm, output_text_llm, output_time_llm, output_image_mm, output_text_mm, output_time_mm, output_image_owlv2, output_text_owlv2, output_time_owlv2, output_image_omdet, output_text_omdet, output_time_omdet]
    generate_btn.click(
        fn=run_detection,
        inputs=inputs,
        outputs=outputs,
    )
    image.upload(
        fn=run_detection,
        inputs=inputs,
        outputs=outputs,
    )

app.launch()