import torch
import gradio as gr
from transformers import Owlv2Processor, Owlv2ForObjectDetection
import spaces

# Use GPU if available
if torch.cuda.is_available():
    device = torch.device("cuda")
else:
    device = torch.device("cpu")

model = Owlv2ForObjectDetection.from_pretrained("google/owlv2-base-patch16-ensemble").to(device)
processor = Owlv2Processor.from_pretrained("google/owlv2-base-patch16-ensemble")

def query_image(Upload_Image, Text, score_threshold):
    Text = Text
    Text = Text.split(",")

    size = max(Upload_Image.shape[:2])
    target_sizes = torch.Tensor([[size, size]])
    inputs = processor(text=Text, images=Upload_Image, return_tensors="pt").to(device)

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

    outputs.logits = outputs.logits.cpu()
    outputs.pred_boxes = outputs.pred_boxes.cpu()
    results = processor.post_process_object_detection(outputs=outputs, target_sizes=target_sizes)
    boxes, scores, labels = results[0]["boxes"], results[0]["scores"], results[0]["labels"]

    result_labels = []
    for box, score, label in zip(boxes, scores, labels):
        box = [int(i) for i in box.tolist()]
        if score < score_threshold:
            continue
        result_labels.append((box, Text[label.item()]))
    return Upload_Image, result_labels


descricao = """
Você pode usar o AnyVision para consultar imagens com descrições de texto de qualquer objeto.
Para usá-lo, basta carregar uma imagem e inserir descrições de texto separadas por vírgula dos objetos que deseja consultar na imagem. 
Você também pode usar o controle deslizante de limiar de pontuação para definir um limite para filtrar previsões de baixa probabilidade.
Você pode obter melhores previsões consultando a imagem com modelos de texto usados no treinamento do modelo original: por exemplo, *"foto de uma bandeira estrelada"*, 
*"imagem de um sapato"*.
"""
demo = gr.Interface(
    query_image,
    inputs=[gr.Image(), "text", gr.Slider(0, 1, value=0.1)],
    outputs="annotatedimage",
    title="AnyVision - Detector de Objetos Zero-Shot com Owl2",
    description=descricao
)
demo.launch()