import spaces

import gradio as gr
import numpy as np
import torch
from transformers import (
    AutoModelForZeroShotObjectDetection,
    AutoProcessor,
    SamModel,
    SamProcessor,
)


CUDA_DEVICE = "cuda"
SAM_MODEL_ID = "facebook/sam-vit-base"
GROUNDING_DINO_MODEL_ID = "IDEA-Research/grounding-dino-base"


# ZeroGPU enables CUDA emulation during startup, so models should be moved to
# CUDA at module scope. Importing `spaces` before `torch` is important.
sam_model = SamModel.from_pretrained(SAM_MODEL_ID).to(CUDA_DEVICE).eval()
sam_processor = SamProcessor.from_pretrained(SAM_MODEL_ID)

dino_processor = AutoProcessor.from_pretrained(GROUNDING_DINO_MODEL_ID)
dino_model = (
    AutoModelForZeroShotObjectDetection.from_pretrained(GROUNDING_DINO_MODEL_ID)
    .to(CUDA_DEVICE)
    .eval()
)


def infer_dino(image: np.ndarray, text_queries: list[str], score_threshold: float):
    query_text = ". ".join(text_queries) + "."
    height, width = image.shape[:2]

    inputs = dino_processor(
        images=image,
        text=query_text,
        return_tensors="pt",
    ).to(CUDA_DEVICE)

    with torch.inference_mode():
        outputs = dino_model(**inputs)

    return dino_processor.post_process_grounded_object_detection(
        outputs=outputs,
        input_ids=inputs.input_ids,
        threshold=score_threshold,
        text_threshold=score_threshold,
        target_sizes=[(height, width)],
    )[0]


@spaces.GPU(duration=120)
def query_image(image: np.ndarray, text_queries: str, dino_threshold: float):
    if image is None:
        raise gr.Error("Please upload an image.")

    queries = [query.strip() for query in text_queries.split(",") if query.strip()]
    if not queries:
        raise gr.Error("Please enter at least one candidate label.")

    detections = infer_dino(image, queries, dino_threshold)
    boxes = detections["boxes"].detach().cpu()
    labels = detections.get("text_labels", detections.get("labels", []))

    if len(boxes) == 0:
        return image, []

    # Process every detected box in one SAM forward pass.
    sam_inputs = sam_processor(
        images=image,
        input_boxes=[boxes.tolist()],
        return_tensors="pt",
    ).to(CUDA_DEVICE)

    with torch.inference_mode():
        sam_outputs = sam_model(**sam_inputs, multimask_output=True)

    masks = sam_processor.image_processor.post_process_masks(
        sam_outputs.pred_masks.detach().cpu(),
        sam_inputs["original_sizes"].detach().cpu(),
        sam_inputs["reshaped_input_sizes"].detach().cpu(),
    )[0]

    iou_scores = sam_outputs.iou_scores.detach().cpu()[0]
    best_mask_indices = iou_scores.argmax(dim=-1)

    annotations = []
    for detection_index, label in enumerate(labels):
        best_mask_index = int(best_mask_indices[detection_index])
        mask = masks[detection_index, best_mask_index].numpy().astype(bool)
        annotations.append((mask, str(label)))

    return image, annotations


description = """
This Space combines Grounding DINO for open-vocabulary object detection with
SAM for text-prompted image segmentation. Enter comma-separated labels such as
`cat, dog, sidewalk, crosswalk`.
"""

demo = gr.Interface(
    fn=query_image,
    inputs=[
        gr.Image(type="numpy", label="Image Input"),
        gr.Textbox(
            label="Candidate Labels",
            placeholder="cat, dog, sidewalk, crosswalk",
        ),
        gr.Slider(
            minimum=0.0,
            maximum=1.0,
            value=0.25,
            step=0.01,
            label="Grounding DINO Confidence Threshold",
        ),
    ],
    outputs=gr.AnnotatedImage(label="Segmentation Result"),
    title="Grounding DINO + SAM Zero-shot Segmentation",
    description=description,
)


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