import gradio as gr
from transformers import pipeline

# Chargement des pipelines
captioneur_blip = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
captioneur_git = pipeline("image-to-text", model="microsoft/git-base")
translateur = pipeline("translation_en_to_fr", model="Helsinki-NLP/opus-mt-en-fr")

# Fonction : Captioning en anglais
def caption_en_anglais(image, nom_model):
    if nom_model == "BLIP":
        resultat = captioneur_blip(image)
    else:
        resultat = captioneur_git(image)
    return resultat[0]['generated_text']

# Fonction : Captioning en anglais + traduction
def caption_en_francais(image, model_name):
    caption_ang = caption_en_anglais(image, model_name)
    caption_fra = translateur(caption_ang)[0]['translation_text']
    return caption_fra

# Interface 1 : Captioning en anglais
anglais_interface = gr.Interface(
    fn=caption_en_anglais,
    inputs=[
        gr.Image(type="pil", label="Téléverser une image"),
        gr.Radio(["BLIP", "GIT"], label="Choisir un modèle", value="BLIP")
    ],
    outputs=gr.Textbox(label="Légende en anglais"),
    title="Image Captioning - Anglais",
    description="Téléversez une image et choisissez un modèle pour obtenir une description en anglais."
)

# Interface 2 : Captioning en français
francais_interface = gr.Interface(
    fn=caption_en_francais,
    inputs=[
        gr.Image(type="pil", label="Téléverser une image"),
        gr.Radio(["BLIP", "GIT"], label="Choisir un modèle", value="BLIP")
    ],
    outputs=gr.Textbox(label="Légende traduite en français"),
    title="Image Captioning - Français",
    description="Téléversez une image, obtenez une légende en anglais automatiquement traduite en français."
)

# Lancer l'application avec navigation entre les deux
demo = gr.TabbedInterface(
    interface_list=[anglais_interface, francais_interface],
    tab_names=["Captioning Anglais", "Captioning Français"],
    theme='gstaff/xkcd'
)

demo.launch()