import torch
import gradio as gr
from transformers import Owlv2Processor, Owlv2ForObjectDetection, pipeline
import spaces
from pathlib import Path

# 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")
translator = pipeline("translation", model="Helsinki-NLP/opus-mt-fr-en")


@spaces.GPU
def query_image(img, description, score_threshold):
    description=description
    
    translation_result = translator(description ,
                                        src_lang="fr", 
                                        tgt_lang="en")
    description = translation_result[0]['translation_text']

    
    description = description.split(",")

    size = max(img.shape[:2])
    target_sizes = torch.Tensor([[size, size]])
    inputs = processor(text=description, images=img, 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, description[label.item()]))
    return img, result_labels


logo = r"""
<center><img src='https://doubiiu.github.io/projects/ToonCrafter/static/logo2.png' alt='ToonCrafter logo' style="width:280px; margin-bottom:2px"></center>
"""


gr.Markdown(logo)

description = """

***logo du catie***.

Détecter des objets en vocabulaire ouvert est une tâche difficile, étant donné la grande variété d'images possibles et de sujets possibles, cependant de récents modèles permettent de faire cela.

En voici notre adaptation en français.
Utilisez ce démonstrateur pour requêter n'importe quelle image avec la description textuelle d'un objet.

Pour s'en servir, chargez votre image et entrez des descritions séparées par des virgules.
Vous pouvez utiliser le seuil pour filtrer les détections avec une faible probabilité.


Si vous avez besoin d'un tel modèle, qui fonctionnera en temps réel, sera spécialisé sur vos propres données, de petite taille donc économe en énergie, contactez-nous : contact@catie.fr

"""
#image_path="./Vaniila.png"
#absolute_path = Path(image_path).resolve()

#gr.Image('Vaniila.png',
#             height=500,width=176,
#             show_download_button=False)


demo = gr.Interface(
    query_image,
    inputs=[gr.Image(), "text", gr.Slider(0, 1, value=0.1)],
    outputs="annotatedimage",
    title="Détection d'objets en vocabulaire ouvert.",
    description=description,
    #css=".gradio-container {background: url('file=Vaniila.png')}",
    examples=[
        ["voitures.jpg", "voiture orange, roues", 0.21],
        ["animos.jpeg", "lion, chimpanzé", 0.37],
        ["dej.jpeg", "jambon, assiette, croissant", 0.3],
        ],
    
)

   
demo.launch()