import gradio as gr

from PIL import Image as PImage
from torch import cuda

from transformers import AutoModel, AutoProcessor, pipeline

from utils.image_utils import draw_results, heatmap_image, heatmap_image_rbf, mask_image
from utils.clip_utils import embed_image, embed_word, idxs_along_axes, idxs_by_dist, make_image
from utils.SigLip2 import SigLip2

DEVICE = "cuda" if cuda.is_available() else "cpu"

SIGLIP_MODEL = "google/siglip2-so400m-patch16-256"
CLIP_MODEL = "openai/clip-vit-large-patch14"
DETR_MODEL = "facebook/detr-resnet-50"
OWL_MODEL = "google/owlv2-base-patch16"

detr = pipeline(task="object-detection",
                model=DETR_MODEL,
                device=DEVICE)

owl = pipeline(task="zero-shot-object-detection",
               model=OWL_MODEL,
               device=DEVICE)

clip_processor = AutoProcessor.from_pretrained(CLIP_MODEL)
clip = AutoModel.from_pretrained(CLIP_MODEL, device_map="auto").to(DEVICE)

siglip = SigLip2(SIGLIP_MODEL)

def run_detr(img):
  predictions = detr(img)
  return draw_results(img, predictions)

def run_owl(img, classes_str):
  classes = [c.strip() for c in classes_str.split(",")]
  predictions = owl(img, candidate_labels=classes)
  return draw_results(img, predictions)

def run_clip(files, word0, word1=""):
  w0e = embed_word(word0, clip_processor, clip, DEVICE)
  w1e = embed_word(word1, clip_processor, clip, DEVICE)
  
  ies = []
  imgs = []
  for f in files:
    img = PImage.open(f.name).convert("RGB")
    img = img.resize((int(256 * img.width/img.height), 256))
    imgs.append(img)
    ies.append(embed_image(img, clip_processor, clip, DEVICE))

  if word1 == "":
    ordered_idxs = idxs_by_dist(ies, w0e)
    return make_image(imgs, ordered_idxs)
  else:
    ordered_idxs = idxs_along_axes(ies, (w0e, w1e))
    return make_image(imgs, ordered_idxs)

def run_siglip2(img, text):
  text = [text]
  similarity_map_np = siglip.get_gradient_activation_map(img, text)

  masked_img = mask_image(img, similarity_map_np)

  # heatmap_img = heatmap_image(similarity_map_np, size=img.size, sampling=PImage.Resampling.BILINEAR)
  # overlay_img = PImage.blend(img, heatmap_img.resize(img.size), 0.65)

  heatmap_img_rbf = heatmap_image_rbf(similarity_map_np, size=img.size)
  overlay_img_rbf = PImage.blend(img, heatmap_img_rbf.resize(img.size), 0.65)

  return [masked_img, overlay_img_rbf]


examples = [
  ("painted portrait young person", "painted portrait old person"),
  ("painted portrait happy person", "painted portrait worried person"),
]

with gr.Blocks() as demo:
  gr.Interface(
    title="Object Detection",
    description="[DETR](https://huggingface.co/facebook/detr-resnet-50) model from facebook (2020), trained on [COCO 2017](https://github.com/amikelive/coco-labels/blob/master/coco-labels-2014_2017.txt) dataset and labels.",
    api_name="object",
    fn=run_detr,
    inputs=gr.Image(type="pil"),
    outputs=gr.Image(format="jpeg"),
    flagging_mode="never",
  )

  gr.Interface(
    title="Zero-Shot Object Detection",
    description="[OWLv2](https://huggingface.co/google/owlv2-large-patch14-ensemble) model from google (2023).",
    api_name="zero",
    fn=run_owl,
    inputs=[gr.Image(type="pil"), gr.Textbox(label="Object", show_label=True)],
    outputs=gr.Image(format="jpeg"),
    flagging_mode="never",
  )

  gr.Interface(
    title="Contrastive Embedding",
    description="[CLIP](https://huggingface.co/openai/clip-vit-large-patch14) model from openai (2021).",
    api_name="clip",
    fn=run_clip,
    inputs=[gr.File(file_count="multiple"),
            gr.Textbox(label="1st Descriptor", show_label=True),
            gr.Textbox(label="2nd Descriptor", show_label=True)],
    outputs=gr.Image(format="jpeg"),
    flagging_mode="never",
  )

  gr.Interface(
    title="Activations",
    description="[SigLip2](https://huggingface.co/google/siglip2-so400m-patch16-256) model from Google (2025).",
    api_name="siglip2",
    fn=run_siglip2,
    inputs=[gr.Image(type="pil"), gr.Textbox(label="Activation Term", show_label=True)],
    outputs=[gr.Image(format="jpeg"), gr.Image(format="jpeg")],
    flagging_mode="never",
  )

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