import torch
import gradio as gr
import numpy as np
from PIL import Image, ImageDraw, ImageFont
from transformers import OwlViTProcessor, OwlViTForObjectDetection

model = OwlViTForObjectDetection.from_pretrained("google/owlvit-base-patch32").eval()
processor = OwlViTProcessor.from_pretrained("google/owlvit-base-patch32")


def query_image(img, text_queries):
    text_queries = text_queries.split(",")
    inputs = processor(text=text_queries, images=img, return_tensors="pt")

    with torch.no_grad():
        outputs = model(**inputs)

    target_sizes = torch.Tensor([[768, 768]])
    results = processor.post_process(outputs=outputs, target_sizes=target_sizes)
    boxes, scores, labels = results[0]["boxes"], results[0]["scores"], results[0]["labels"]

    draw = ImageDraw.Draw(img)  
    font = ImageFont.truetype("assets/Helvetica.ttf", size=22)

    score_threshold = 0.1
    for box, score, label in zip(boxes, scores, labels):
        box = [int(i) for i in box.tolist()]

        if score >= score_threshold:
            draw.rectangle(box, outline="red", width=4)
            text_loc =[box[0]+5, box[3]+10]
            draw.text(text_loc, text_queries[label], fill="red", font=font, stroke_width=1)
    
    img = np.array(img)
    return img


description = """
Gradio demo for <a href="https://huggingface.co/docs/transformers/main/en/model_doc/owlvit">OWL-ViT</a>, 
introduced in <a href="https://arxiv.org/abs/2205.06230">Simple Open-Vocabulary Object Detection
with Vision Transformers</a>. 
\n\nYou can use OWL-ViT to query images with text descriptions of any object. 
To use it, simply upload an image and enter comma separated text descriptions of objects you want to query the image for. 
\n\n<a href="https://colab.research.google.com/github/huggingface/notebooks/blob/main/examples/zeroshot_object_detection_with_owlvit.ipynb">Colab demo</a>
"""
demo = gr.Interface(
    query_image, 
    inputs=[gr.Image(shape=(768, 768), type="pil"), "text"], 
    outputs="image",
    title="Zero-Shot Object Detection with OWL-ViT",
    description=description,
    examples=[["assets/astronaut.png", "human face, rocket, flag, nasa badge"], ["assets/coffee.png", "coffee mug, spoon, plate"]]
)
demo.launch(debug=True)