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 = """

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 :
- spécialisé sur vos données et qui sera donc très performant sur celles-ci,
- qui fonctionne en temps réel, de petite taille donc économe en énergie,
- qui sera intégré dans votre infrastucture 

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(1).jpeg", "lion, chimpanzé", 0.37],
        ["dej(1).jpeg", "jambon, assiette, croissant", 0.3],
        ],
    
)
gr.Image("/file=catie(3).png", height=250,width=80, show_download_button=False)

demo.launch(allowed_paths=["catie(3).png"])